gateway/modules/serviceCenter/services/serviceAgent/conversationManager.py

# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Conversation manager for the Agent service.
Handles message history, context window management, and progressive summarization."""

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

from modules.serviceCenter.services.serviceAgent.datamodelAgent import ToolDefinition

logger = logging.getLogger(__name__)

FIRST_SUMMARY_ROUND = 4
META_SUMMARY_ROUND = 7
KEEP_RECENT_MESSAGES = 4
MAX_ESTIMATED_TOKENS = 60000


class ConversationManager:
    """Manages the conversation history and context window for agent runs.

    Progressive summarization strategy:
      - Rounds 1-3: full conversation retained
      - Round 4+: older messages compressed into a running summary
      - Round 7+: meta-summary replaces prior summaries
    Supports RAG context injection before each round via injectRagContext."""

    def __init__(self, systemPrompt: str):
        self._messages: List[Dict[str, Any]] = [
            {"role": "system", "content": systemPrompt}
        ]
        self._summaries: List[Dict[str, Any]] = []
        self._lastSummarizedRound: int = 0
        self._ragContextInjected: bool = False

    @property
    def messages(self) -> List[Dict[str, Any]]:
        """Current messages for the next AI call (internal markers stripped)."""
        return [
            {k: v for k, v in msg.items() if not k.startswith("_")}
            for msg in self._messages
        ]

    def addUserMessage(self, content: str):
        """Add a user message."""
        self._messages.append({"role": "user", "content": content})

    def addAssistantMessage(self, content: str, toolCalls: List[Dict[str, Any]] = None):
        """Add an assistant message, optionally with tool calls."""
        msg: Dict[str, Any] = {"role": "assistant", "content": content}
        if toolCalls:
            msg["tool_calls"] = toolCalls
        self._messages.append(msg)

    def addToolResults(self, results: List[Dict[str, Any]]):
        """Add tool results to the conversation.
        Each result: {toolCallId, toolName, content}."""
        for result in results:
            self._messages.append({
                "role": "tool",
                "tool_call_id": result["toolCallId"],
                "content": result["content"]
            })

    def addToolResultsAsText(self, resultText: str):
        """Add combined tool results as a user message (text-based fallback)."""
        self._messages.append({
            "role": "user",
            "content": f"Tool Results:\n{resultText}"
        })

    def injectRagContext(self, ragContext: str):
        """Inject RAG context as a system message right after the main system prompt.

        Called before each agent round by the agent loop if KnowledgeService is available.
        Replaces any previously injected RAG context to keep the context fresh."""
        if not ragContext:
            return

        ragMessage = {
            "role": "system",
            "content": f"Relevant Knowledge (from indexed documents and workflow context):\n{ragContext}",
            "_isRagContext": True,
        }

        # Replace existing RAG message if present, otherwise insert after system prompt
        for i, msg in enumerate(self._messages):
            if msg.get("_isRagContext"):
                self._messages[i] = ragMessage
                self._ragContextInjected = True
                return

        # Insert after the first system prompt
        self._messages.insert(1, ragMessage)
        self._ragContextInjected = True

    def getMessageCount(self) -> int:
        """Get the number of messages (excluding system prompt)."""
        return len(self._messages) - 1

    def estimateTokenCount(self) -> int:
        """Rough estimate of total tokens in the conversation (4 chars ≈ 1 token)."""
        totalChars = sum(len(str(m.get("content", ""))) for m in self._messages)
        return totalChars // 4

    def needsSummarization(self, currentRound: int) -> bool:
        """Check if progressive summarization should be triggered.

        Triggers:
          - At round FIRST_SUMMARY_ROUND (4) if not yet summarized
          - At round META_SUMMARY_ROUND (7) for meta-summary
          - Every 5 rounds after that
          - When estimated token count exceeds MAX_ESTIMATED_TOKENS
        """
        if currentRound >= FIRST_SUMMARY_ROUND and self._lastSummarizedRound < currentRound:
            if currentRound == FIRST_SUMMARY_ROUND or currentRound == META_SUMMARY_ROUND:
                return True
            if (currentRound - META_SUMMARY_ROUND) % 5 == 0 and currentRound > META_SUMMARY_ROUND:
                return True
        if self.estimateTokenCount() > MAX_ESTIMATED_TOKENS:
            return True
        return False

    async def summarize(self, currentRound: int, aiCallFn) -> Optional[str]:
        """Perform progressive summarization of older messages.

        Rounds 1-3: full history retained, no summarization.
        Round 4+: compress older messages into a running summary.
        Round 7+: meta-summary that consolidates prior summaries.
        """
        if currentRound < FIRST_SUMMARY_ROUND and self.estimateTokenCount() <= MAX_ESTIMATED_TOKENS:
            return None

        systemMsgs = [m for m in self._messages if m.get("role") == "system"]
        nonSystemMessages = [m for m in self._messages if m.get("role") != "system"]

        keepRecent = min(KEEP_RECENT_MESSAGES, len(nonSystemMessages))
        if len(nonSystemMessages) <= keepRecent + 1:
            return None

        splitIdx = len(nonSystemMessages) - keepRecent
        # Ensure the split doesn't orphan tool messages from their assistant.
        # Walk backwards from splitIdx: if we're landing in the middle of a
        # tool-call sequence (assistant+tool_calls → tool → tool …), include
        # the entire sequence in recentMessages.
        while splitIdx > 0 and nonSystemMessages[splitIdx].get("role") == "tool":
            splitIdx -= 1
        # Also include the assistant message that triggered the tool calls.
        if splitIdx > 0 and splitIdx < len(nonSystemMessages) and \
           nonSystemMessages[splitIdx].get("role") == "assistant" and \
           nonSystemMessages[splitIdx].get("tool_calls"):
            pass  # splitIdx already points at the assistant; keep it in recent
        elif splitIdx == 0:
            return None  # nothing to summarize

        messagesToSummarize = nonSystemMessages[:splitIdx]
        recentMessages = nonSystemMessages[splitIdx:]

        summaryInput = _formatMessagesForSummary(messagesToSummarize)
        previousSummary = self._summaries[-1]["content"] if self._summaries else ""

        isMetaSummary = currentRound >= META_SUMMARY_ROUND and len(self._summaries) >= 2
        summaryPrompt = _buildSummaryPrompt(summaryInput, previousSummary, isMetaSummary)

        try:
            summaryText = await aiCallFn(summaryPrompt)
        except Exception as e:
            logger.error(f"Progressive summarization failed: {e}")
            return None

        self._summaries.append({
            "round": currentRound,
            "content": summaryText,
            "isMeta": isMetaSummary,
        })
        self._lastSummarizedRound = currentRound

        mainSystem = systemMsgs[0] if systemMsgs else {"role": "system", "content": ""}
        ragMessages = [m for m in systemMsgs if m.get("_isRagContext")]

        self._messages = [
            mainSystem,
            *ragMessages,
            {"role": "system", "content": f"Conversation Summary (rounds 1-{currentRound - keepRecent}):\n{summaryText}"},
            *recentMessages,
        ]

        logger.info(
            f"Progressive summarization at round {currentRound}: "
            f"compressed {len(messagesToSummarize)} messages into "
            f"{'meta-' if isMetaSummary else ''}summary"
        )
        return summaryText


def _formatMessagesForSummary(messages: List[Dict[str, Any]]) -> str:
    """Format messages into a text block for summarization."""
    parts = []
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        if role == "tool":
            toolName = msg.get("tool_call_id", "tool")
            parts.append(f"[Tool Result ({toolName})]:\n{content}")
        elif role == "assistant" and msg.get("tool_calls"):
            calls = msg["tool_calls"]
            callNames = [c.get("function", {}).get("name", "?") for c in calls]
            parts.append(f"[Assistant → Tool Calls: {', '.join(callNames)}]")
            if content:
                parts.append(f"[Assistant]: {content}")
        else:
            parts.append(f"[{role.capitalize()}]: {content}")
    return "\n\n".join(parts)


def _buildSummaryPrompt(messagesText: str, previousSummary: str, isMetaSummary: bool = False) -> str:
    """Build the prompt for progressive summarization."""
    if isMetaSummary:
        prompt = (
            "Create a comprehensive meta-summary consolidating the previous summary "
            "and the new messages. Preserve all key facts, decisions, entities (names, "
            "numbers, dates), tool results, and action outcomes. Be concise but complete.\n\n"
        )
    else:
        prompt = (
            "Summarize the following conversation concisely. Preserve all key facts, "
            "decisions, entities (names, numbers, dates), and tool results. "
            "Do not lose any important information.\n\n"
        )
    if previousSummary:
        prompt += f"Previous Summary:\n{previousSummary}\n\n"
    prompt += f"New Messages to Summarize:\n{messagesText}\n\nProvide a concise, factual summary:"
    return prompt


_LANGUAGE_NAMES = {
    "de": "German", "en": "English", "fr": "French", "it": "Italian",
    "es": "Spanish", "pt": "Portuguese", "nl": "Dutch", "ja": "Japanese",
    "zh": "Chinese", "ko": "Korean", "ar": "Arabic", "ru": "Russian",
}


def buildSystemPrompt(
    tools: List[ToolDefinition],
    toolsFormatted: str = None,
    userLanguage: str = "",
) -> str:
    """Build the system prompt for the agent.

    Args:
        tools: Available tool definitions.
        toolsFormatted: Pre-formatted tool descriptions for text-based fallback.
        userLanguage: ISO 639-1 language code (e.g. "de", "en"). The agent will
            respond in this language.
    """
    langName = _LANGUAGE_NAMES.get(userLanguage, "")
    langInstruction = (
        f"IMPORTANT: Always respond in {langName} ({userLanguage}). "
        f"The user's language is {langName}. All your messages, explanations, "
        f"and summaries MUST be in {langName}. "
        f"Only use English for tool call arguments and technical identifiers.\n\n"
    ) if langName else ""

    prompt = (
        f"{langInstruction}"
        "You are an AI agent with access to tools. "
        "Use the provided tools to accomplish the user's task. "
        "Think step by step. Call tools when you need information or need to perform actions. "
        "When you have enough information to answer, respond directly without calling tools.\n\n"
    )
    if toolsFormatted:
        prompt += f"Available Tools:\n{toolsFormatted}\n\n"
        prompt += (
            "To call a tool, use this format:\n"
            "```tool_call\n"
            "tool: <tool_name>\n"
            'args: {"param": "value"}\n'
            "```\n\n"
        )
    return prompt