X-Agents/agent/app/xbot/loop.py

"""Agent loop for tool-calling conversation."""

import asyncio
import json
import re
from typing import Any, Callable, Optional

from loguru import logger


class AgentLoop:
    """
    Agent loop with tool-calling capability.

    This is the core of the nanobot agent - it handles:
    - Multi-turn conversation with the LLM
    - Tool execution when the model requests it
    - Progress callbacks for streaming responses
    """

    _TOOL_RESULT_MAX_CHARS = 50000

    def __init__(
        self,
        provider: Any,
        model: str,
        tools: Any,
        max_iterations: int = 50,
    ):
        """
        Initialize the agent loop.

        Args:
            provider: LLM provider (must implement chat_with_retry)
            model: Model name
            tools: Tool registry (must have get_definitions() and execute())
            max_iterations: Maximum tool call iterations
        """
        self.provider = provider
        self.model = model
        self.tools = tools
        self.max_iterations = max_iterations

    @staticmethod
    def _strip_think(text: Optional[str]) -> Optional[str]:
        """Strip model thinking blocks from content."""
        if not text:
            return None
        # Strip <thinking> tags commonly used by models like DeepSeek
        pattern = r"<thinking>[\s\S]*?</thinking>"
        text = re.sub(pattern, "", text)
        return text.strip() or None

    @staticmethod
    def _tool_hint(tool_calls: list) -> str:
        """Format tool calls as concise hint."""
        def _fmt(tc):
            args = tc.arguments or {}
            val = next(iter(args.values()), None) if isinstance(args, dict) else None
            if not isinstance(val, str):
                return tc.name
            return f'{tc.name}("{val[:40]}...")' if len(val) > 40 else f'{tc.name}("{val}")'
        return ", ".join(_fmt(tc) for tc in tool_calls)

    async def run_loop(
        self,
        initial_messages: list[dict],
        system_prompt: str = "",
        on_progress: Optional[Callable[..., Any]] = None,
    ) -> tuple[Optional[str], list[str], list[dict]]:
        """
        Run the agent iteration loop.

        Args:
            initial_messages: Starting message list
            system_prompt: System prompt to prepend
            on_progress: Optional callback for progress updates

        Returns:
            Tuple of (final_content, tools_used, all_messages)
        """
        # Prepend system prompt if provided
        if system_prompt:
            messages = [{"role": "system", "content": system_prompt}] + initial_messages
        else:
            messages = initial_messages

        iteration = 0
        final_content = None
        tools_used: list[str] = []

        while iteration < self.max_iterations:
            iteration += 1

            tool_defs = self.tools.get_definitions() if self.tools else []

            response = await self.provider.chat_with_retry(
                messages=messages,
                tools=tool_defs,
                model=self.model,
            )

            if response.has_tool_calls:
                # Send progress update
                if on_progress:
                    thought = self._strip_think(response.content)
                    if thought:
                        await on_progress(thought)
                    await on_progress(self._tool_hint(response.tool_calls), tool_hint=True)

                # Add assistant message with tool calls
                tool_call_dicts = [
                    tc.to_openai_tool_call() if hasattr(tc, 'to_openai_tool_call') else tc
                    for tc in response.tool_calls
                ]

                messages = self._add_assistant_message(
                    messages, response.content, tool_call_dicts,
                    reasoning_content=getattr(response, 'reasoning_content', None),
                )

                # Execute tools
                for tool_call in response.tool_calls:
                    tools_used.append(tool_call.name)
                    args_str = json.dumps(tool_call.arguments, ensure_ascii=False)
                    logger.info("Tool call: {}({})", tool_call.name, args_str[:200])

                    result = await self.tools.execute(tool_call.name, tool_call.arguments)
                    messages = self._add_tool_result(messages, tool_call.id, tool_call.name, result)
            else:
                clean = self._strip_think(response.content)

                # Handle error responses
                if response.finish_reason == "error":
                    logger.error("LLM returned error: {}", (clean or "")[:200])
                    final_content = clean or "Sorry, I encountered an error calling the AI model."
                    break

                messages = self._add_assistant_message(
                    messages, clean,
                    reasoning_content=getattr(response, 'reasoning_content', None),
                )
                final_content = clean
                break

        if final_content is None and iteration >= self.max_iterations:
            logger.warning("Max iterations ({}) reached", self.max_iterations)
            final_content = (
                f"I reached the maximum number of tool call iterations ({self.max_iterations}) "
                "without completing the task."
            )

        return final_content, tools_used, messages

    def _add_assistant_message(
        self,
        messages: list[dict],
        content: Optional[str],
        tool_calls: Optional[list[dict]] = None,
        reasoning_content: Optional[str] = None,
    ) -> list[dict]:
        """Add an assistant message to the message list."""
        msg: dict[str, Any] = {"role": "assistant", "content": content}
        if tool_calls:
            msg["tool_calls"] = tool_calls
        if reasoning_content is not None:
            msg["reasoning_content"] = reasoning_content
        messages.append(msg)
        return messages

    def _add_tool_result(
        self,
        messages: list[dict],
        tool_call_id: str,
        tool_name: str,
        result: Any,
    ) -> list[dict]:
        """Add a tool result message to the message list."""
        # Truncate large results
        content = str(result)
        if len(content) > self._TOOL_RESULT_MAX_CHARS:
            content = content[:self._TOOL_RESULT_MAX_CHARS] + "\n... (truncated)"

        messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "name": tool_name,
            "content": content,
        })
        return messages