SupplyAgent/supply_agent/agent/loop.py

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Callable, Iterator
from typing import TYPE_CHECKING

from supply_agent.llm.client import LLMClient
from supply_agent.tools.registry import ToolRegistry
from supply_agent.types import (
    AgentEvent,
    AgentEventType,
    AgentResult,
    Message,
    Role,
)

if TYPE_CHECKING:
    from supply_agent.logging.logger import AgentLogger


class AgentLoop:
    """
    ReAct-style agent loop: Reason → Act (tool call) → Observe → Repeat.

    Implements the standard tool-calling pattern used by modern agent frameworks.
    """

    def __init__(
        self,
        llm: LLMClient,
        tools: ToolRegistry,
        system_message: Message,
        messages: list[Message],
        max_iterations: int = 20,
        temperature: float | None = None,
        logger: AgentLogger | None = None,
        active_skills: list[str] | None = None,
        system_message_builder: Callable[[], Message] | None = None,
    ) -> None:
        self.llm = llm
        self.tools = tools
        self.system_message = system_message
        self.system_message_builder = system_message_builder
        self.messages = messages
        self.max_iterations = max_iterations
        self.temperature = temperature
        self.logger = logger
        self.active_skills = active_skills or []
        self.tool_calls_made = 0

    def _all_messages(self) -> list[Message]:
        return [self.system_message, *self.messages]

    def _on_skill_loaded(self, arguments: str) -> None:
        """Refresh system message after a skill is loaded."""
        try:
            args = json.loads(arguments)
            skill_name = args.get("name", "")
            if skill_name and skill_name not in self.active_skills:
                self.active_skills.append(skill_name)
        except json.JSONDecodeError:
            pass

        if self.system_message_builder:
            self.system_message = self.system_message_builder()

    def run(self) -> AgentResult:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            response = self.llm.chat(
                self._all_messages(),
                tools=self.tools.definitions or None,
                temperature=self.temperature,
                iteration=iterations,
            )
            self.messages.append(response)

            if not response.tool_calls:
                return self._build_result(response.content or "", iterations)

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                result = self.tools.execute(tc.id, tc.name, tc.arguments)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations, tc.name, tc.arguments, result.content, result.is_error
                    )
                if tc.name == "load_skill" and not result.is_error:
                    self._on_skill_loaded(tc.arguments)
                    if self.logger:
                        self.logger.log_skill_loaded(iterations, tc.arguments)
                self.messages.append(
                    Message(
                        role=Role.TOOL,
                        content=result.content,
                        tool_call_id=result.tool_call_id,
                        name=result.name,
                    )
                )

        self.messages.append(
            Message(
                role=Role.USER,
                content="Maximum iterations reached. Please provide your best answer now.",
            )
        )
        final = self.llm.chat(
            self._all_messages(),
            temperature=self.temperature,
            iteration=iterations + 1,
        )
        self.messages.append(final)
        return self._build_result(final.content or "", iterations)

    async def arun(self) -> AgentResult:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            response = await self.llm.achat(
                self._all_messages(),
                tools=self.tools.definitions or None,
                temperature=self.temperature,
                iteration=iterations,
            )
            self.messages.append(response)

            if not response.tool_calls:
                return self._build_result(response.content or "", iterations)

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                result = await self.tools.aexecute(tc.id, tc.name, tc.arguments)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations, tc.name, tc.arguments, result.content, result.is_error
                    )
                if tc.name == "load_skill" and not result.is_error:
                    self._on_skill_loaded(tc.arguments)
                    if self.logger:
                        self.logger.log_skill_loaded(iterations, tc.arguments)
                self.messages.append(
                    Message(
                        role=Role.TOOL,
                        content=result.content,
                        tool_call_id=result.tool_call_id,
                        name=result.name,
                    )
                )

        self.messages.append(
            Message(
                role=Role.USER,
                content="Maximum iterations reached. Please provide your best answer now.",
            )
        )
        final = await self.llm.achat(
            self._all_messages(),
            temperature=self.temperature,
            iteration=iterations + 1,
        )
        self.messages.append(final)
        return self._build_result(final.content or "", iterations)

    def stream(self) -> Iterator[AgentEvent]:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            yield AgentEvent(
                type=AgentEventType.THINKING,
                data={"iteration": iterations},
            )

            response = self.llm.chat(
                self._all_messages(),
                tools=self.tools.definitions or None,
                temperature=self.temperature,
                iteration=iterations,
            )
            self.messages.append(response)

            if not response.tool_calls:
                yield AgentEvent(
                    type=AgentEventType.MESSAGE,
                    data={"content": response.content or ""},
                )
                yield AgentEvent(
                    type=AgentEventType.DONE,
                    data=self._build_result(response.content or "", iterations).model_dump(),
                )
                return

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                yield AgentEvent(
                    type=AgentEventType.TOOL_CALL,
                    data={"name": tc.name, "arguments": tc.arguments},
                )
                result = self.tools.execute(tc.id, tc.name, tc.arguments)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations, tc.name, tc.arguments, result.content, result.is_error
                    )
                yield AgentEvent(
                    type=AgentEventType.TOOL_RESULT,
                    data={"name": result.name, "content": result.content, "is_error": result.is_error},
                )
                self.messages.append(
                    Message(
                        role=Role.TOOL,
                        content=result.content,
                        tool_call_id=result.tool_call_id,
                        name=result.name,
                    )
                )

        yield AgentEvent(type=AgentEventType.DONE, data={"content": "Max iterations reached"})

    async def astream(self) -> AsyncIterator[AgentEvent]:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            yield AgentEvent(
                type=AgentEventType.THINKING,
                data={"iteration": iterations},
            )

            response = await self.llm.achat(
                self._all_messages(),
                tools=self.tools.definitions or None,
                temperature=self.temperature,
                iteration=iterations,
            )
            self.messages.append(response)

            if not response.tool_calls:
                yield AgentEvent(
                    type=AgentEventType.MESSAGE,
                    data={"content": response.content or ""},
                )
                yield AgentEvent(
                    type=AgentEventType.DONE,
                    data=self._build_result(response.content or "", iterations).model_dump(),
                )
                return

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                yield AgentEvent(
                    type=AgentEventType.TOOL_CALL,
                    data={"name": tc.name, "arguments": tc.arguments},
                )
                result = await self.tools.aexecute(tc.id, tc.name, tc.arguments)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations, tc.name, tc.arguments, result.content, result.is_error
                    )
                yield AgentEvent(
                    type=AgentEventType.TOOL_RESULT,
                    data={"name": result.name, "content": result.content, "is_error": result.is_error},
                )
                self.messages.append(
                    Message(
                        role=Role.TOOL,
                        content=result.content,
                        tool_call_id=result.tool_call_id,
                        name=result.name,
                    )
                )

        yield AgentEvent(type=AgentEventType.DONE, data={"content": "Max iterations reached"})

    def _build_result(self, content: str, iterations: int) -> AgentResult:
        return AgentResult(
            content=content,
            messages=self.messages,
            iterations=iterations,
            tool_calls_made=self.tool_calls_made,
            skills_used=list(self.active_skills),
        )