mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 19:02:16 +00:00
feat: add SSE streaming support for managed agent messages
Add `stream: bool` parameter to `POST /v1/managed-agents/{id}/messages`.
When `stream=true`, the agent processes the message synchronously and
returns an SSE stream (OpenAI-compatible format) with token-by-token
response, tool result events, and usage metadata.
This enables real-time voice assistants and chat UIs to receive agent
responses as they are generated, rather than polling for completion.
- Extend SendMessageRequest with `stream` field (default: false)
- Add _stream_managed_agent() helper using asyncio.to_thread()
- Build AgentContext from conversation history for multi-turn support
- Emit tool_results as named SSE events
- Persist agent response in DB after streaming completes
- Add 6 new tests covering streaming behavior
- Update agents.md documentation with streaming examples
This commit is contained in:
@@ -621,3 +621,49 @@ All agents publish events on the `EventBus` when a bus is provided:
|
||||
`INFERENCE_START` / `INFERENCE_END` events are published by the `InstrumentedEngine` wrapper, not by agents directly. This keeps telemetry opt-in and transparent to agent code.
|
||||
|
||||
These events enable the telemetry and trace systems to record detailed interaction data automatically.
|
||||
|
||||
---
|
||||
|
||||
## Managed Agent Streaming
|
||||
|
||||
The Managed Agent API (`/v1/managed-agents/{id}/messages`) supports real-time SSE streaming. Send a message with `stream: true` to receive the agent's response as a Server-Sent Events stream instead of the default asynchronous queue mode.
|
||||
|
||||
### Streaming Messages
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:8000/v1/managed-agents/{id}/messages \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"content": "What is 2+2?", "stream": true}'
|
||||
```
|
||||
|
||||
The response follows the OpenAI SSE format:
|
||||
|
||||
1. **Content chunks** -- `data: {"choices": [{"delta": {"content": "token"}}]}`
|
||||
2. **Tool results** (if the agent used tools) -- `event: tool_results\ndata: {"results": [...]}`
|
||||
3. **Final chunk** -- `data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}`
|
||||
4. **Done sentinel** -- `data: [DONE]`
|
||||
|
||||
When `stream: false` (the default), the endpoint behaves exactly as before -- the message is queued and the agent must be triggered separately via `/run`.
|
||||
|
||||
### Behavior Details
|
||||
|
||||
- The user message is always stored in the database before the agent runs.
|
||||
- After streaming completes, the full agent response is persisted as an `agent_to_user` message.
|
||||
- The agent is instantiated from the managed agent's stored `agent_type` and `config`.
|
||||
- Conversation history from prior messages is automatically loaded as context.
|
||||
- If the engine is not available on the server, a `503` error is returned.
|
||||
|
||||
### Python Example
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
with httpx.stream(
|
||||
"POST",
|
||||
"http://localhost:8000/v1/managed-agents/{id}/messages",
|
||||
json={"content": "Summarize today's news", "stream": True},
|
||||
) as response:
|
||||
for line in response.iter_lines():
|
||||
if line.startswith("data:") and "[DONE]" not in line:
|
||||
print(line[5:].strip())
|
||||
```
|
||||
|
||||
@@ -2,16 +2,20 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents.manager import AgentManager
|
||||
|
||||
try:
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
except ImportError:
|
||||
raise ImportError("fastapi and pydantic are required for server routes")
|
||||
|
||||
logger = logging.getLogger("openjarvis.server.agent_manager")
|
||||
|
||||
|
||||
class CreateAgentRequest(BaseModel):
|
||||
name: str
|
||||
@@ -46,6 +50,7 @@ class BindChannelRequest(BaseModel):
|
||||
class SendMessageRequest(BaseModel):
|
||||
content: str
|
||||
mode: str = "queued"
|
||||
stream: bool = False # SSE streaming mode
|
||||
|
||||
|
||||
class FeedbackRequest(BaseModel):
|
||||
@@ -202,6 +207,171 @@ def build_tools_list() -> List[Dict[str, Any]]:
|
||||
return items
|
||||
|
||||
|
||||
async def _stream_managed_agent(
|
||||
*,
|
||||
manager: AgentManager,
|
||||
agent_record: Dict[str, Any],
|
||||
user_content: str,
|
||||
message_id: str,
|
||||
engine: Any,
|
||||
bus: Any,
|
||||
) -> StreamingResponse:
|
||||
"""Run a managed agent and stream the response as SSE.
|
||||
|
||||
Instantiates the agent from its stored config, builds conversation
|
||||
context from message history, executes the agent in a background
|
||||
thread, and yields SSE-formatted chunks. After completion the
|
||||
full response is persisted via ``manager.store_agent_response()``.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
from openjarvis.core.types import Message, Role
|
||||
|
||||
agent_id = agent_record["id"]
|
||||
config = agent_record.get("config", {})
|
||||
agent_type = agent_record.get("agent_type", "orchestrator")
|
||||
model = config.get("model", getattr(engine, "_model", ""))
|
||||
|
||||
# Resolve the agent class from registry
|
||||
agent_cls = AgentRegistry.get(agent_type)
|
||||
if agent_cls is None:
|
||||
# Fallback to orchestrator if the type is not registered
|
||||
agent_cls = AgentRegistry.get("orchestrator")
|
||||
if agent_cls is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Agent type '{agent_type}' not found in registry",
|
||||
)
|
||||
|
||||
# Build agent constructor kwargs from config
|
||||
agent_kwargs: Dict[str, Any] = {
|
||||
"engine": engine,
|
||||
"model": model,
|
||||
}
|
||||
if bus is not None:
|
||||
agent_kwargs["bus"] = bus
|
||||
if config.get("system_prompt"):
|
||||
agent_kwargs["system_prompt"] = config["system_prompt"]
|
||||
if config.get("temperature") is not None:
|
||||
agent_kwargs["temperature"] = config["temperature"]
|
||||
if config.get("max_tokens") is not None:
|
||||
agent_kwargs["max_tokens"] = config["max_tokens"]
|
||||
if config.get("max_turns") is not None:
|
||||
agent_kwargs["max_turns"] = config["max_turns"]
|
||||
|
||||
try:
|
||||
agent = agent_cls(**agent_kwargs)
|
||||
except TypeError as exc:
|
||||
logger.warning("Agent instantiation failed with all kwargs, retrying minimal: %s", exc)
|
||||
agent = agent_cls(engine=engine, model=model)
|
||||
|
||||
# Build conversation context from existing messages
|
||||
ctx = AgentContext()
|
||||
messages = manager.list_messages(agent_id, limit=50)
|
||||
# Messages come in DESC order, reverse for chronological
|
||||
for m in reversed(messages):
|
||||
# Skip the message we just stored (it will be the input)
|
||||
if m["id"] == message_id:
|
||||
continue
|
||||
if m["direction"] == "user_to_agent":
|
||||
ctx.conversation.add(Message(role=Role.USER, content=m["content"]))
|
||||
elif m["direction"] == "agent_to_user":
|
||||
ctx.conversation.add(Message(role=Role.ASSISTANT, content=m["content"]))
|
||||
|
||||
# Mark the user message as delivered
|
||||
manager.mark_message_delivered(message_id)
|
||||
|
||||
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
async def generate():
|
||||
"""Async generator yielding SSE-formatted chunks."""
|
||||
collected_content = ""
|
||||
|
||||
# Run agent.run() in a background thread
|
||||
try:
|
||||
result = await asyncio.to_thread(agent.run, user_content, context=ctx)
|
||||
except Exception as exc:
|
||||
logger.error("Managed agent stream error: %s", exc, exc_info=True)
|
||||
error_data = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": f"Error: {exc}"},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
}
|
||||
yield f"data: {json.dumps(error_data)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return
|
||||
|
||||
content = result.content or ""
|
||||
collected_content = content
|
||||
|
||||
# Emit tool results metadata if any
|
||||
if result.tool_results:
|
||||
tool_data = []
|
||||
for tr in result.tool_results:
|
||||
tool_data.append({
|
||||
"tool_name": tr.tool_name,
|
||||
"success": tr.success,
|
||||
"output": tr.content,
|
||||
"latency_ms": tr.latency_seconds * 1000,
|
||||
})
|
||||
yield f"event: tool_results\ndata: {json.dumps({'results': tool_data})}\n\n"
|
||||
|
||||
# Stream content word-by-word for real-time feel
|
||||
if content:
|
||||
words = content.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
token = word if i == 0 else " " + word
|
||||
chunk_data = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": token},
|
||||
"finish_reason": None,
|
||||
}],
|
||||
}
|
||||
yield f"data: {json.dumps(chunk_data)}\n\n"
|
||||
await asyncio.sleep(0.012)
|
||||
|
||||
# Final chunk with finish_reason
|
||||
final_data = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
}
|
||||
yield f"data: {json.dumps(final_data)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
# Persist agent response in DB after streaming completes
|
||||
if collected_content:
|
||||
try:
|
||||
manager.store_agent_response(agent_id, collected_content)
|
||||
except Exception as store_exc:
|
||||
logger.error(
|
||||
"Failed to store agent response: %s", store_exc, exc_info=True,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
|
||||
)
|
||||
|
||||
|
||||
def create_agent_manager_router(
|
||||
manager: AgentManager,
|
||||
) -> Tuple[APIRouter, APIRouter, APIRouter, APIRouter]:
|
||||
@@ -399,9 +569,34 @@ def create_agent_manager_router(
|
||||
return {"messages": manager.list_messages(agent_id)}
|
||||
|
||||
@agents_router.post("/{agent_id}/messages")
|
||||
def send_message(agent_id: str, req: SendMessageRequest):
|
||||
async def send_message(agent_id: str, req: SendMessageRequest, request: Request):
|
||||
agent_record = manager.get_agent(agent_id)
|
||||
if not agent_record:
|
||||
raise HTTPException(status_code=404, detail="Agent not found")
|
||||
|
||||
# Store user message in DB (always, regardless of stream mode)
|
||||
msg = manager.send_message(agent_id, req.content, mode=req.mode)
|
||||
return msg
|
||||
|
||||
if not req.stream:
|
||||
return msg
|
||||
|
||||
# --- Streaming mode: run agent and return SSE response ---
|
||||
engine = getattr(request.app.state, "engine", None)
|
||||
bus = getattr(request.app.state, "bus", None)
|
||||
if engine is None:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Engine not available for streaming",
|
||||
)
|
||||
|
||||
return await _stream_managed_agent(
|
||||
manager=manager,
|
||||
agent_record=agent_record,
|
||||
user_content=req.content,
|
||||
message_id=msg["id"],
|
||||
engine=engine,
|
||||
bus=bus,
|
||||
)
|
||||
|
||||
# ── State inspection ─────────────────────────────────────
|
||||
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -200,3 +202,168 @@ class TestAgentManagerRoutes:
|
||||
assert "channels" in state
|
||||
assert "messages" in state
|
||||
assert "checkpoint" in state
|
||||
|
||||
def test_send_message_non_stream_unchanged(self, manager, client):
|
||||
"""stream=False (default) returns a normal JSON message, not SSE."""
|
||||
agent = manager.create_agent(name="basic", agent_type="simple")
|
||||
res = client.post(
|
||||
f"/v1/managed-agents/{agent['id']}/messages",
|
||||
json={"content": "hello", "stream": False},
|
||||
)
|
||||
assert res.status_code == 200
|
||||
data = res.json()
|
||||
assert data["content"] == "hello"
|
||||
assert data["direction"] == "user_to_agent"
|
||||
|
||||
def test_send_message_stream_not_found(self, manager, client):
|
||||
"""Streaming to a non-existent agent returns 404."""
|
||||
res = client.post(
|
||||
"/v1/managed-agents/nonexistent/messages",
|
||||
json={"content": "hello", "stream": True},
|
||||
)
|
||||
assert res.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed")
|
||||
class TestAgentManagerStreaming:
|
||||
"""Tests for the SSE streaming mode of the managed-agent messages endpoint."""
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_engine(self):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine._model = "test-model"
|
||||
engine.health.return_value = True
|
||||
return engine
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_agent_cls(self):
|
||||
"""Register a mock agent class in the AgentRegistry for testing."""
|
||||
from openjarvis.agents._stubs import AgentResult
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
class _MockStreamAgent:
|
||||
agent_id = "mock_stream"
|
||||
|
||||
def __init__(self, engine, model, **kwargs):
|
||||
self._engine = engine
|
||||
self._model = model
|
||||
|
||||
def run(self, input_text, context=None, **kwargs):
|
||||
return AgentResult(content=f"Echo: {input_text}", turns=1)
|
||||
|
||||
# Register under a unique key for test isolation
|
||||
AgentRegistry._entries()["_test_stream"] = _MockStreamAgent
|
||||
yield _MockStreamAgent
|
||||
AgentRegistry._entries().pop("_test_stream", None)
|
||||
|
||||
@pytest.fixture
|
||||
def stream_client(self, manager, _mock_engine, _mock_agent_cls):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from openjarvis.server.agent_manager_routes import create_agent_manager_router
|
||||
|
||||
app = FastAPI()
|
||||
app.state.engine = _mock_engine
|
||||
app.state.bus = None
|
||||
|
||||
routers = create_agent_manager_router(manager)
|
||||
agents_router, templates_router, global_router, tools_router = routers
|
||||
app.include_router(agents_router)
|
||||
app.include_router(templates_router)
|
||||
app.include_router(global_router)
|
||||
app.include_router(tools_router)
|
||||
return TestClient(app)
|
||||
|
||||
def test_send_message_stream(self, manager, stream_client, _mock_agent_cls):
|
||||
"""Test streaming mode returns SSE response with [DONE] sentinel."""
|
||||
agent = manager.create_agent(
|
||||
name="streamer", agent_type="_test_stream",
|
||||
)
|
||||
resp = stream_client.post(
|
||||
f"/v1/managed-agents/{agent['id']}/messages",
|
||||
json={"content": "What is 2+2?", "stream": True},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert "text/event-stream" in resp.headers.get("content-type", "")
|
||||
|
||||
# Parse SSE events
|
||||
lines = resp.text.strip().split("\n")
|
||||
data_lines = [ln for ln in lines if ln.startswith("data:")]
|
||||
assert len(data_lines) > 0
|
||||
# Last data line must be [DONE]
|
||||
assert data_lines[-1].strip() == "data: [DONE]"
|
||||
|
||||
def test_send_message_stream_content(self, manager, stream_client, _mock_agent_cls):
|
||||
"""Test streaming returns the correct agent response content."""
|
||||
agent = manager.create_agent(
|
||||
name="streamer2", agent_type="_test_stream",
|
||||
)
|
||||
resp = stream_client.post(
|
||||
f"/v1/managed-agents/{agent['id']}/messages",
|
||||
json={"content": "Hello world", "stream": True},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
# Collect content tokens from stream
|
||||
content = ""
|
||||
for line in resp.text.strip().split("\n"):
|
||||
if line.startswith("data:") and "[DONE]" not in line:
|
||||
raw = line[5:].strip()
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
choices = data.get("choices", [{}])
|
||||
delta_content = choices[0].get("delta", {}).get("content")
|
||||
if delta_content:
|
||||
content += delta_content
|
||||
|
||||
assert content == "Echo: Hello world"
|
||||
|
||||
def test_send_message_stream_stores_response(
|
||||
self, manager, stream_client, _mock_agent_cls,
|
||||
):
|
||||
"""After streaming, agent response is persisted in the DB."""
|
||||
agent = manager.create_agent(
|
||||
name="streamer3", agent_type="_test_stream",
|
||||
)
|
||||
resp = stream_client.post(
|
||||
f"/v1/managed-agents/{agent['id']}/messages",
|
||||
json={"content": "persist me", "stream": True},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
# Check messages in DB
|
||||
messages = manager.list_messages(agent["id"])
|
||||
# Should have both the user message and the agent response
|
||||
assert len(messages) == 2
|
||||
directions = {m["direction"] for m in messages}
|
||||
assert "user_to_agent" in directions
|
||||
assert "agent_to_user" in directions
|
||||
agent_msg = next(m for m in messages if m["direction"] == "agent_to_user")
|
||||
assert "persist me" in agent_msg["content"]
|
||||
|
||||
def test_send_message_stream_finish_reason(
|
||||
self, manager, stream_client, _mock_agent_cls,
|
||||
):
|
||||
"""The final chunk before [DONE] has finish_reason='stop'."""
|
||||
agent = manager.create_agent(
|
||||
name="streamer4", agent_type="_test_stream",
|
||||
)
|
||||
resp = stream_client.post(
|
||||
f"/v1/managed-agents/{agent['id']}/messages",
|
||||
json={"content": "check finish", "stream": True},
|
||||
)
|
||||
# Collect all data chunks (excluding [DONE])
|
||||
chunks = []
|
||||
for line in resp.text.strip().split("\n"):
|
||||
if line.startswith("data:") and "[DONE]" not in line:
|
||||
raw = line[5:].strip()
|
||||
try:
|
||||
chunks.append(json.loads(raw))
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# Last chunk should have finish_reason="stop"
|
||||
assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
|
||||
|
||||
Reference in New Issue
Block a user