core/engine/src/runtime/agent_runtime.py - main - alexefan136/flowstack

core/engine/src/runtime/agent_runtime.py
282 строки10 KB
Alexander Efanov
upd fix
31 июл 2026, 19:17
31 июл 2026, 19:17d146d86
"""Agent Runtime — выполнение агентов с LLM loop и function calling."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any

import structlog

from src.llm.client import LLMClient, get_llm_client
from src.llm.config import LLMConfig
from src.primitives import AgentContext, AgentState

logger = structlog.get_logger()

# Тип исполнителя инструментов: (tool_name, arguments) -> результат (str)
ToolExecutor = Callable[[str, dict[str, Any]], Awaitable[str]]


@dataclass
class AgentConfig:
    """Конфигурация агента."""

    agent_id: str
    agent_name: str = ""
    system_prompt: str = ""
    model: str = "deepseek-ai/DeepSeek-V4-Pro"
    tools: list[dict[str, Any]] = field(default_factory=list)  # OpenAI tools format
    temperature: float = 0.7
    max_tokens: int = 4096
    max_iterations: int = 10  # максимум итераций agent loop (защита от циклов)


class AgentRuntime:
    """
    Исполнитель агентов.

    Запускает ReAct-style loop:
    1. Запрос к LLM (с tools)
    2. Если LLM вернул tool_calls → выполнить tools → повторить
    3. Если LLM вернул финальный ответ → завершить
    """

    def __init__(
        self,
        llm_client: LLMClient | None = None,
        tool_executor: ToolExecutor | None = None,
    ) -> None:
        self.llm_client = llm_client or get_llm_client()
        self.tool_executor = tool_executor
        logger.info("agent_runtime.initialized")

    async def start(self) -> None:
        """Запустить runtime."""
        logger.info("agent_runtime.started")

    async def stop(self) -> None:
        """Остановить runtime."""
        logger.info("agent_runtime.stopped")

    def _build_llm_config(self, config: AgentConfig) -> LLMConfig:
        """Собрать LLMConfig из AgentConfig."""
        return LLMConfig(
            model=config.model,
            max_tokens=config.max_tokens,
            temperature=config.temperature,
            top_p=0.9,
            timeout=300.0,
        )

    def _build_messages(
        self, config: AgentConfig, messages: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Собрать список сообщений (system prompt + входные)."""
        result: list[dict[str, Any]] = []
        if config.system_prompt:
            result.append({"role": "system", "content": config.system_prompt})
        result.extend(messages)
        return result

    async def _execute_tool(self, tool_call: dict[str, Any], context: AgentContext | None) -> str:
        """Выполнить tool call и вернуть результат (строкой)."""
        func = tool_call.get("function", {})
        tool_name = func.get("name", "")
        raw_args = func.get("arguments", "{}")

        try:
            arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
        except json.JSONDecodeError:
            arguments = {}

        if self.tool_executor is None:
            return json.dumps({"error": f"No tool executor for '{tool_name}'"})

        try:
            return await self.tool_executor(tool_name, arguments)
        except Exception as e:
            logger.error("agent_runtime.tool_failed", tool=tool_name, error=str(e))
            return json.dumps({"error": str(e)})

    # ========================================================================
    # Non-streaming
    # ========================================================================

    async def run(
        self,
        config: AgentConfig,
        messages: list[dict[str, Any]],
        context: AgentContext | None = None,
    ) -> AgentState:
        """
        Выполнить агента (non-streaming, полный agent loop).

        Returns:
            AgentState с результатом.
        """
        state = AgentState(
            agent_id=config.agent_id,
            agent_name=config.agent_name or config.agent_id,
            model=config.model,
            temperature=config.temperature,
            max_tokens=config.max_tokens,
            system_prompt=config.system_prompt,
        )
        state.start()

        msgs = self._build_messages(config, messages)
        llm_config = self._build_llm_config(config)
        tools = config.tools or None

        try:
            for _ in range(config.max_iterations):
                state.add_llm_call()
                response = await self.llm_client.chat_completion(
                    msgs, config=llm_config, tools=tools
                )
                state.add_tokens(response.get("tokens_total", 0))

                tool_calls = response.get("tool_calls")
                content = response.get("content", "")

                # Есть tool calls → выполнить и продолжить loop
                if tool_calls:
                    assistant_msg: dict[str, Any] = {"role": "assistant", "content": content}
                    assistant_msg["tool_calls"] = tool_calls
                    msgs.append(assistant_msg)

                    for tc in tool_calls:
                        tool_name = tc.get("function", {}).get("name", "")
                        result = await self._execute_tool(tc, context)
                        state.add_tool_call(tool_name)
                        msgs.append(
                            {
                                "role": "tool",
                                "tool_call_id": tc.get("id", ""),
                                "content": result,
                            }
                        )
                    continue

                # Финальный ответ
                state.add_message("assistant", content)
                state.complete(
                    {
                        "output": content,
                        "reasoning_content": response.get("reasoning_content"),
                    }
                )
                return state

            # Превышен лимит итераций
            state.fail(f"Max iterations ({config.max_iterations}) exceeded")

        except Exception as e:
            logger.error("agent_runtime.run_failed", agent_id=config.agent_id, error=str(e))
            state.fail(str(e))

        return state

    # ========================================================================
    # Streaming
    # ========================================================================

    async def run_stream(
        self,
        config: AgentConfig,
        messages: list[dict[str, Any]],
        context: AgentContext | None = None,
    ) -> AsyncIterator[dict[str, Any]]:
        """
        Выполнить агента со стримингом событий.

        Yields события:
            agent_start, reasoning, agent_message, tool_call, tool_result,
            agent_done, error
        """
        state = AgentState(
            agent_id=config.agent_id,
            agent_name=config.agent_name or config.agent_id,
            model=config.model,
        )
        state.start()

        yield {"type": "agent_start", "agent": state.agent_name, "agent_id": config.agent_id}

        msgs = self._build_messages(config, messages)
        llm_config = self._build_llm_config(config)
        tools = config.tools or None

        try:
            for _ in range(config.max_iterations):
                state.add_llm_call()
                full_content = ""
                tool_calls: list[dict[str, Any]] = []

                async for event in self.llm_client.chat_completion_stream(
                    msgs, config=llm_config, tools=tools
                ):
                    etype = event.get("type")

                    if etype == "reasoning":
                        yield {"type": "reasoning", "content": event.get("content", "")}

                    elif etype == "content":
                        piece = event.get("content", "")
                        full_content += piece
                        yield {"type": "agent_message", "agent": state.agent_name, "content": piece}

                    elif etype == "tool_calls":
                        tool_calls = event.get("tool_calls", [])

                    elif etype == "done":
                        state.add_tokens(event.get("tokens_total", 0))

                    elif etype == "error":
                        yield {"type": "error", "error": event.get("error", "Unknown error")}
                        state.fail(event.get("error", "Unknown error"))
                        return

                # Есть tool calls → выполнить и продолжить
                if tool_calls:
                    assistant_msg: dict[str, Any] = {"role": "assistant", "content": full_content}
                    assistant_msg["tool_calls"] = tool_calls
                    msgs.append(assistant_msg)

                    for tc in tool_calls:
                        tool_name = tc.get("function", {}).get("name", "")
                        yield {"type": "tool_call", "tool": tool_name, "id": tc.get("id")}

                        result = await self._execute_tool(tc, context)
                        state.add_tool_call(tool_name)
                        yield {"type": "tool_result", "tool": tool_name, "result": result}

                        msgs.append(
                            {"role": "tool", "tool_call_id": tc.get("id", ""), "content": result}
                        )
                    continue

                # Финальный ответ
                state.add_message("assistant", full_content)
                state.complete({"output": full_content})
                yield {
                    "type": "agent_done",
                    "agent": state.agent_name,
                    "output": full_content,
                    "tokens": state.tokens_used,
                }
                return

            # Превышен лимит итераций
            yield {"type": "error", "error": f"Max iterations ({config.max_iterations}) exceeded"}
            state.fail("Max iterations exceeded")

        except Exception as e:
            logger.error("agent_runtime.stream_failed", agent_id=config.agent_id, error=str(e))
            yield {"type": "error", "error": str(e)}
            state.fail(str(e))


__all__ = ["AgentConfig", "AgentRuntime", "ToolExecutor"]