SupplyAgent/supply_agent/agent/loop.py

from __future__ import annotations

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

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

if TYPE_CHECKING:
    from supply_agent.logging.logger import AgentLogger

# Outcome of one assistant turn before tools / completion.
_TurnKind = Literal["done", "tools"]


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

    Sync/async and stream/non-stream entry points share the same turn and
    tool-handling logic so skill injection and max-iteration behavior stay
    consistent.
    """

    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
        # Use `is not None` so an empty shared list from Agent is kept by identity.
        self.active_skills = active_skills if active_skills is not None else []
        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 _append_malformed_tool_feedback(self) -> None:
        self.messages.append(
            Message(
                role=Role.USER,
                content=(
                    "上一步连续生成了无效工具参数。请继续任务，下一步只调用一个"
                    "最必要的工具，严格使用其 schema，不得添加未定义字段。"
                ),
            )
        )

    def _available_tool_definitions(self) -> list[ToolDefinition] | None:
        return self.tools.definitions or None

    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),
        )

    def _record_tool_result(
        self,
        tool_call: ToolCall,
        result: ToolResult,
        iterations: int,
    ) -> None:
        if self.logger:
            self.logger.log_tool_call(
                iterations,
                tool_call.name,
                tool_call.arguments,
                result.content,
                result.is_error,
                tool_call_id=tool_call.id,
            )
        if tool_call.name == "load_skill" and not result.is_error:
            self._on_skill_loaded(tool_call.arguments)
            if self.logger:
                self.logger.log_skill_loaded(iterations, tool_call.arguments)
        self.messages.append(
            Message(
                role=Role.TOOL,
                content=result.content,
                tool_call_id=result.tool_call_id,
                name=result.name,
            )
        )

    def _resolve_tool_call_sync(self, tool_call: ToolCall) -> ToolResult:
        return self.tools.execute(
            tool_call.id, tool_call.name, tool_call.arguments
        )

    async def _resolve_tool_call_async(self, tool_call: ToolCall) -> ToolResult:
        return await self.tools.aexecute(
            tool_call.id, tool_call.name, tool_call.arguments
        )

    def _handle_assistant_message(
        self,
        response: Message,
        iterations: int,
    ) -> tuple[_TurnKind, AgentResult | None]:
        """Classify an assistant turn: finish, or run tools."""
        self.messages.append(response)
        if response.tool_calls:
            return "tools", None
        return "done", self._build_result(response.content or "", iterations)

    def _nudge_best_answer_sync(self, iterations: int) -> AgentResult:
        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 _nudge_best_answer_async(self, iterations: int) -> AgentResult:
        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 run(self) -> AgentResult:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            try:
                response = self.llm.chat(
                    self._all_messages(),
                    tools=self._available_tool_definitions(),
                    temperature=self.temperature,
                    iteration=iterations,
                )
            except MalformedFunctionCallError:
                self._append_malformed_tool_feedback()
                continue

            kind, done = self._handle_assistant_message(response, iterations)
            if kind == "done" and done is not None:
                return done

            assert response.tool_calls is not None
            for tc in response.tool_calls:
                self.tool_calls_made += 1
                result = self._resolve_tool_call_sync(tc)
                self._record_tool_result(tc, result, iterations)

        return self._nudge_best_answer_sync(iterations)

    async def arun(self) -> AgentResult:
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            try:
                response = await self.llm.achat(
                    self._all_messages(),
                    tools=self._available_tool_definitions(),
                    temperature=self.temperature,
                    iteration=iterations,
                )
            except MalformedFunctionCallError:
                self._append_malformed_tool_feedback()
                continue

            kind, done = self._handle_assistant_message(response, iterations)
            if kind == "done" and done is not None:
                return done

            assert response.tool_calls is not None
            for tc in response.tool_calls:
                self.tool_calls_made += 1
                result = await self._resolve_tool_call_async(tc)
                self._record_tool_result(tc, result, iterations)

        return await self._nudge_best_answer_async(iterations)

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

            try:
                response = self.llm.chat(
                    self._all_messages(),
                    tools=self._available_tool_definitions(),
                    temperature=self.temperature,
                    iteration=iterations,
                )
            except MalformedFunctionCallError:
                self._append_malformed_tool_feedback()
                continue

            kind, done = self._handle_assistant_message(response, iterations)
            if kind == "done" and done is not None:
                yield AgentEvent(
                    type=AgentEventType.MESSAGE,
                    data={"content": done.content},
                )
                yield AgentEvent(
                    type=AgentEventType.DONE,
                    data=done.model_dump(),
                )
                return

            assert response.tool_calls is not None
            for tc in response.tool_calls:
                self.tool_calls_made += 1
                yield AgentEvent(
                    type=AgentEventType.TOOL_CALL,
                    data={
                        "name": tc.name,
                        "arguments": tc.arguments,
                        "id": tc.id,
                    },
                )
                result = self._resolve_tool_call_sync(tc)
                self._record_tool_result(tc, result, iterations)
                yield AgentEvent(
                    type=AgentEventType.TOOL_RESULT,
                    data={
                        "name": result.name,
                        "content": result.content,
                        "is_error": result.is_error,
                    },
                )

        final = self._nudge_best_answer_sync(iterations)
        yield AgentEvent(
            type=AgentEventType.MESSAGE,
            data={"content": final.content},
        )
        yield AgentEvent(
            type=AgentEventType.DONE,
            data=final.model_dump(),
        )

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

            try:
                response = await self.llm.achat(
                    self._all_messages(),
                    tools=self._available_tool_definitions(),
                    temperature=self.temperature,
                    iteration=iterations,
                )
            except MalformedFunctionCallError:
                self._append_malformed_tool_feedback()
                continue

            kind, done = self._handle_assistant_message(response, iterations)
            if kind == "done" and done is not None:
                yield AgentEvent(
                    type=AgentEventType.MESSAGE,
                    data={"content": done.content},
                )
                yield AgentEvent(
                    type=AgentEventType.DONE,
                    data=done.model_dump(),
                )
                return

            assert response.tool_calls is not None
            for tc in response.tool_calls:
                self.tool_calls_made += 1
                yield AgentEvent(
                    type=AgentEventType.TOOL_CALL,
                    data={
                        "name": tc.name,
                        "arguments": tc.arguments,
                        "id": tc.id,
                    },
                )
                result = await self._resolve_tool_call_async(tc)
                self._record_tool_result(tc, result, iterations)
                yield AgentEvent(
                    type=AgentEventType.TOOL_RESULT,
                    data={
                        "name": result.name,
                        "content": result.content,
                        "is_error": result.is_error,
                    },
                )

        final = await self._nudge_best_answer_async(iterations)
        yield AgentEvent(
            type=AgentEventType.MESSAGE,
            data={"content": final.content},
        )
        yield AgentEvent(
            type=AgentEventType.DONE,
            data=final.model_dump(),
        )