ai-lawyer-agent/backend/core/streaming.py

16 lines
700 B
Python

from typing import AsyncGenerator
from agents import Agent, Runner
from openai.types.responses import ResponseTextDeltaEvent
async def stream_response(agent: Agent, prompt: list[dict] | str) -> AsyncGenerator[str, None]:
"""Stream agent response and update the UI."""
try:
async with agent.mcp_servers[0]:
result = Runner.run_streamed(agent, input=prompt)
async for event in result.stream_events():
if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
yield event.data.delta # <-- sends the next piece of response text
except Exception as e:
yield f"⚠️🖨 Error: {e}"