app/agent.py

from pathlib import Path

from pydantic import BaseModel, Field
from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.messages import (
    ModelMessage,
    ModelRequest,
    ModelResponse,
    TextPart,
    ToolReturnPart,
)
from pydantic_ai.models.google import GoogleModel, GoogleModelSettings

from app import models
from app.tools import dailymed, literature


class Context(BaseModel):
    thoughts: list[str]
    sources: dict[str, dict]


class Statement(BaseModel):
    text: str
    sources: list[str] | None = Field(
        default=None, description="ID of the sources that support this statement."
    )


def get_context(messages: list[ModelMessage]) -> Context:
    thoughts: list[str] = []
    sources: dict[str, dict] = {}

    for message in messages:
        if isinstance(message, ModelResponse):
            for part in message.parts:
                if isinstance(part, TextPart):
                    thoughts.append(part.content)
        elif isinstance(message, ModelRequest):
            for part in message.parts:
                if isinstance(part, ToolReturnPart) and part.tool_name in {
                    "search_medical_literature",
                    "find_drug_set_ids",
                }:
                    for item in part.content:
                        sources[item["id"]] = item

    return Context(thoughts=thoughts, sources=sources)


def create_response(ctx: RunContext, output: list[Statement]) -> models.Statements:
    context = get_context(ctx.messages)

    statements = []
    for statement in output:
        sources = []
        for source_id in statement.sources or []:
            try:
                sources.append(context.sources[source_id])
            except KeyError as err:
                raise ModelRetry(
                    f"Source ID '{source_id}' not found in literature."
                ) from err
        statements.append({"text": statement.text, "sources": sources})

    return models.Statements.model_validate(
        {
            "statements": statements,
            "thoughts": "\n\n".join(context.thoughts),
        }
    )


model = GoogleModel("gemini-2.5-flash-preview-05-20")
settings = GoogleModelSettings(
    # temperature=0.1,
    google_thinking_config={"thinking_budget": 1024, "include_thoughts": True},
)
agent = Agent(
    model=model,
    name="elna",
    model_settings=settings,
    output_type=create_response,
    system_prompt=(Path(__file__).parent / "system_instruction.txt").read_text(),
    tools=[
        dailymed.find_drug_set_ids,
        dailymed.find_drug_instruction,
        literature.search_medical_literature,
    ],
)