SupplyAgent/supply_agent/agent/core.py

from __future__ import annotations

from collections.abc import AsyncIterator, Callable, Iterator

from supply_agent.agent.loop import AgentLoop
from supply_agent.config import Settings, get_settings
from supply_agent.llm.client import LLMClient
from supply_agent.logging.logger import AgentLogger
from supply_agent.logging.publish import publish_run_artifacts
from supply_agent.skills.registry import SkillRegistry
from supply_agent.tools.registry import ToolRegistry
from supply_agent.types import AgentEvent, AgentEventType, AgentResult, Message, Role


DEFAULT_SYSTEM_PROMPT = """\
你是一个通用 AI 助手，根据用户的要求理解任务并完成目标。

## 工作方式
1. 先理解用户真正想要什么，必要时先澄清关键信息
2. 将复杂任务拆解为可执行的步骤，逐步推进
3. 需要外部信息或操作时，主动调用可用工具
4. 完成后给出清晰、可直接使用的结果

## 工具使用
- 你有权访问一组工具，仅在确实需要时调用
- 调用前想清楚：需要什么输入、期望得到什么输出
- 工具返回后，结合结果继续推理，不要重复无效调用

## 技能（Skills）
- 若系统提示中列出了可用 Skills，且任务与某项技能匹配，先调用 `load_skill` 加载对应技能
- 技能包含专业流程和输出规范，加载后严格遵循，不要自行猜测格式

## 输出要求
- 回答紧扣用户要求，避免无关内容
- 结构清晰，重点突出，便于阅读和直接使用
- 无法完成时，说明原因并给出可行的替代方案
"""


class Agent:
    """
    Main Agent class — orchestrates LLM, tools, and skills.

    Usage:
        agent = Agent(model="anthropic/claude-sonnet-5")
        agent.tools.register(my_tool)
        result = agent.run("What is 2+2?")
    """

    def __init__(
        self,
        settings: Settings | None = None,
        *,
        name: str | None = None,
        model: str | None = None,
        system_prompt: str | None = None,
        tools: ToolRegistry | None = None,
        skills: SkillRegistry | None = None,
        max_iterations: int | None = None,
        temperature: float | None = None,
        reasoning_effort: str | None = None,
        logger: AgentLogger | None = None,
        completion_guard: Callable[[list[Message]], str | None] | None = None,
        completion_guard_blocks: bool = False,
        tool_call_budgets: dict[str, tuple[set[str], int]] | None = None,
        tool_repeat_requires_change: dict[str, set[str]] | None = None,
    ) -> None:
        self.name = name
        self.settings = settings or get_settings()
        self.logger = logger or AgentLogger(
            self.settings.logs_dir,
            enabled=self.settings.log_enabled,
        )
        self.llm = LLMClient(
            self.settings,
            logger=self.logger,
            reasoning_effort=reasoning_effort,
        )
        self.tools = tools or ToolRegistry()
        self.skills = skills or SkillRegistry(self.settings.skills_dir)
        self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
        self.max_iterations = max_iterations or self.settings.agent_max_iterations
        self.temperature = temperature
        self.completion_guard = completion_guard
        self.completion_guard_blocks = completion_guard_blocks
        self.tool_call_budgets = tool_call_budgets or {}
        self.tool_repeat_requires_change = tool_repeat_requires_change or {}

        if model:
            self.llm.set_model(model)

        self._active_skills: list[str] = []
        self._setup_builtin_tools()

    @property
    def model(self) -> str:
        return self.llm.model

    @model.setter
    def model(self, value: str) -> None:
        self.llm.set_model(value)

    def _setup_builtin_tools(self) -> None:
        """Register built-in tools for skill loading."""

        def load_skill(name: str) -> str:
            """Load a skill by name to get specialized instructions for a task."""
            context = self.skills.get_skill_context(name)
            if context is None:
                available = ", ".join(self.skills.list_skills()) or "none"
                return f'{{"error": "Skill \'{name}\' not found. Available: {available}"}}'
            if name not in self._active_skills:
                self._active_skills.append(name)
            return context

        load_skill.__doc__ = "Load a skill by name to get specialized instructions for a task."
        from supply_agent.tools.base import tool as tool_decorator

        decorated = tool_decorator(name="load_skill")(load_skill)
        self.tools.register(decorated, name="load_skill")

    def _build_system_message(self) -> Message:
        parts = [self.system_prompt]
        catalog = self.skills.get_catalog()
        if catalog:
            parts.append(catalog)
        for skill_name in self._active_skills:
            ctx = self.skills.get_skill_context(skill_name)
            if ctx:
                parts.append(ctx)
        return Message(role=Role.SYSTEM, content="\n\n".join(parts))

    def _create_loop(self, messages: list[Message]) -> AgentLoop:
        return AgentLoop(
            llm=self.llm,
            tools=self.tools,
            system_message=self._build_system_message(),
            system_message_builder=self._build_system_message,
            messages=messages,
            max_iterations=self.max_iterations,
            temperature=self.temperature,
            logger=self.logger,
            active_skills=self._active_skills,
            completion_guard=self.completion_guard,
            completion_guard_blocks=self.completion_guard_blocks,
            tool_call_budgets=self.tool_call_budgets,
            tool_repeat_requires_change=self.tool_repeat_requires_change,
        )

    def _finish_run(self, result: AgentResult) -> None:
        """Close run logs and publish visualization artifacts."""
        self.logger.end_run(result)
        publish_run_artifacts(self.logger)

    def run(self, user_input: str, *, history: list[Message] | None = None) -> AgentResult:
        """Run the agent synchronously with a user message."""
        self.logger.start_run(user_input, model=self.model, agent_name=self.name)
        messages = list(history or [])
        messages.append(Message(role=Role.USER, content=user_input))
        loop = self._create_loop(messages)
        result = loop.run()
        self._finish_run(result)
        return result

    async def arun_core(
        self, user_input: str, *, history: list[Message] | None = None
    ) -> AgentResult:
        """Run the agent loop without closing logs or publishing artifacts."""
        self.logger.start_run(user_input, model=self.model, agent_name=self.name)
        messages = list(history or [])
        messages.append(Message(role=Role.USER, content=user_input))
        loop = self._create_loop(messages)
        return await loop.arun()

    async def arun(
        self, user_input: str, *, history: list[Message] | None = None
    ) -> AgentResult:
        """Run the agent asynchronously."""
        result = await self.arun_core(user_input, history=history)
        self._finish_run(result)
        return result

    def stream(
        self, user_input: str, *, history: list[Message] | None = None
    ) -> Iterator[AgentEvent]:
        """Stream agent events during execution."""
        self.logger.start_run(user_input, model=self.model, agent_name=self.name)
        messages = list(history or [])
        messages.append(Message(role=Role.USER, content=user_input))
        loop = self._create_loop(messages)
        final_result: AgentResult | None = None
        for event in loop.stream():
            if event.type == AgentEventType.DONE:
                data = event.data
                final_result = AgentResult(
                    content=data.get("content", ""),
                    messages=loop.messages,
                    iterations=data.get("iterations", 0),
                    tool_calls_made=data.get("tool_calls_made", 0),
                    skills_used=list(self._active_skills),
                )
            yield event
        if final_result:
            self._finish_run(final_result)

    async def astream(
        self, user_input: str, *, history: list[Message] | None = None
    ) -> AsyncIterator[AgentEvent]:
        """Async stream agent events."""
        self.logger.start_run(user_input, model=self.model, agent_name=self.name)
        messages = list(history or [])
        messages.append(Message(role=Role.USER, content=user_input))
        loop = self._create_loop(messages)
        final_result: AgentResult | None = None
        async for event in loop.astream():
            if event.type == AgentEventType.DONE:
                data = event.data
                final_result = AgentResult(
                    content=data.get("content", ""),
                    messages=loop.messages,
                    iterations=data.get("iterations", 0),
                    tool_calls_made=data.get("tool_calls_made", 0),
                    skills_used=list(self._active_skills),
                )
            yield event
        if final_result:
            self._finish_run(final_result)

    def reset(self) -> None:
        """Reset active skills state."""
        self._active_skills.clear()