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