136 lines
3.7 KiB
Python
136 lines
3.7 KiB
Python
"""
|
|
REST API routes for agents.
|
|
"""
|
|
import json
|
|
from fastapi import APIRouter, HTTPException
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from src.domains.agents.base import get_agent, list_agents
|
|
|
|
# Import agents to ensure they're registered
|
|
import src.domains.agents.explore # noqa: F401
|
|
import src.domains.agents.plan # noqa: F401
|
|
import src.domains.agents.task # noqa: F401
|
|
from src.domains.agents.schemas import (
|
|
AgentRunRequest,
|
|
AgentRunResponse,
|
|
AgentInfo,
|
|
AgentListResponse,
|
|
)
|
|
from src.shared.logging import logged, get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
router = APIRouter(prefix="/agents", tags=["Agents"])
|
|
|
|
|
|
@router.get("/", response_model=AgentListResponse)
|
|
async def list_available_agents() -> AgentListResponse:
|
|
"""List all available agents."""
|
|
agents = list_agents()
|
|
return AgentListResponse(
|
|
agents=[AgentInfo(**a) for a in agents]
|
|
)
|
|
|
|
|
|
@router.post("/run", response_model=AgentRunResponse)
|
|
@logged()
|
|
async def run_agent(request: AgentRunRequest) -> AgentRunResponse:
|
|
"""
|
|
Run an agent with the given prompt.
|
|
|
|
The agent will use tools to explore the codebase and answer questions.
|
|
"""
|
|
# Get the requested agent
|
|
agent = get_agent(request.agent_type)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Unknown agent type: {request.agent_type}"
|
|
)
|
|
|
|
try:
|
|
# Run the agent
|
|
response = await agent.run(
|
|
request.prompt,
|
|
working_dir=request.working_dir,
|
|
)
|
|
|
|
return AgentRunResponse(
|
|
response=response,
|
|
agent_type=request.agent_type,
|
|
success=True,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Agent execution failed: {e}")
|
|
return AgentRunResponse(
|
|
response="",
|
|
agent_type=request.agent_type,
|
|
success=False,
|
|
error=str(e),
|
|
)
|
|
|
|
|
|
@router.post("/stream")
|
|
@logged()
|
|
async def stream_agent(request: AgentRunRequest) -> StreamingResponse:
|
|
"""
|
|
Run an agent with streaming response.
|
|
|
|
Returns Server-Sent Events (SSE) with text chunks.
|
|
Event types:
|
|
- "chunk": Text chunk from the agent
|
|
- "done": Stream complete
|
|
- "error": Error occurred
|
|
"""
|
|
agent = get_agent(request.agent_type)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Unknown agent type: {request.agent_type}"
|
|
)
|
|
|
|
async def generate():
|
|
try:
|
|
async for chunk in agent.run_stream(
|
|
request.prompt,
|
|
working_dir=request.working_dir,
|
|
):
|
|
# SSE format: data: {json}\n\n
|
|
event = {"event": "chunk", "data": chunk}
|
|
yield f"data: {json.dumps(event)}\n\n"
|
|
|
|
# Signal completion
|
|
yield f"data: {json.dumps({'event': 'done'})}\n\n"
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Stream error: {e}")
|
|
error_event = {"event": "error", "data": str(e)}
|
|
yield f"data: {json.dumps(error_event)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
generate(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
}
|
|
)
|
|
|
|
|
|
@router.get("/{agent_type}", response_model=AgentInfo)
|
|
async def get_agent_info(agent_type: str) -> AgentInfo:
|
|
"""Get information about a specific agent."""
|
|
agent = get_agent(agent_type)
|
|
if not agent:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"Agent not found: {agent_type}"
|
|
)
|
|
|
|
return AgentInfo(
|
|
name=agent.name,
|
|
description=agent.description,
|
|
)
|