atlas/home-voice-agent/conversation/summarization/summarizer.py

"""
Conversation summarization using LLM.

Summarizes long conversations to reduce context size while preserving important information.
"""

import logging
from typing import List, Dict, Any, Optional
from datetime import datetime

logger = logging.getLogger(__name__)


class ConversationSummarizer:
    """Summarizes conversations to reduce context size."""
    
    def __init__(self, llm_client=None):
        """
        Initialize summarizer.
        
        Args:
            llm_client: LLM client for summarization (optional, can be set later)
        """
        self.llm_client = llm_client
    
    def should_summarize(self, 
                        message_count: int,
                        total_tokens: int,
                        max_messages: int = 20,
                        max_tokens: int = 4000) -> bool:
        """
        Determine if conversation should be summarized.
        
        Args:
            message_count: Number of messages in conversation
            total_tokens: Total token count
            max_messages: Maximum messages before summarization
            max_tokens: Maximum tokens before summarization
            
        Returns:
            True if summarization is needed
        """
        return message_count > max_messages or total_tokens > max_tokens
    
    def create_summary_prompt(self, messages: List[Dict[str, Any]]) -> str:
        """
        Create prompt for summarization.
        
        Args:
            messages: List of conversation messages
            
        Returns:
            Summarization prompt
        """
        # Format messages
        conversation_text = "\n".join([
            f"{msg['role'].upper()}: {msg['content']}"
            for msg in messages
        ])
        
        prompt = f"""Please summarize the following conversation, preserving:
1. Important facts and information mentioned
2. Decisions made or actions taken
3. User preferences or requests
4. Any tasks or reminders created
5. Key context for future conversations

Conversation:
{conversation_text}

Provide a concise summary that captures the essential information:"""
        
        return prompt
    
    def summarize(self, 
                  messages: List[Dict[str, Any]],
                  agent_type: str = "family") -> Dict[str, Any]:
        """
        Summarize a conversation.
        
        Args:
            messages: List of conversation messages
            agent_type: Agent type ("work" or "family")
            
        Returns:
            Summary dict with summary text and metadata
        """
        if not self.llm_client:
            # Fallback: simple extraction if no LLM available
            return self._simple_summary(messages)
        
        try:
            prompt = self.create_summary_prompt(messages)
            
            # Use LLM to summarize
            # This would call the LLM client - for now, return structured response
            summary_response = {
                "summary": "Summary would be generated by LLM",
                "key_points": [],
                "timestamp": datetime.now().isoformat(),
                "message_count": len(messages),
                "original_tokens": self._estimate_tokens(messages)
            }
            
            # TODO: Integrate with actual LLM client
            # summary_response = self.llm_client.generate(prompt, agent_type=agent_type)
            
            return summary_response
            
        except Exception as e:
            logger.error(f"Error summarizing conversation: {e}")
            return self._simple_summary(messages)
    
    def _simple_summary(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Create a simple summary without LLM."""
        user_messages = [msg for msg in messages if msg.get("role") == "user"]
        assistant_messages = [msg for msg in messages if msg.get("role") == "assistant"]
        
        summary = f"Conversation with {len(user_messages)} user messages and {len(assistant_messages)} assistant responses."
        
        # Extract key phrases
        key_points = []
        for msg in user_messages:
            content = msg.get("content", "")
            if len(content) > 50:
                key_points.append(content[:100] + "...")
        
        return {
            "summary": summary,
            "key_points": key_points[:5],  # Top 5 points
            "timestamp": datetime.now().isoformat(),
            "message_count": len(messages),
            "original_tokens": self._estimate_tokens(messages)
        }
    
    def _estimate_tokens(self, messages: List[Dict[str, Any]]) -> int:
        """Estimate token count (rough: 4 chars per token)."""
        total_chars = sum(len(str(msg.get("content", ""))) for msg in messages)
        return total_chars // 4
    
    def prune_messages(self,
                      messages: List[Dict[str, Any]],
                      keep_recent: int = 10,
                      summary: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """
        Prune messages, keeping recent ones and adding summary.
        
        Args:
            messages: List of messages
            keep_recent: Number of recent messages to keep
            summary: Optional summary to add at the beginning
            
        Returns:
            Pruned message list with summary
        """
        # Keep recent messages
        recent_messages = messages[-keep_recent:] if len(messages) > keep_recent else messages
        
        # Add summary as system message if available
        pruned = []
        if summary:
            pruned.append({
                "role": "system",
                "content": f"[Previous conversation summary: {summary.get('summary', '')}]"
            })
        
        pruned.extend(recent_messages)
        
        return pruned


# Global summarizer instance
_summarizer = ConversationSummarizer()


def get_summarizer() -> ConversationSummarizer:
    """Get the global summarizer instance."""
    return _summarizer