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.
|
All agents are built on PydanticAI and registered in a central registry.
|
||||||
"""
|
"""
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Protocol, runtime_checkable
|
from typing import Any, Protocol, runtime_checkable
|
||||||
|
|
||||||
@@ -112,6 +113,22 @@ class BaseAgent(ABC):
|
|||||||
"""Execute the agent."""
|
"""Execute the agent."""
|
||||||
pass
|
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 ===
|
# === Agent Registry ===
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ Fast codebase exploration with read-only tools.
|
|||||||
Uses sanitized Ollama provider for reliable tool calling.
|
Uses sanitized Ollama provider for reliable tool calling.
|
||||||
"""
|
"""
|
||||||
import os
|
import os
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -113,6 +114,34 @@ class ExploreAgentImpl(BaseAgent):
|
|||||||
raise
|
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
|
# Create and register the singleton instance
|
||||||
explore_agent = ExploreAgentImpl()
|
explore_agent = ExploreAgentImpl()
|
||||||
register_agent(explore_agent)
|
register_agent(explore_agent)
|
||||||
@@ -125,3 +154,13 @@ async def explore(
|
|||||||
) -> str:
|
) -> str:
|
||||||
"""Run exploration query."""
|
"""Run exploration query."""
|
||||||
return await explore_agent.run(prompt, working_dir=working_dir, **kwargs)
|
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.
|
REST API routes for agents.
|
||||||
"""
|
"""
|
||||||
|
import json
|
||||||
from fastapi import APIRouter, HTTPException
|
from fastapi import APIRouter, HTTPException
|
||||||
|
from fastapi.responses import StreamingResponse
|
||||||
|
|
||||||
from src.domains.agents.base import get_agent, list_agents
|
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)
|
@router.get("/{agent_type}", response_model=AgentInfo)
|
||||||
async def get_agent_info(agent_type: str) -> AgentInfo:
|
async def get_agent_info(agent_type: str) -> AgentInfo:
|
||||||
"""Get information about a specific agent."""
|
"""Get information about a specific agent."""
|
||||||
|
|||||||
@@ -3,7 +3,9 @@ Webber API client.
|
|||||||
|
|
||||||
Communicates with the Webber API backend for agent execution.
|
Communicates with the Webber API backend for agent execution.
|
||||||
"""
|
"""
|
||||||
|
import json
|
||||||
import httpx
|
import httpx
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -132,6 +134,52 @@ class WebberClient:
|
|||||||
error=data.get("error"),
|
error=data.get("error"),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def run_agent_stream(
|
||||||
|
self,
|
||||||
|
agent_type: str,
|
||||||
|
prompt: str,
|
||||||
|
working_dir: str = ".",
|
||||||
|
) -> AsyncIterator[str]:
|
||||||
|
"""
|
||||||
|
Run an agent with streaming response.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agent_type: Type of agent (e.g., "explore")
|
||||||
|
prompt: User prompt/query
|
||||||
|
working_dir: Working directory for the agent
|
||||||
|
|
||||||
|
Yields:
|
||||||
|
Text chunks as they arrive
|
||||||
|
"""
|
||||||
|
# Use a fresh client for streaming with longer timeout
|
||||||
|
async with httpx.AsyncClient(
|
||||||
|
base_url=self.base_url,
|
||||||
|
timeout=httpx.Timeout(300.0, connect=10.0),
|
||||||
|
) as client:
|
||||||
|
async with client.stream(
|
||||||
|
"POST",
|
||||||
|
"/agents/stream",
|
||||||
|
json={
|
||||||
|
"agent_type": agent_type,
|
||||||
|
"prompt": prompt,
|
||||||
|
"working_dir": working_dir,
|
||||||
|
},
|
||||||
|
) as response:
|
||||||
|
response.raise_for_status()
|
||||||
|
async for line in response.aiter_lines():
|
||||||
|
if line.startswith("data: "):
|
||||||
|
try:
|
||||||
|
data = json.loads(line[6:])
|
||||||
|
event = data.get("event")
|
||||||
|
if event == "chunk":
|
||||||
|
yield data.get("data", "")
|
||||||
|
elif event == "error":
|
||||||
|
raise Exception(data.get("data", "Unknown error"))
|
||||||
|
elif event == "done":
|
||||||
|
break
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
continue
|
||||||
|
|
||||||
async def __aenter__(self) -> "WebberClient":
|
async def __aenter__(self) -> "WebberClient":
|
||||||
"""Async context manager entry."""
|
"""Async context manager entry."""
|
||||||
return self
|
return self
|
||||||
|
|||||||
@@ -9,9 +9,11 @@ Usage:
|
|||||||
"""
|
"""
|
||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import typer
|
import typer
|
||||||
|
from rich.live import Live
|
||||||
from rich.markdown import Markdown
|
from rich.markdown import Markdown
|
||||||
from rich.panel import Panel
|
from rich.panel import Panel
|
||||||
|
|
||||||
@@ -74,6 +76,12 @@ def chat(
|
|||||||
"--agent",
|
"--agent",
|
||||||
help="Agent to use",
|
help="Agent to use",
|
||||||
),
|
),
|
||||||
|
stream: bool = typer.Option(
|
||||||
|
True,
|
||||||
|
"--stream/--no-stream",
|
||||||
|
"-s",
|
||||||
|
help="Stream responses in real-time",
|
||||||
|
),
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Start interactive chat session.
|
Start interactive chat session.
|
||||||
@@ -87,12 +95,14 @@ def chat(
|
|||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
asyncio.run(_chat_loop(api_url, agent, working_dir))
|
asyncio.run(_chat_loop(api_url, agent, working_dir, stream))
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
console.print("\n[dim]Goodbye![/]")
|
console.print("\n[dim]Goodbye![/]")
|
||||||
|
|
||||||
|
|
||||||
async def _chat_loop(api_url: str, agent_type: str, working_dir: str) -> None:
|
async def _chat_loop(
|
||||||
|
api_url: str, agent_type: str, working_dir: str, stream: bool = True
|
||||||
|
) -> None:
|
||||||
"""Interactive chat loop."""
|
"""Interactive chat loop."""
|
||||||
theme = get_theme()
|
theme = get_theme()
|
||||||
|
|
||||||
@@ -118,6 +128,8 @@ async def _chat_loop(api_url: str, agent_type: str, working_dir: str) -> None:
|
|||||||
console.print(f"[title]Webber CLI[/] [dim]→ {api_url}[/]")
|
console.print(f"[title]Webber CLI[/] [dim]→ {api_url}[/]")
|
||||||
console.print(f"[dim]Working in:[/] [path]{working_dir}[/]")
|
console.print(f"[dim]Working in:[/] [path]{working_dir}[/]")
|
||||||
console.print(f"[dim]Agent:[/] {agent_info.name} - {agent_info.description}")
|
console.print(f"[dim]Agent:[/] {agent_info.name} - {agent_info.description}")
|
||||||
|
mode = "streaming" if stream else "batch"
|
||||||
|
console.print(f"[dim]Mode:[/] {mode}")
|
||||||
console.print()
|
console.print()
|
||||||
console.print("[dim]Type 'exit' to quit, 'clear' to clear screen.[/]")
|
console.print("[dim]Type 'exit' to quit, 'clear' to clear screen.[/]")
|
||||||
console.print()
|
console.print()
|
||||||
@@ -148,15 +160,31 @@ async def _chat_loop(api_url: str, agent_type: str, working_dir: str) -> None:
|
|||||||
console.print(f"[error]Directory not found:[/] {new_dir}")
|
console.print(f"[error]Directory not found:[/] {new_dir}")
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Call the API
|
|
||||||
with console.status("[info]Thinking...[/]", spinner=theme.spinner):
|
|
||||||
result = await client.run_agent(agent_type, user_input, working_dir)
|
|
||||||
|
|
||||||
console.print()
|
console.print()
|
||||||
if result.success:
|
|
||||||
console.print(Markdown(result.response))
|
if stream:
|
||||||
|
# Stream response in real-time
|
||||||
|
full_response = ""
|
||||||
|
try:
|
||||||
|
async for chunk in client.run_agent_stream(
|
||||||
|
agent_type, user_input, working_dir
|
||||||
|
):
|
||||||
|
sys.stdout.write(chunk)
|
||||||
|
sys.stdout.flush()
|
||||||
|
full_response += chunk
|
||||||
|
console.print() # Newline after streaming
|
||||||
|
except Exception as e:
|
||||||
|
console.print(f"\n[error]Stream error:[/] {e}")
|
||||||
else:
|
else:
|
||||||
console.print(f"[error]Error:[/] {result.error}")
|
# Batch mode with spinner
|
||||||
|
with console.status("[info]Thinking...[/]", spinner=theme.spinner):
|
||||||
|
result = await client.run_agent(agent_type, user_input, working_dir)
|
||||||
|
|
||||||
|
if result.success:
|
||||||
|
console.print(Markdown(result.response))
|
||||||
|
else:
|
||||||
|
console.print(f"[error]Error:[/] {result.error}")
|
||||||
|
|
||||||
console.print()
|
console.print()
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
@@ -181,6 +209,12 @@ def explore(
|
|||||||
"-a",
|
"-a",
|
||||||
help="Webber API URL",
|
help="Webber API URL",
|
||||||
),
|
),
|
||||||
|
stream: bool = typer.Option(
|
||||||
|
True,
|
||||||
|
"--stream/--no-stream",
|
||||||
|
"-s",
|
||||||
|
help="Stream responses in real-time",
|
||||||
|
),
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
One-shot codebase exploration.
|
One-shot codebase exploration.
|
||||||
@@ -193,10 +227,12 @@ def explore(
|
|||||||
console.print(f"[error]Error:[/] Directory not found: {working_dir}")
|
console.print(f"[error]Error:[/] Directory not found: {working_dir}")
|
||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
|
|
||||||
asyncio.run(_explore(api_url, query, working_dir))
|
asyncio.run(_explore(api_url, query, working_dir, stream))
|
||||||
|
|
||||||
|
|
||||||
async def _explore(api_url: str, query: str, working_dir: str) -> None:
|
async def _explore(
|
||||||
|
api_url: str, query: str, working_dir: str, stream: bool = True
|
||||||
|
) -> None:
|
||||||
"""Execute exploration query."""
|
"""Execute exploration query."""
|
||||||
theme = get_theme()
|
theme = get_theme()
|
||||||
|
|
||||||
@@ -211,21 +247,32 @@ async def _explore(api_url: str, query: str, working_dir: str) -> None:
|
|||||||
console.print(f"[dim]Query:[/] {query}")
|
console.print(f"[dim]Query:[/] {query}")
|
||||||
console.print()
|
console.print()
|
||||||
|
|
||||||
with console.status("[info]Searching...[/]", spinner=theme.spinner):
|
if stream:
|
||||||
result = await client.run_agent("explore", query, working_dir)
|
# Stream response in real-time
|
||||||
|
try:
|
||||||
if result.success:
|
async for chunk in client.run_agent_stream("explore", query, working_dir):
|
||||||
console.print(Panel(
|
sys.stdout.write(chunk)
|
||||||
Markdown(result.response),
|
sys.stdout.flush()
|
||||||
title="[success]Findings[/]",
|
console.print() # Newline after streaming
|
||||||
border_style=theme.colors.border_success,
|
except Exception as e:
|
||||||
))
|
console.print(f"\n[error]Stream error:[/] {e}")
|
||||||
else:
|
else:
|
||||||
console.print(Panel(
|
# Batch mode with spinner
|
||||||
f"[error]{result.error}[/]",
|
with console.status("[info]Searching...[/]", spinner=theme.spinner):
|
||||||
title="[error]Error[/]",
|
result = await client.run_agent("explore", query, working_dir)
|
||||||
border_style=theme.colors.border_error,
|
|
||||||
))
|
if result.success:
|
||||||
|
console.print(Panel(
|
||||||
|
Markdown(result.response),
|
||||||
|
title="[success]Findings[/]",
|
||||||
|
border_style=theme.colors.border_success,
|
||||||
|
))
|
||||||
|
else:
|
||||||
|
console.print(Panel(
|
||||||
|
f"[error]{result.error}[/]",
|
||||||
|
title="[error]Error[/]",
|
||||||
|
border_style=theme.colors.border_error,
|
||||||
|
))
|
||||||
|
|
||||||
|
|
||||||
@app.command()
|
@app.command()
|
||||||
|
|||||||
Reference in New Issue
Block a user