Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions datamind/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -235,12 +235,19 @@ async def chat(
async def stream() -> AsyncIterator[bytes]:
context = RequestContext(session_id=session, profile=st.settings.data.profile)
with bind_context(context):
async for event in st.system.retrieve.loop.stream_turn(
user_message=req.message,
history=req.history or [],
):
try:
async for event in st.system.retrieve.loop.stream_turn(
user_message=req.message,
history=req.history or [],
):
payload = json.dumps(
{"type": event.type, **event.data},
ensure_ascii=False,
)
yield f"data: {payload}\n\n".encode("utf-8")
except Exception as exc: # noqa: BLE001 - convert provider errors to SSE
payload = json.dumps(
{"type": event.type, **event.data},
{"type": "error", "message": f"stream failed: {type(exc).__name__}"},
ensure_ascii=False,
)
yield f"data: {payload}\n\n".encode("utf-8")
Expand Down
21 changes: 21 additions & 0 deletions datamind/tests/test_http_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,12 @@ async def stream_turn(self, *, user_message: str, history=None, final_contract=N
})


class _FailingStreamLoop(_FakeLoop):
async def stream_turn(self, *, user_message: str, history=None, final_contract=None):
yield AgentEvent("text", {"delta": "partial"})
raise RuntimeError("provider unavailable")


class _FakeKB:
async def list_documents(self):
return [{"source": "demo.md", "chunks": 1}]
Expand Down Expand Up @@ -197,3 +203,18 @@ async def test_http_api_routes_round_trip(configured_app):

graph = await client.get("/api/graph/stats")
assert graph.json() == {"nodes": 2, "edges": 1}


@pytest.mark.asyncio
async def test_chat_converts_stream_failure_to_sse_error(configured_app):
test_app, _settings = configured_app
test_app.state.datamind.system.retrieve.loop = _FailingStreamLoop("retrieve")
transport = httpx.ASGITransport(app=test_app, raise_app_exceptions=True)

async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
response = await client.post("/api/chat", json={"message": "stream this"})

assert response.status_code == 200
assert '"type": "text"' in response.text
assert '"type": "error"' in response.text
assert '"message": "stream failed: RuntimeError"' in response.text