legal-ai-assistant/backend/agent/agent.py

59 lines
1.7 KiB
Python

from agents import Agent, OpenAIChatCompletionsModel, AsyncOpenAI, ModelSettings
from agents.mcp import MCPServerStreamableHttp
from backend.agent.sys_prompt import get_system_prompt
from backend.agent.hooks import LegalAssistantHooks
from configs import (
LITELLM_BASE_URL,
LITELLM_API_KEY,
LLM_TIMEOUT,
MCP_SERVER_URL,
AGENT_TEMPERATURE,
DEFAULT_MODEL
)
def make_client() -> AsyncOpenAI:
"""Async client pointing to LiteLLM proxy."""
return AsyncOpenAI (
base_url=LITELLM_BASE_URL,
api_key=LITELLM_API_KEY,
timeout=LLM_TIMEOUT,
max_retries=0
)
def make_mcp_server() -> MCPServerStreamableHttp:
"""MCP server exposing Slovak Justice API tools."""
return MCPServerStreamableHttp(
name="Slovak Justice API",
params={"url": MCP_SERVER_URL},
cache_tools_list=True
)
def build_agent(mcp_server: MCPServerStreamableHttp, model_name: str = DEFAULT_MODEL) -> Agent:
"""
Build and return a Legal AI Assistant agent.
Args:
model_name: Model identifier routed through LiteLLM.
Returns:
Configured Agent ready to run.
"""
return Agent (
name="Legal AI Assistant",
instructions=get_system_prompt(),
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=make_client()
),
model_settings=ModelSettings(
temperature=AGENT_TEMPERATURE,
tool_choice="auto",
parallel_tool_calls=False,
include_usage=True
),
mcp_servers=[mcp_server],
tool_use_behavior="run_llm_again",
reset_tool_choice=True,
hooks=LegalAssistantHooks()
)