diff --git a/webber-api/src/domains/agents/base.py b/webber-api/src/domains/agents/base.py index d834928..592a4fd 100644 --- a/webber-api/src/domains/agents/base.py +++ b/webber-api/src/domains/agents/base.py @@ -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 === diff --git a/webber-api/src/domains/agents/explore/agent.py b/webber-api/src/domains/agents/explore/agent.py index a183520..ddb023a 100644 --- a/webber-api/src/domains/agents/explore/agent.py +++ b/webber-api/src/domains/agents/explore/agent.py @@ -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 diff --git a/webber-api/src/domains/agents/router.py b/webber-api/src/domains/agents/router.py index c7b2ad6..f32fd3b 100644 --- a/webber-api/src/domains/agents/router.py +++ b/webber-api/src/domains/agents/router.py @@ -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.""" diff --git a/webber-cli/webber_cli/client.py b/webber-cli/webber_cli/client.py index 6aac015..d5ff1a5 100644 --- a/webber-cli/webber_cli/client.py +++ b/webber-cli/webber_cli/client.py @@ -3,7 +3,9 @@ Webber API client. Communicates with the Webber API backend for agent execution. """ +import json import httpx +from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any @@ -132,6 +134,52 @@ class WebberClient: 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 context manager entry.""" return self diff --git a/webber-cli/webber_cli/main.py b/webber-cli/webber_cli/main.py index 979c1d0..277f427 100644 --- a/webber-cli/webber_cli/main.py +++ b/webber-cli/webber_cli/main.py @@ -9,9 +9,11 @@ Usage: """ import asyncio import os +import sys from pathlib import Path import typer +from rich.live import Live from rich.markdown import Markdown from rich.panel import Panel @@ -74,6 +76,12 @@ def chat( "--agent", help="Agent to use", ), + stream: bool = typer.Option( + True, + "--stream/--no-stream", + "-s", + help="Stream responses in real-time", + ), ) -> None: """ Start interactive chat session. @@ -87,12 +95,14 @@ def chat( raise typer.Exit(1) try: - asyncio.run(_chat_loop(api_url, agent, working_dir)) + asyncio.run(_chat_loop(api_url, agent, working_dir, stream)) except KeyboardInterrupt: 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.""" 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"[dim]Working in:[/] [path]{working_dir}[/]") 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("[dim]Type 'exit' to quit, 'clear' to clear screen.[/]") 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}") 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() - 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: - 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() except KeyboardInterrupt: @@ -181,6 +209,12 @@ def explore( "-a", help="Webber API URL", ), + stream: bool = typer.Option( + True, + "--stream/--no-stream", + "-s", + help="Stream responses in real-time", + ), ) -> None: """ One-shot codebase exploration. @@ -193,10 +227,12 @@ def explore( console.print(f"[error]Error:[/] Directory not found: {working_dir}") 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.""" 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() - with console.status("[info]Searching...[/]", spinner=theme.spinner): - result = await client.run_agent("explore", query, working_dir) - - if result.success: - console.print(Panel( - Markdown(result.response), - title="[success]Findings[/]", - border_style=theme.colors.border_success, - )) + if stream: + # Stream response in real-time + try: + async for chunk in client.run_agent_stream("explore", query, working_dir): + sys.stdout.write(chunk) + sys.stdout.flush() + console.print() # Newline after streaming + except Exception as e: + console.print(f"\n[error]Stream error:[/] {e}") else: - console.print(Panel( - f"[error]{result.error}[/]", - title="[error]Error[/]", - border_style=theme.colors.border_error, - )) + # Batch mode with spinner + with console.status("[info]Searching...[/]", spinner=theme.spinner): + result = await client.run_agent("explore", query, working_dir) + + 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()