shira-hermes/api/services/agent_runner.py

"""Agent runner — wraps OpenAI client with tool-calling loop via ai-gateway."""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable

from openai import AsyncOpenAI

from mcp_server.espocrm_client import EspoCrmClient
from mcp_server.tools.crm_tools import register_crm_tools
from mcp_server.tools.document_tools import register_document_tools
from mcp_server.tools.legal_tools import register_legal_tools
from mcp_server.tools.legal_kb_tools import register_legal_kb_tools
from mcp_server.tools.signature_tools import register_signature_tools
from mcp_server.tools.billing_tools import register_billing_tools
from mcp_server.tools.legal_aid_tools import register_legal_aid_tools
from api.services.skills import register_skill_tools
from api.services.memory import register_memory_tools
from api.services.delegate import register_delegate_tool

logger = logging.getLogger("shira.agent")


# Hebrew labels for the on_event callback in run(). These are surfaced
# verbatim in the KnowledgeBase ask UI to give users live progress
# (instead of just an elapsed-time counter). Keep them short — they
# render in a single line of the progress log.

def _thinking_label(iteration: int) -> str:
    if iteration == 0:
        return "קוראת את השאלה"
    return "ממשיכה לחקור"


def _tool_label(name: str, args: dict) -> str:
    if name == "search_legal_kb":
        q = (args.get("query") or "").strip()
        if q:
            qs = q if len(q) <= 60 else q[:60] + "…"
            return f'מחפשת בבסיס הידע: "{qs}"'
        return "מחפשת בבסיס הידע"
    if name == "list_skills" or name == "view_skill":
        return "בודקת מיומנויות זמינות"
    if name == "save_memory" or name == "view_memory":
        return "מסתכלת בזיכרון"
    if name.startswith("delegate"):
        return "מאצילה משימת משנה לסוכן עזר"
    if name.startswith("create_") or name.startswith("update_") or name.startswith("save_"):
        return f"כותבת ב-CRM: {name}"
    return f"מפעילה כלי: {name}"


def _tool_done_label(name: str, result) -> str:
    if name == "search_legal_kb":
        # legal_kb_tools formats the result as a Hebrew string starting
        # with a "## נמצאו N קטעים" line. Surface that count.
        s = str(result or "")
        for line in s.splitlines():
            line = line.strip().lstrip("#").strip()
            if not line:
                continue
            return line[:80] if len(line) > 80 else line
        return "סיימה לחפש"
    return "הושלם"


class AgentRunner:
    """
    Runs an agentic conversation loop using the OpenAI SDK pointed at ai-gateway.

    The loop:
    1. Send messages + tools to the model
    2. If the model returns tool_calls, execute them and feed results back
    3. Repeat until the model returns a text response (no tool_calls)
    4. Return the final text
    """

    def __init__(
        self,
        ai_gateway_url: str,
        ai_gateway_api_key: str,
        model: str = "sonnet",
        max_tokens: int = 4096,
        max_iterations: int = 10,
        total_timeout: float = 570.0,
    ):
        self.client = AsyncOpenAI(
            base_url=f"{ai_gateway_url.rstrip('/')}/v1",
            api_key=ai_gateway_api_key,
        )
        self.model = model
        self.max_tokens = max_tokens
        self.max_iterations = max_iterations
        self.total_timeout = total_timeout

    async def run(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        context: dict,
        espocrm_url: str,
        espocrm_api_key: str,
        depth: int = 0,
        blocked_tools: set[str] | None = None,
        allowed_toolsets: list[str] | None = None,
        kb_sources_used: list[dict] | None = None,
        kb_topic_id: int | None = None,
        kb_topic_name: str | None = None,
        on_event: Callable[[dict], None] | None = None,
    ) -> str:
        """Run a conversation with tool-calling loop, bounded by total_timeout seconds."""
        try:
            return await asyncio.wait_for(
                self._run_impl(
                    system_prompt=system_prompt,
                    messages=messages,
                    context=context,
                    espocrm_url=espocrm_url,
                    espocrm_api_key=espocrm_api_key,
                    depth=depth,
                    blocked_tools=blocked_tools,
                    allowed_toolsets=allowed_toolsets,
                    kb_sources_used=kb_sources_used,
                    kb_topic_id=kb_topic_id,
                    kb_topic_name=kb_topic_name,
                    on_event=on_event,
                ),
                timeout=self.total_timeout,
            )
        except asyncio.TimeoutError:
            logger.warning("[agent] total timeout (%.0fs) reached", self.total_timeout)
            return "⏱️ המשימה לקחה יותר זמן מהמותר. נסי לפצל אותה לחלקים קטנים — למשל: קודם בקשי לקרוא מסמך ספציפי, ואז לכתוב את הסיכום."

    async def _run_impl(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        context: dict,
        espocrm_url: str,
        espocrm_api_key: str,
        depth: int = 0,
        blocked_tools: set[str] | None = None,
        allowed_toolsets: list[str] | None = None,
        kb_sources_used: list[dict] | None = None,
        kb_topic_id: int | None = None,
        kb_topic_name: str | None = None,
        on_event: Callable[[dict], None] | None = None,
    ) -> str:
        """Internal tool-calling loop. Returns the final text response.

        Args:
            depth: Current delegation depth. 0 = main conversation, 1 = child agent.
            blocked_tools: Tool names to exclude (used by delegation to prevent recursion).
            allowed_toolsets: If set, only register tools from these categories (crm, documents, legal, signature, billing, legal_aid).
            on_event: Optional sync callback fired between iterations and tool calls.
                Receives a dict {type, label, ...} where type ∈ {thinking,
                tool_start, tool_done, tool_error}. Used by /kb/ask/stream
                to surface live progress to the UI. Exceptions in the
                callback are logged and ignored — the agent never fails
                because the caller's display logic threw.
        """
        def emit(event: dict) -> None:
            if on_event is None:
                return
            try:
                on_event(event)
            except Exception as e:
                logger.warning("[agent] on_event callback failed: %s", e)

        # Build tools from context
        crm = EspoCrmClient(espocrm_url, espocrm_api_key)
        case_id = (context.get("case") or {}).get("id") or (context.get("drillDown", {}).get("case", {}).get("id"))
        user_id = (context.get("currentUser") or {}).get("id")

        tools_registry: dict[str, dict] = {}

        # Core CRM/document/legal tools
        should_register_all = allowed_toolsets is None
        if should_register_all or "crm" in (allowed_toolsets or []):
            register_crm_tools(tools_registry, crm, case_id, user_id, context)
        if should_register_all or "documents" in (allowed_toolsets or []):
            register_document_tools(tools_registry, crm, case_id, context)
        if should_register_all or "legal" in (allowed_toolsets or []):
            register_legal_tools(tools_registry, crm, case_id, context)
            register_legal_kb_tools(
                tools_registry,
                sources_used=kb_sources_used,
                topic_id=kb_topic_id,
                topic_name=kb_topic_name,
            )
        if should_register_all or "signature" in (allowed_toolsets or []):
            register_signature_tools(tools_registry, crm, case_id, user_id, context)
        if should_register_all or "billing" in (allowed_toolsets or []):
            register_billing_tools(tools_registry, crm, case_id, user_id, context)
        if should_register_all or "legal_aid" in (allowed_toolsets or []):
            register_legal_aid_tools(tools_registry, crm, case_id, user_id, context)

        # Phase 3 tools (only for main conversation, not child agents)
        if depth == 0:
            register_skill_tools(tools_registry)
            register_memory_tools(tools_registry, case_id, user_id)
            register_delegate_tool(tools_registry, context, espocrm_url, espocrm_api_key)

        # Remove blocked tools (for delegation safety)
        if blocked_tools:
            for name in blocked_tools:
                tools_registry.pop(name, None)

        # Build OpenAI tools format
        openai_tools = []
        for name, tool in tools_registry.items():
            openai_tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool["description"],
                    "parameters": tool["parameters"],
                },
            })

        # Build messages with system prompt
        openai_messages = [{"role": "system", "content": system_prompt}]
        openai_messages.extend(messages)
        self._last_messages = openai_messages

        # Agentic loop
        for iteration in range(self.max_iterations):
            logger.info("[agent] depth=%d iteration %d, messages=%d", depth, iteration + 1, len(openai_messages))

            emit({
                "type": "thinking",
                "iteration": iteration + 1,
                "label": _thinking_label(iteration),
            })

            try:
                response = await self.client.chat.completions.create(
                    model=self.model,
                    messages=openai_messages,
                    tools=openai_tools if openai_tools else None,
                    max_tokens=self.max_tokens,
                )
            except Exception as e:
                logger.error("[agent] API error: %s", e)
                return f"שגיאה בתקשורת עם שירות ה-AI: {e}"

            choice = response.choices[0]
            message = choice.message

            # If no tool calls, return the text
            if not message.tool_calls:
                emit({"type": "writing", "label": "מנסחת תשובה"})
                return message.content or "לא הצלחתי להבין את הבקשה. נסי שוב."

            # Process tool calls
            openai_messages.append(message.model_dump())

            for tool_call in message.tool_calls:
                fn_name = tool_call.function.name
                try:
                    fn_args = json.loads(tool_call.function.arguments)
                except json.JSONDecodeError:
                    fn_args = {}

                logger.info("[agent] tool call: %s(%s)", fn_name, json.dumps(fn_args, ensure_ascii=False)[:100])
                emit({
                    "type": "tool_start",
                    "name": fn_name,
                    "label": _tool_label(fn_name, fn_args),
                })

                tool_def = tools_registry.get(fn_name)
                if not tool_def:
                    result = f"❌ כלי לא ידוע: {fn_name}"
                    emit({"type": "tool_error", "name": fn_name, "label": f"כלי לא מוכר: {fn_name}"})
                else:
                    try:
                        handler = tool_def["handler"]
                        result = await handler(**fn_args)
                        emit({
                            "type": "tool_done",
                            "name": fn_name,
                            "label": _tool_done_label(fn_name, result),
                        })
                    except Exception as e:
                        logger.error("[agent] tool error %s: %s", fn_name, e)
                        result = f"❌ שגיאה בהפעלת {fn_name}: {e}"
                        emit({"type": "tool_error", "name": fn_name, "label": f"שגיאה ב-{fn_name}"})

                openai_messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": str(result),
                })

        # Exhausted iterations
        logger.warning("[agent] max iterations (%d) reached", self.max_iterations)
        return "הגעתי למספר המקסימלי של פעולות. נסי לשאול שוב בצורה ממוקדת יותר."

    def get_messages(self) -> list[dict]:
        """Return the conversation messages from the last run (for post-conversation learning)."""
        return getattr(self, "_last_messages", [])