59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
from pydantic_ai import Agent, RunContext
|
|
|
|
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
|
|
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
|
|
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.utils import get_model
|
|
|
|
|
|
def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisResult]:
|
|
"""Create an analysis agent with code execution capability.
|
|
|
|
The analysis agent can write and execute Python code in a sandboxed
|
|
environment to solve problems that require computation, aggregation,
|
|
or complex traversal across documents.
|
|
|
|
Args:
|
|
config: Application configuration.
|
|
|
|
Returns:
|
|
A pydantic-ai Agent configured for analysis execution.
|
|
"""
|
|
model = get_model(config.analysis.model, config)
|
|
|
|
agent: Agent[AnalysisDeps, RawAnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
|
model,
|
|
deps_type=AnalysisDeps,
|
|
output_type=RawAnalysisResult,
|
|
instructions=ANALYSIS_SYSTEM_PROMPT,
|
|
retries=3,
|
|
)
|
|
|
|
@agent.tool
|
|
async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution:
|
|
"""Execute Python code in a sandboxed interpreter.
|
|
|
|
The code has access to search() and llm() functions, and a
|
|
virtual filesystem at /documents/ with document content and structure.
|
|
|
|
Use print() to output results.
|
|
|
|
Args:
|
|
code: Python code to execute.
|
|
|
|
Returns:
|
|
Structured result with success status, stdout, and stderr.
|
|
"""
|
|
result = await ctx.deps.sandbox.execute(code)
|
|
|
|
execution = CodeExecution(
|
|
code=code,
|
|
stdout=result.stdout,
|
|
stderr=result.stderr,
|
|
success=result.success,
|
|
)
|
|
|
|
return execution
|
|
|
|
return agent
|