feat: add streaming responses to API and CLI
Add real-time streaming support for agent responses using Server-Sent Events (SSE). Responses now appear as they're generated instead of waiting for completion. - Add run_stream method to BaseAgent and ExploreAgentImpl - Add /agents/stream SSE endpoint to API router - Add run_agent_stream method to CLI client - Add --stream flag to chat and explore commands (enabled by default) - Use --no-stream for batch mode with spinner Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -4,6 +4,7 @@ Base classes and registry for agent implementations.
|
||||
All agents are built on PydanticAI and registered in a central registry.
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
@@ -112,6 +113,22 @@ class BaseAgent(ABC):
|
||||
"""Execute the agent."""
|
||||
pass
|
||||
|
||||
async def run_stream(
|
||||
self, prompt: str, **kwargs: Any
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Execute the agent with streaming output.
|
||||
|
||||
Default implementation falls back to non-streaming run().
|
||||
Override this for true streaming support.
|
||||
|
||||
Yields:
|
||||
Text chunks as they become available
|
||||
"""
|
||||
# Default: fall back to non-streaming
|
||||
result = await self.run(prompt, **kwargs)
|
||||
yield result
|
||||
|
||||
|
||||
# === Agent Registry ===
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ Fast codebase exploration with read-only tools.
|
||||
Uses sanitized Ollama provider for reliable tool calling.
|
||||
"""
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
@@ -113,6 +114,34 @@ class ExploreAgentImpl(BaseAgent):
|
||||
raise
|
||||
|
||||
|
||||
async def run_stream(
|
||||
self,
|
||||
prompt: str,
|
||||
working_dir: str | None = None,
|
||||
allowed_paths: list[str] | None = None,
|
||||
**kwargs: Any
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Run the explore agent with streaming output.
|
||||
|
||||
Yields text chunks as they become available.
|
||||
"""
|
||||
ctx = ExploreContext(
|
||||
working_dir=working_dir or os.getcwd(),
|
||||
allowed_paths=allowed_paths or self._settings.effective_allowed_paths,
|
||||
timeout_seconds=self._settings.tool_timeout_seconds,
|
||||
)
|
||||
|
||||
async with trace_span("explore_agent_stream"):
|
||||
try:
|
||||
async with self.agent.run_stream(prompt, deps=ctx) as result:
|
||||
async for chunk in result.stream_text():
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.exception(f"Explore agent stream error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
# Create and register the singleton instance
|
||||
explore_agent = ExploreAgentImpl()
|
||||
register_agent(explore_agent)
|
||||
@@ -125,3 +154,13 @@ async def explore(
|
||||
) -> str:
|
||||
"""Run exploration query."""
|
||||
return await explore_agent.run(prompt, working_dir=working_dir, **kwargs)
|
||||
|
||||
|
||||
async def explore_stream(
|
||||
prompt: str,
|
||||
working_dir: str | None = None,
|
||||
**kwargs: Any
|
||||
) -> AsyncIterator[str]:
|
||||
"""Run exploration query with streaming."""
|
||||
async for chunk in explore_agent.run_stream(prompt, working_dir=working_dir, **kwargs):
|
||||
yield chunk
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
"""
|
||||
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
|
||||
|
||||
@@ -68,6 +70,53 @@ async def run_agent(request: AgentRunRequest) -> AgentRunResponse:
|
||||
)
|
||||
|
||||
|
||||
@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."""
|
||||
|
||||
Reference in New Issue
Block a user