agents.py

from types import NoneType
from typing import Literal, Sequence, TypeVar
from config import settings
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.models.anthropic import AnthropicModel
from openai import AsyncOpenAI
from model import Queries, Feedback, ReportState, Sections, SectionState
from model.base import BaseModel
from prompts import (
    report_planner_query_writer_instructions,
    query_writer_instructions,
    report_planner_instructions,
    section_writer_instructions,
    section_grader_instructions,
    final_section_writer_instructions,
)
from pydantic_ai import Agent, RunContext
from pydantic_ai.agent import EndStrategy
from config import ModelProvider
from anthropic import AsyncAnthropic

TA = TypeVar("TA", BaseModel, NoneType)
TB = TypeVar("TB", BaseModel, str)

openai_client: AsyncOpenAI | None = None
anthropic_client: AsyncAnthropic | None = None


def get_model_by_provider_and_mode(
    provider: ModelProvider, mode: Literal["planner", "writer"]
) -> OpenAIModel | AnthropicModel:
    global openai_client, anthropic_client
    match provider:
        case ModelProvider.OPENAI:
            if not openai_client:
                openai_client = AsyncOpenAI(
                    api_key=settings.openai_api_key, base_url=settings.openai_base_url
                )
            match mode:
                case "planner":
                    return OpenAIModel(
                        settings.planner_model, openai_client=openai_client
                    )
                case "writer":
                    return OpenAIModel(
                        settings.writer_model, openai_client=openai_client
                    )
        case ModelProvider.ANTHROPIC:
            if not anthropic_client:
                anthropic_client = AsyncAnthropic(
                    api_key=settings.anthropic_api_key,
                    base_url=settings.anthropic_base_url,
                )
            match mode:
                case "planner":
                    return AnthropicModel(
                        settings.planner_model, anthropic_client=anthropic_client
                    )
                case "writer":
                    return AnthropicModel(
                        settings.writer_model, anthropic_client=anthropic_client
                    )


def build_agent(
    model: OpenAIModel | AnthropicModel = get_model_by_provider_and_mode(
        settings.writer_provider, "writer"
    ),
    result_type: type[TB] = BaseModel,
    deps_type: type[TA] = NoneType,
    name: str | None = None,
    retries=settings.max_retries,
    system_prompt: str | Sequence[str] = (),
    # tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
    end_strategy: EndStrategy = "early",
) -> Agent[TA, TB]:
    agent = Agent(
        model=model,
        result_type=result_type,
        deps_type=deps_type,
        name=name,
        retries=retries,
        system_prompt=system_prompt,
        end_strategy=end_strategy,
    )
    return agent


agent_for_gen_query_in_report_planner: Agent[ReportState, Queries] = build_agent(
    result_type=Queries,
    deps_type=ReportState,
)

agent_for_gen_section_in_report_planner: Agent[ReportState, Sections] = build_agent(
    get_model_by_provider_and_mode(settings.planner_provider, "planner"),
    result_type=Sections,
    deps_type=ReportState,
)

agent_for_gen_query: Agent[SectionState, Queries] = build_agent(
    result_type=Queries,
    deps_type=SectionState,
)

agent_for_gen_section: Agent[SectionState, str] = build_agent(
    result_type=str,
    deps_type=SectionState,
)

agent_for_gen_section_feedback: Agent[SectionState, Feedback] = build_agent(
    result_type=Feedback,
    deps_type=SectionState,
)

agent_for_gen_final_section: Agent[SectionState, str] = build_agent(
    result_type=str,
    deps_type=SectionState,
)


@agent_for_gen_query_in_report_planner.system_prompt
async def gen_query_in_report_planner_system_prompt(ctx: RunContext[ReportState]):
    return report_planner_query_writer_instructions.format(
        topic=ctx.deps.topic,
        report_organization=settings.report_structure,
        number_of_queries=settings.number_of_queries,
        feedback=ctx.deps.feedback_on_report_plan,
    )


@agent_for_gen_section_in_report_planner.system_prompt
async def gen_section_in_report_planner_system_prompt(ctx: RunContext[ReportState]):
    return report_planner_instructions.format(
        topic=ctx.deps.topic,
        report_organization=settings.report_structure,
        feedback=ctx.deps.feedback_on_report_plan,
    )


@agent_for_gen_query.system_prompt
async def gen_query_system_prompt(ctx: RunContext[SectionState]):
    return query_writer_instructions.format(
        section_topic=ctx.deps.section.description,
        number_of_queries=settings.number_of_queries,
    )


@agent_for_gen_section.system_prompt
async def gen_section_system_prompt(ctx: RunContext[SectionState]):
    return section_writer_instructions.format(
        section_title=ctx.deps.section.name,
        section_topic=ctx.deps.section.description,
        context=ctx.deps.search_results_xml,
        section_content=ctx.deps.section.content,
    )


@agent_for_gen_section_feedback.system_prompt
async def gen_section_feedback_system_prompt(ctx: RunContext[SectionState]):
    return section_grader_instructions.format(
        section_topic=ctx.deps.section.description,
        section=ctx.deps.section.content,
    )


@agent_for_gen_final_section.system_prompt
async def gen_final_section_system_prompt(ctx: RunContext[SectionState]):
    return final_section_writer_instructions.format(
        section_title=ctx.deps.section.name,
        section_topic=ctx.deps.section.description,
        context=ctx.deps.report_sections_from_research,
    )