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


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,
        completion_guard: Callable[[list[Message]], str | None] | None = None,
        tool_call_budgets: dict[str, tuple[set[str], int]] | None = None,
        tool_repeat_requires_change: dict[str, set[str]] | 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.completion_guard = completion_guard
        self.tool_call_budgets = tool_call_budgets or {}
        self.tool_repeat_requires_change = tool_repeat_requires_change or {}
        self._tool_budget_counts = {
            group: 0 for group in self.tool_call_budgets
        }
        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 _completion_block_reason(self) -> str | None:
        if self.completion_guard is None:
            return None
        return self.completion_guard(self.messages)

    def _append_completion_feedback(self, reason: str) -> None:
        self.messages.append(
            Message(
                role=Role.USER,
                content=(
                    "完成守卫拒绝当前最终回答："
                    f"{reason}。请继续调用必要工具修复这些问题；"
                    "只有守卫条件全部满足后才能输出最终答案。"
                ),
            )
        )

    def _append_malformed_tool_feedback(self) -> None:
        self.messages.append(
            Message(
                role=Role.USER,
                content=(
                    "上一步连续生成了无效工具参数。请继续任务，下一步只调用一个"
                    "最必要的工具，严格使用其 schema，不得添加未定义字段。"
                ),
            )
        )

    def _consume_tool_budget(self, tool_name: str) -> str | None:
        matching = [
            (group, limit)
            for group, (tool_names, limit) in self.tool_call_budgets.items()
            if tool_name in tool_names
        ]
        for group, limit in matching:
            used = self._tool_budget_counts[group]
            if used >= limit:
                return (
                    f"工具预算已耗尽: group={group}, limit={limit}。"
                    "请使用已有结果完成筛选、持久化和审计，不要继续扩大搜索。"
                )
        for group, _ in matching:
            self._tool_budget_counts[group] += 1
        return None

    def _repeat_policy_error(self, tool_name: str) -> str | None:
        dependencies = self.tool_repeat_requires_change.get(tool_name)
        if not dependencies:
            return None
        last_call_index = -1
        last_change_index = -1
        for index, message in enumerate(self.messages):
            if message.role != Role.TOOL or not message.name:
                continue
            try:
                payload = json.loads(message.content or "")
            except (TypeError, json.JSONDecodeError):
                continue
            if isinstance(payload, dict) and payload.get("error"):
                continue
            if message.name == tool_name:
                last_call_index = index
            if message.name in dependencies:
                last_change_index = index
        if last_call_index >= 0 and last_change_index < last_call_index:
            return (
                f"工具 {tool_name} 在状态未变化时不得重复调用。"
                f"必须先成功执行以下任一状态变更工具: {sorted(dependencies)}"
            )
        return None

    def _available_tool_definitions(self) -> list[ToolDefinition] | None:
        exhausted_tools: set[str] = set()
        for group, (tool_names, limit) in self.tool_call_budgets.items():
            if self._tool_budget_counts[group] >= limit:
                exhausted_tools.update(tool_names)
        definitions = [
            definition
            for definition in self.tools.definitions
            if definition.name not in exhausted_tools
            and self._repeat_policy_error(definition.name) is None
        ]
        return definitions or None

    def _refund_tool_budget_for_input_error(
        self,
        tool_name: str,
        result: ToolResult,
    ) -> None:
        try:
            payload = json.loads(result.content)
        except (TypeError, json.JSONDecodeError):
            return
        if not isinstance(payload, dict) or payload.get("input_error") is not True:
            return
        for group, (tool_names, _) in self.tool_call_budgets.items():
            if tool_name in tool_names and self._tool_budget_counts[group] > 0:
                self._tool_budget_counts[group] -= 1

    @staticmethod
    def _budget_error_result(tool_call: ToolCall, error: str) -> ToolResult:
        return ToolResult(
            tool_call_id=tool_call.id,
            name=tool_call.name,
            content=json.dumps(
                {
                    "error": error,
                    "budget_exhausted": True,
                    "title": "工具预算拒绝调用",
                },
                ensure_ascii=False,
            ),
            is_error=True,
        )

    @staticmethod
    def _policy_error_result(tool_call: ToolCall, error: str) -> ToolResult:
        return ToolResult(
            tool_call_id=tool_call.id,
            name=tool_call.name,
            content=json.dumps(
                {
                    "error": error,
                    "policy_rejected": True,
                    "title": "工具状态策略拒绝调用",
                },
                ensure_ascii=False,
            ),
            is_error=True,
        )

    def _max_iterations_result(self, iterations: int) -> AgentResult:
        reason = self._completion_block_reason()
        content = "Max iterations reached"
        if reason:
            content = f"Max iterations reached before completion: {reason}"
        return self._build_result(content, 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
            self.messages.append(response)

            if not response.tool_calls:
                reason = self._completion_block_reason()
                if reason:
                    self._append_completion_feedback(reason)
                    continue
                return self._build_result(response.content or "", iterations)

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                policy_error = self._repeat_policy_error(tc.name)
                budget_error = (
                    None if policy_error else self._consume_tool_budget(tc.name)
                )
                result = (
                    self._policy_error_result(tc, policy_error)
                    if policy_error
                    else (
                        self._budget_error_result(tc, budget_error)
                        if budget_error
                        else self.tools.execute(tc.id, tc.name, tc.arguments)
                    )
                )
                if not policy_error and not budget_error:
                    self._refund_tool_budget_for_input_error(tc.name, result)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations,
                        tc.name,
                        tc.arguments,
                        result.content,
                        result.is_error,
                        tool_call_id=tc.id,
                    )
                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,
                    )
                )

        if self.completion_guard is not None:
            return self._max_iterations_result(iterations)

        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
            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
            self.messages.append(response)

            if not response.tool_calls:
                reason = self._completion_block_reason()
                if reason:
                    self._append_completion_feedback(reason)
                    continue
                return self._build_result(response.content or "", iterations)

            for tc in response.tool_calls:
                self.tool_calls_made += 1
                policy_error = self._repeat_policy_error(tc.name)
                budget_error = (
                    None if policy_error else self._consume_tool_budget(tc.name)
                )
                result = (
                    self._policy_error_result(tc, policy_error)
                    if policy_error
                    else (
                        self._budget_error_result(tc, budget_error)
                        if budget_error
                        else await self.tools.aexecute(
                            tc.id, tc.name, tc.arguments
                        )
                    )
                )
                if not policy_error and not budget_error:
                    self._refund_tool_budget_for_input_error(tc.name, result)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations,
                        tc.name,
                        tc.arguments,
                        result.content,
                        result.is_error,
                        tool_call_id=tc.id,
                    )
                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,
                    )
                )

        if self.completion_guard is not None:
            return self._max_iterations_result(iterations)

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

            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
            self.messages.append(response)

            if not response.tool_calls:
                reason = self._completion_block_reason()
                if reason:
                    self._append_completion_feedback(reason)
                    continue
                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, "id": tc.id},
                )
                policy_error = self._repeat_policy_error(tc.name)
                budget_error = (
                    None if policy_error else self._consume_tool_budget(tc.name)
                )
                result = (
                    self._policy_error_result(tc, policy_error)
                    if policy_error
                    else (
                        self._budget_error_result(tc, budget_error)
                        if budget_error
                        else self.tools.execute(tc.id, tc.name, tc.arguments)
                    )
                )
                if not policy_error and not budget_error:
                    self._refund_tool_budget_for_input_error(tc.name, result)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations,
                        tc.name,
                        tc.arguments,
                        result.content,
                        result.is_error,
                        tool_call_id=tc.id,
                    )
                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=self._max_iterations_result(iterations).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
            self.messages.append(response)

            if not response.tool_calls:
                reason = self._completion_block_reason()
                if reason:
                    self._append_completion_feedback(reason)
                    continue
                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, "id": tc.id},
                )
                policy_error = self._repeat_policy_error(tc.name)
                budget_error = (
                    None if policy_error else self._consume_tool_budget(tc.name)
                )
                result = (
                    self._policy_error_result(tc, policy_error)
                    if policy_error
                    else (
                        self._budget_error_result(tc, budget_error)
                        if budget_error
                        else await self.tools.aexecute(
                            tc.id, tc.name, tc.arguments
                        )
                    )
                )
                if not policy_error and not budget_error:
                    self._refund_tool_budget_for_input_error(tc.name, result)
                if self.logger:
                    self.logger.log_tool_call(
                        iterations,
                        tc.name,
                        tc.arguments,
                        result.content,
                        result.is_error,
                        tool_call_id=tc.id,
                    )
                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=self._max_iterations_result(iterations).model_dump(),
        )

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