X-Agents/core/agents/agent/loop.py

"""Agent run loop - complete implementation."""

import asyncio
import json
import logging
import re
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Awaitable, AsyncGenerator

from agents.agent.context import ContextBuilder
from agents.agent.memory import AgentMemory
from agents.llm import LLMProvider, LLMResponse, ProviderFactory
from agents.tools import ToolRegistry

logger = logging.getLogger(__name__)


class AgentLoop:
    """Agent loop with message processing, LLM calls, tool execution, and streaming."""

    _TOOL_RESULT_MAX_CHARS = 10000

    def __init__(
        self,
        provider: LLMProvider,
        model: str,
        workspace: Path | None = None,
        max_iterations: int = 10,
        tools: ToolRegistry | None = None,
    ):
        """Initialize the agent loop.

        Args:
            provider: LLM provider (OpenAI, Anthropic, etc.)
            model: Model name to use
            workspace: Workspace directory for memory and configs
            max_iterations: Maximum tool call iterations
            tools: Tool registry (creates default if None)
        """
        self.provider = provider
        self.model = model
        self.workspace = workspace or Path.cwd()
        self.max_iterations = max_iterations
        self.tools = tools

        self.context = ContextBuilder(self.workspace)
        self.memory = AgentMemory(self.workspace)

    async def chat(
        self,
        message: str,
        history: list[dict[str, Any]] | None = None,
        session_key: str = "default",
        on_progress: Callable[[str], Awaitable[None]] | None = None,
        model_id: str | None = None,
        model_name: str | None = None,
        model_provider: str | None = None,
        api_key: str | None = None,
        base_url: str | None = None,
        use_xbot: bool = False,
    ) -> str:
        """Process a chat message and return the response.

        Args:
            message: User message
            history: Conversation history
            session_key: Session identifier
            on_progress: Optional callback for progress updates
            model_id: Model ID (optional)
            model_name: Model name (optional)
            model_provider: Model provider (optional)
            api_key: API key (optional)
            base_url: Custom API base URL (optional)
            use_xbot: Use xbot mode (optional)

        Returns:
            Agent response content
        """
        history = history or []

        # Load history from session if session_key is provided
        if session_key and session_key != "default":
            loaded_history = self.memory.get_history(session_key, max_messages=20)
            if loaded_history:
                logger.info(f"Loaded {len(loaded_history)} messages from session history")
                # Merge loaded history with provided history (loaded takes precedence if empty)
                if not history:
                    history = loaded_history
                else:
                    # Append loaded history before current messages
                    history = loaded_history + history

        # Check if dynamic provider parameters are provided
        if api_key or model_provider:
            logger.info(f"Using dynamic provider: model_provider={model_provider}, model_name={model_name}, base_url={base_url}")
            # Create temporary provider with dynamic parameters
            temp_provider = ProviderFactory.create(
                provider=model_provider or "openai",
                api_key=api_key,
                api_base=base_url,
            )
            # Use temporary provider and model
            temp_model = model_name or temp_provider.get_default_model()
            logger.info(f"Created temp provider with model: {temp_model}")
            return await self._chat_with_provider(
                message=message,
                history=history,
                session_key=session_key,
                on_progress=on_progress,
                provider=temp_provider,
                model=temp_model,
            )

        # Build messages for LLM
        messages = self.context.build_messages(
            history=history,
            current_message=message,
        )

        # Log which provider is being used
        logger.info(f"Using static provider: {type(self.provider).__name__}, model={self.model}")

        # Run the agent loop
        final_content, tools_used, all_messages = await self._run_loop(
            messages, on_progress
        )

        # Save to history
        self._save_history(session_key, all_messages, len(history))

        return final_content or "No response generated."

    async def _chat_with_provider(
        self,
        message: str,
        history: list[dict[str, Any]] | None = None,
        session_key: str = "default",
        on_progress: Callable[[str], Awaitable[None]] | None = None,
        provider: LLMProvider | None = None,
        model: str | None = None,
    ) -> str:
        """Chat with a specific provider (used for dynamic provider support).

        Args:
            message: User message
            history: Conversation history
            session_key: Session identifier
            on_progress: Optional callback for progress updates
            provider: LLM provider to use
            model: Model name to use

        Returns:
            Agent response content
        """
        history = history or []

        # Load history from session if session_key is provided
        if session_key and session_key != "default":
            loaded_history = self.memory.get_history(session_key, max_messages=20)
            if loaded_history:
                logger.info(f"Loaded {len(loaded_history)} messages from session history")
                # Merge loaded history with provided history (loaded takes precedence if empty)
                if not history:
                    history = loaded_history
                else:
                    # Append loaded history before current messages
                    history = loaded_history + history

        provider = provider or self.provider
        model = model or self.model

        # Build messages for LLM
        messages = self.context.build_messages(
            history=history,
            current_message=message,
        )

        # Run the agent loop with custom provider
        final_content, tools_used, all_messages = await self._run_loop(
            messages, on_progress, provider=provider, model=model
        )

        # Save to history
        self._save_history(session_key, all_messages, len(history))

        return final_content or "No response generated."

    async def chat_stream(
        self,
        message: str,
        history: list[dict[str, Any]] | None = None,
        session_key: str = "default",
        model_id: str | None = None,
        model_name: str | None = None,
        model_provider: str | None = None,
        api_key: str | None = None,
        base_url: str | None = None,
        use_xbot: bool = False,
    ) -> AsyncGenerator[str, None]:
        """Process a chat message with streaming response.

        Args:
            message: User message
            history: Conversation history
            session_key: Session identifier
            model_id: Model ID (optional)
            model_name: Model name (optional)
            model_provider: Model provider (optional)
            api_key: API key (optional)
            base_url: Custom API base URL (optional)
            use_xbot: Use xbot mode (optional)

        Yields:
            Response content chunks
        """
        history = history or []

        # Load history from session if session_key is provided
        if session_key and session_key != "default":
            loaded_history = self.memory.get_history(session_key, max_messages=20)
            if loaded_history:
                logger.info(f"[stream] Loaded {len(loaded_history)} messages from session history")
                # Merge loaded history with provided history (loaded takes precedence if empty)
                if not history:
                    history = loaded_history
                else:
                    # Append loaded history before current messages
                    history = loaded_history + history

        # Check if dynamic provider parameters are provided
        if api_key or model_provider:
            logger.info(f"[stream] Using dynamic provider: model_provider={model_provider}, model_name={model_name}, base_url={base_url}")
            # Create temporary provider with dynamic parameters
            temp_provider = ProviderFactory.create(
                provider=model_provider or "openai",
                api_key=api_key,
                api_base=base_url,
            )
            # Use temporary provider and model
            temp_model = model_name or temp_provider.get_default_model()
            logger.info(f"[stream] Created temp provider with model: {temp_model}")
            async for chunk in self._chat_stream_with_provider(
                message=message,
                history=history,
                session_key=session_key,
                provider=temp_provider,
                model=temp_model,
            ):
                yield chunk
            return

        # Build messages for LLM
        messages = self.context.build_messages(
            history=history,
            current_message=message,
        )

        # Stream the response
        async for chunk in self._run_loop_stream(messages):
            yield chunk

    async def _chat_stream_with_provider(
        self,
        message: str,
        history: list[dict[str, Any]] | None = None,
        session_key: str = "default",
        provider: LLMProvider | None = None,
        model: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream chat with a specific provider (used for dynamic provider support).

        Args:
            message: User message
            history: Conversation history
            session_key: Session identifier
            provider: LLM provider to use
            model: Model name to use

        Yields:
            Response content chunks
        """
        history = history or []

        # Load history from session if session_key is provided
        if session_key and session_key != "default":
            loaded_history = self.memory.get_history(session_key, max_messages=20)
            if loaded_history:
                logger.info(f"[stream] Loaded {len(loaded_history)} messages from session history")
                # Merge loaded history with provided history (loaded takes precedence if empty)
                if not history:
                    history = loaded_history
                else:
                    # Append loaded history before current messages
                    history = loaded_history + history

        provider = provider or self.provider
        model = model or self.model

        # Build messages for LLM
        messages = self.context.build_messages(
            history=history,
            current_message=message,
        )

        # Stream the response with custom provider
        async for chunk in self._run_loop_stream(messages, provider=provider, model=model):
            yield chunk

    async def _run_loop(
        self,
        initial_messages: list[dict],
        on_progress: Callable[..., Awaitable[None]] | None = None,
        provider: LLMProvider | None = None,
        model: str | None = None,
    ) -> tuple[str | None, list[str], list[dict]]:
        """Run the agent iteration loop.

        Args:
            initial_messages: Initial message list
            on_progress: Progress callback
            provider: Optional LLM provider to use (defaults to self.provider)
            model: Optional model name to use (defaults to self.model)

        Returns:
            Tuple of (final_content, tools_used, all_messages)
        """
        messages = initial_messages
        iteration = 0
        final_content = None
        tools_used: list[str] = []
        provider = provider or self.provider
        model = model or self.model

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

        while iteration < self.max_iterations:
            iteration += 1

            # Call LLM
            response = await provider.chat_with_retry(
                messages=messages,
                tools=tool_defs if tool_defs else None,
                model=model,
            )

            if response.has_tool_calls:
                # Progress callback for tool calls
                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_dict() for tc in response.tool_calls]
                messages = self.context.add_assistant_message(
                    messages,
                    response.content,
                    tool_call_dicts,
                    reasoning_content=response.reasoning_content,
                )

                # Execute tools
                for tool_call in response.tool_calls:
                    tools_used.append(tool_call.name)
                    args = tool_call.arguments
                    logger.info(f"Tool call: {tool_call.name}({args})")

                    # Execute tool
                    result = await self._execute_tool(tool_call.name, args)

                    # Truncate large results
                    if len(result) > self._TOOL_RESULT_MAX_CHARS:
                        result = result[:self._TOOL_RESULT_MAX_CHARS] + "\n... (truncated)"

                    # Add tool result
                    messages = self.context.add_tool_result(
                        messages, tool_call.id, tool_call.name, result
                    )
            else:
                # No tool calls - return the response
                clean = self._strip_think(response.content)

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

                messages = self.context.add_assistant_message(
                    messages, clean, reasoning_content=response.reasoning_content
                )
                final_content = clean
                break

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

        return final_content, tools_used, messages

    async def _run_loop_stream(
        self,
        initial_messages: list[dict],
        provider: LLMProvider | None = None,
        model: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """Run the agent loop with streaming response.

        Args:
            initial_messages: Initial message list
            provider: Optional LLM provider to use (defaults to self.provider)
            model: Optional model name to use (defaults to self.model)

        Yields:
            Response content chunks
        """
        provider = provider or self.provider
        model = model or self.model
        tool_defs = self.tools.get_definitions() if self.tools else []

        # First call to check for tool calls
        response = await provider.chat_with_retry(
            messages=initial_messages,
            tools=tool_defs if tool_defs else None,
            model=model,
        )

        if response.has_tool_calls:
            # Execute tools first
            for tool_call in response.tool_calls:
                logger.info(f"Tool call: {tool_call.name}")
                result = await self._execute_tool(tool_call.name, tool_call.arguments)

                # Add to messages
                initial_messages = self.context.add_tool_result(
                    initial_messages, tool_call.id, tool_call.name, result
                )

            # Recursive call after tool execution
            async for chunk in self._run_loop_stream(initial_messages, provider=provider, model=model):
                yield chunk
        else:
            # Stream the content
            content = self._strip_think(response.content)
            if content:
                yield content

    async def _execute_tool(self, tool_name: str, args: dict) -> str:
        """Execute a tool.

        Args:
            tool_name: Name of the tool to execute
            args: Tool arguments

        Returns:
            Tool execution result
        """
        if self.tools:
            return await self.tools.execute(tool_name, args)
        return json.dumps({"error": "No tools registered"})

    @staticmethod
    def _strip_think(text: str | None) -> str | None:
        """Strip think blocks that some models embed in content."""
        if not text:
            return None
        import re
        # Match content between [/INST] or [/CONTINUE] tags commonly used in thinking
        patterns = [
            r"<think>[\s\S]*?</think>",
            r"<\/?think>",
        ]
        for pattern in patterns:
            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)

    def _save_history(
        self,
        session_key: str,
        messages: list[dict],
        skip: int = 0,
    ) -> None:
        """Save messages to history.

        Args:
            session_key: Session identifier
            messages: Messages to save
            skip: Number of messages to skip
        """
        for m in messages[skip:]:
            role = m.get("role")
            content = m.get("content")

            if role == "user" and content:
                self.memory.add_to_history("user", str(content)[:1000], session_key)
            elif role == "assistant" and content:
                self.memory.add_to_history("assistant", str(content)[:1000], session_key)
            # Save tool_calls for assistant messages (needed for multi-turn tool calls)
            elif role == "assistant" and m.get("tool_calls"):
                # Save the assistant message with tool_calls
                tool_calls_str = json.dumps(m.get("tool_calls", []))
                self.memory.add_to_history("assistant", f"[tool_calls]{tool_calls_str}", session_key)
            # Save tool results (needed for multi-turn conversations)
            elif role == "tool":
                tool_call_id = m.get("tool_call_id", "")
                tool_name = m.get("name", "")
                tool_content = m.get("content", "")
                tool_result_str = json.dumps({
                    "tool_call_id": tool_call_id,
                    "name": tool_name,
                    "content": tool_content
                })
                self.memory.add_to_history("tool", f"[tool_result]{tool_result_str}", session_key)