core/engine/src/agents/base.py - main - alexefan136/flowstack

core/engine/src/agents/base.py
148 строк5 KB
Alexander Efanov
Обновление репозитория
15 июл 2026, 12:19
15 июл 2026, 12:1976704c6
# src/agents/base.py
"""Base agent с интеграцией реального LLM."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator
import logging

from openai.types.chat import ChatCompletionMessageParam

from ..llm.client import get_llm_client

logger = logging.getLogger(__name__)

# =============================================================================
# Алиасы для совместимости с researcher.py и другой легаси-архитектурой
# =============================================================================
AgentState = dict[str, Any]  # type: ignore[misc]


class BaseAgent(ABC):
    """Базовый класс для всех агентов."""

    def __init__(self, name: str = "Agent", system_prompt: str = ""):
        self.name = name
        self.system_prompt = system_prompt

    @abstractmethod
    async def run(self, input_data: dict) -> dict:
        """Запустить агента (не-стриминг)."""
        pass

    async def run_stream(self, input_data: dict) -> AsyncGenerator[dict, None]:
        """Запустить агента со стримингом."""
        yield {"type": "error", "error": "Streaming not implemented"}


# Алиас для обратной совместимости (researcher.py использует `Agent`)
Agent = BaseAgent


class LLMAgent(BaseAgent):
    """Агент, использующий LLM для генерации ответов."""

    async def run(self, input_data: dict) -> dict:
        """Не-стриминг запуск."""
        client = get_llm_client()

        messages: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": str(input_data.get("topic", ""))},
        ]

        result = await client.chat_completion(messages)

        return {
            "agent": self.name,
            "output": result["content"],
            "tokens": result["tokens_total"],
            "model": result["model"],
        }

    async def run_stream(self, input_data: dict) -> AsyncGenerator[dict, None]:
        """Стриминг запуск."""
        client = get_llm_client()

        messages: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": str(input_data.get("topic", ""))},
        ]

        async for chunk in client.chat_completion_stream(messages):
            yield chunk


# =============================================================================
# Предустановленные агенты
# =============================================================================
# У каждого класса есть DEFAULT_SYSTEM_PROMPT, но конструктор принимает
# опциональные name и system_prompt — это позволяет api/agents.py создавать
# кастомные агенты через AgentClass(system_prompt=...), а также использовать
# предустановленные конфигурации.

class ResearchAgent(LLMAgent):
    """Агент-исследователь."""

    DEFAULT_SYSTEM_PROMPT = """You are a research assistant. Your task is to:
1. Analyze the user's question
2. Provide comprehensive, factual information
3. Structure your response clearly with headings and bullet points
4. Cite sources when possible

Be thorough but concise. Focus on accuracy."""

    def __init__(
        self,
        name: str = "Researcher",
        system_prompt: str | None = None,
    ):
        super().__init__(
            name=name,
            system_prompt=system_prompt or self.DEFAULT_SYSTEM_PROMPT,
        )


class WriterAgent(LLMAgent):
    """Агент-писатель."""

    DEFAULT_SYSTEM_PROMPT = """You are a professional writer. Your task is to:
1. Take the research provided
2. Write clear, engaging content
3. Use proper grammar and style
4. Format with markdown for readability

Make the content accessible to a general audience."""

    def __init__(
        self,
        name: str = "Writer",
        system_prompt: str | None = None,
    ):
        super().__init__(
            name=name,
            system_prompt=system_prompt or self.DEFAULT_SYSTEM_PROMPT,
        )


class CriticAgent(LLMAgent):
    """Агент-критик."""

    DEFAULT_SYSTEM_PROMPT = """You are a critical reviewer. Your task is to:
1. Analyze the content for accuracy and completeness
2. Identify potential issues or gaps
3. Suggest improvements
4. Be constructive but thorough

Point out both strengths and weaknesses."""

    def __init__(
        self,
        name: str = "Critic",
        system_prompt: str | None = None,
    ):
        super().__init__(
            name=name,
            system_prompt=system_prompt or self.DEFAULT_SYSTEM_PROMPT,
        )