dev-rcfa-memory/be-ircfa/src/ircfa_interaction/lib/context_compactor.py

import logging
import json
from typing import List, Dict, Any, Tuple
from datetime import datetime
from qwen_agent.llm import get_chat_model
from src.ircfa_interaction.llm.gemma_llamacpp import GemmaLlamaCppLLM
from src.ircfa_interaction.config import DEFAULT_CONTEXT_WINDOW

logger = logging.getLogger(__name__)

class ContextCompactor:
    def __init__(self, llm_config: dict):
        self.llm_config = llm_config
        self.llm = None
            
        # Determine context size limit from config or user spec (default: DEFAULT_CONTEXT_WINDOW)
        self.num_ctx = DEFAULT_CONTEXT_WINDOW
        model_name = llm_config.get("name", "").lower()
        if "gemma" in model_name:
            self.num_ctx = DEFAULT_CONTEXT_WINDOW
        else:
            # Check config defaults or settings
            self.num_ctx = llm_config.get("generate_cfg", {}).get("num_ctx", DEFAULT_CONTEXT_WINDOW)
            
        # Threshold at 75% of context window size
        self.threshold_tokens = int(self.num_ctx * 0.75)

    def _init_llm(self):
        if self.llm is not None:
            return
        if "gemma" in self.llm_config.get("name", "").lower() and "llamacpp" in self.llm_config.get("name", "").lower():
            self.llm = GemmaLlamaCppLLM(self.llm_config)
        else:
            self.llm = get_chat_model(self.llm_config)

    def estimate_tokens(self, text: str) -> int:
        """Estimate tokens based on characters and words (1 token approx 4 characters)."""
        if not text:
            return 0
        return max(len(text) // 4, int(len(text.split()) * 1.3))

    def estimate_messages_tokens(self, messages: List[Dict[str, str]]) -> int:
        """Estimate total tokens in a list of chat messages."""
        total = 0
        for m in messages:
            total += self.estimate_tokens(m.get("content", ""))
        return total

    async def _call_llm(self, prompt: str) -> str:
        """Call the LLM to get a response (handling sync, generator, and async responses)."""
        self._init_llm()
        try:
            # We run the chat model
            resp = self.llm.chat(
                messages=[{"role": "user", "content": prompt}],
                stream=False
            )
            
            content = ""
            # Handle if generator/iterator (e.g. streaming adapter returned by chat)
            if hasattr(resp, "__next__"):
                for chunk in resp:
                    if isinstance(chunk, list) and chunk:
                        # qwen-agent returns a list of Messages
                        msg = chunk[-1]
                        content = msg.get("content", "") if hasattr(msg, "get") else getattr(msg, "content", "")
            elif hasattr(resp, "__aiter__"):
                async for chunk in resp:
                    if isinstance(chunk, list) and chunk:
                        msg = chunk[-1]
                        content = msg.get("content", "") if hasattr(msg, "get") else getattr(msg, "content", "")
            else:
                if isinstance(resp, list) and resp:
                    msg = resp[-1]
                    content = msg.get("content", "") if hasattr(msg, "get") else getattr(msg, "content", "")
                elif hasattr(resp, "content"):
                    content = getattr(resp, "content", "")
                else:
                    content = str(resp)
                    
            return content.strip()
        except Exception as e:
            logger.error(f"Error in ContextCompactor LLM call: {e}")
            raise

    async def summarize_history(self, messages_to_summarize: List[Dict[str, str]]) -> str:
        """Generate a concise summary of the provided messages."""
        if not messages_to_summarize:
            return ""

        formatted_history = ""
        for m in messages_to_summarize:
            role = "User" if m["role"] == "user" else "Asisten"
            formatted_history += f"{role}: {m['content']}\n\n"

        prompt = (
            "Ringkas percakapan berikut secara sangat singkat dan padat dalam Bahasa Indonesia. "
            "Fokus pada gejala kerusakan, equipment/location tag yang dibahas, "
            "hasil analisis, dan kesimpulan sementara. Jangan hilangkan detail penting:\n\n"
            f"{formatted_history}\n\n"
            "Ringkasan Singkat:"
        )

        try:
            summary = await self._call_llm(prompt)
            return summary
        except Exception as e:
            logger.error(f"Failed to generate summary: {e}")
            return "[Percakapan sebelumnya dipadatkan secara otomatis untuk menghemat memori.]"

    async def combine_summaries(self, prev_summary: str, new_summary: str) -> str:
        """Combine old summary and new summary into a single condensed summary."""
        prompt = (
            "Gabungkan dua ringkasan percakapan berikut menjadi satu ringkasan yang padat, "
            "kohesif, dan kronologis dalam Bahasa Indonesia:\n\n"
            f"Ringkasan 1:\n{prev_summary}\n\n"
            f"Ringkasan 2:\n{new_summary}\n\n"
            "Ringkasan Gabungan:"
        )
        try:
            combined = await self._call_llm(prompt)
            return combined
        except Exception as e:
            logger.error(f"Failed to combine summaries: {e}")
            return f"{prev_summary}\n{new_summary}"

    async def compact_context(
        self,
        system_prompt: str,
        history_context: List[Dict[str, str]],
        doc_summary: str,
        current_query: str,
        rcfa_state: Dict[str, Any]
    ) -> Tuple[List[Dict[str, str]], Dict[str, Any], bool]:
        """
        Check context token size and perform compaction (summarization or trimming) if threshold is exceeded.
        Returns:
            - Updated history_context
            - Updated rcfa_state
            - Boolean flag indicating if compaction occurred
        """
        # Calculate current token usage
        sys_tokens = self.estimate_tokens(system_prompt)
        doc_tokens = self.estimate_tokens(doc_summary)
        query_tokens = self.estimate_tokens(current_query)
        history_tokens = self.estimate_messages_tokens(history_context)

        total_tokens = sys_tokens + doc_tokens + query_tokens + history_tokens
        
        logger.info(
            f"Context check: total={total_tokens} tokens (sys={sys_tokens}, doc={doc_tokens}, "
            f"query={query_tokens}, history={history_tokens}). Limit={self.num_ctx}, Threshold={self.threshold_tokens}"
        )

        if total_tokens <= self.threshold_tokens:
            return history_context, rcfa_state, False

        # Context window is full! We need to compact
        compaction_count = rcfa_state.get("context_compaction_count", 0)
        logger.warning(f"Context window limit threatened ({total_tokens} tokens). Current compaction count: {compaction_count}")

        if compaction_count < 3:
            # We can compact (summarize older messages)
            # We keep the last 3 messages (approx 1.5 turns) intact to maintain conversation continuity
            keep_count = min(3, len(history_context))
            messages_to_summarize = history_context[:-keep_count] if keep_count > 0 else history_context
            messages_to_keep = history_context[-keep_count:] if keep_count > 0 else []

            # Filter out any pre-existing system summary message from what we summarize 
            # (since that's already in the accumulated summary)
            messages_to_summarize = [m for m in messages_to_summarize if m["role"] != "system"]

            if not messages_to_summarize:
                # No older messages to summarize (history is already minimal, maybe docs are huge)
                # Fallback: Trim documents or history_context
                return history_context, rcfa_state, False

            logger.info(f"Summarizing {len(messages_to_summarize)} older history messages...")
            new_summary = await self.summarize_history(messages_to_summarize)
            
            # Combine with any previous summary stored in rcfa_state
            prev_summary = rcfa_state.get("accumulated_conversation_summary", "")
            if prev_summary:
                logger.info("Combining new summary with previous accumulated summary...")
                accumulated_summary = await self.combine_summaries(prev_summary, new_summary)
            else:
                accumulated_summary = new_summary

            # Find the latest timestamp of the summarized messages to set the cutoff
            cutoff_timestamp = None
            for m in reversed(messages_to_summarize):
                if m.get("timestamp"):
                    cutoff_timestamp = m.get("timestamp")
                    break

            # Increment compaction count and update state
            compaction_count += 1
            rcfa_state["context_compaction_count"] = compaction_count
            rcfa_state["accumulated_conversation_summary"] = accumulated_summary
            if cutoff_timestamp:
                rcfa_state["last_compacted_timestamp"] = cutoff_timestamp

            logger.info(f"Compaction successful. New compaction count: {compaction_count}, cutoff: {cutoff_timestamp}")
            return messages_to_keep, rcfa_state, True

        else:
            # Compaction count has reached 3! We must trim/discard older messages instead of summarizing
            logger.warning("Max compaction limit (3) reached. Trimming oldest history messages to free space.")
            # Keep only the last 4 messages (approx 2 turns)
            keep_count = min(4, len(history_context))
            new_history = history_context[-keep_count:] if keep_count > 0 else []
            
            # If the oldest kept message has a timestamp, update the database cutoff so they don't load next time
            cutoff_timestamp = None
            # Find the message just before the kept messages to discard them
            discarded_messages = history_context[:-keep_count] if keep_count > 0 else history_context
            for m in reversed(discarded_messages):
                if m.get("timestamp"):
                    cutoff_timestamp = m.get("timestamp")
                    break
            
            if cutoff_timestamp:
                rcfa_state["last_compacted_timestamp"] = cutoff_timestamp
            rcfa_state["context_trimmed"] = True
            
            return new_history, rcfa_state, True