Add SSE streaming for chat responses #8

README.md

@@ -46,6 +46,7 @@ The platform is composed of three services:
 | Method | Route | Purpose |
 |--------|-------|---------|
 | `POST` | `/chat` | Start or continue a chat session |
+| `POST` | `/chat/stream` | Stream chat responses as Server-Sent Events |
 | `DELETE` | `/chat/{chat_id}` | Delete a chat session |
 | `GET` | `/memory/{user_id}` | Retrieve a stored user profile (transparency) |
 | `DELETE` | `/memory/{user_id}` | Completely erase a stored user profile |
@@ -87,6 +88,17 @@ uvicorn main:app --reload
 
 The API runs at `http://localhost:8000` — interactive docs at `http://localhost:8000/docs`.
 
+### Streaming chat endpoint
+
+Use the streaming endpoint to receive incremental chat output as Server-Sent Events:
+
+```bash
+curl -N -X POST http://localhost:8000/chat/stream \
+  -H 'Content-Type: application/json' \
+  -d '{"message": "Explain the concept of tawakkul in Islam", "chat_id": "demo"}'
+```
+
+Each `delta` event contains a JSON payload like `{"delta": "..."}` and the terminal `done` event contains the `chat_id` and the completed text.
 ### Docker
 
 The included `Dockerfile` produces a production-ready image with Python 3.12,

main.py

@@ -1,3 +1,9 @@
+import json
+import os
+from typing import List, Optional, Dict, Any
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 import asyncio
 import json
 import logging
@@ -946,6 +952,55 @@ async def event_generator():
         raise HTTPException(status_code=500, detail="Internal server error") from e
 
 
+@app.post("/chat/stream")
+async def chat_stream(request: ChatRequest):
+    if not GEMINI_API_KEY:
+        raise HTTPException(status_code=500, detail="GEMINI_API_KEY is not configured.")
+
+    chat_id = request.chat_id or "default"
+    if chat_id not in sessions:
+        model = get_model()
+        sessions[chat_id] = model.start_chat(history=[])
+
+    chat_session = sessions[chat_id]
+
+    async def event_stream():
+        full_text_parts: List[str] = []
+        try:
+            response_stream = None
+            if hasattr(chat_session, "send_message_async"):
+                response_stream = await chat_session.send_message_async(request.message, stream=True)
+            else:
+                response_stream = chat_session.send_message(request.message, stream=True)
+
+            async for chunk in response_stream:
+                text = getattr(chunk, "text", None)
+                if text is None and isinstance(chunk, dict):
+                    text = chunk.get("text")
+                if text is None:
+                    continue
+
+                full_text_parts.append(text)
+                payload = json.dumps({"delta": text})
+                yield f"event: delta\ndata: {payload}\n\n"
+        except Exception as exc:  # pragma: no cover - defensive path
+            error_payload = json.dumps({"error": str(exc)})
+            yield f"event: error\ndata: {error_payload}\n\n"
+            return
+
+        final_payload = json.dumps({"chat_id": chat_id, "text": "".join(full_text_parts)})
+        yield f"event: done\ndata: {final_payload}\n\n"
+
+    return StreamingResponse(
+        event_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
 @app.delete("/chat/{chat_id}")
 async def delete_chat(chat_id: str):
     try:

requirements.txt

@@ -4,6 +4,9 @@ python-dotenv==0.21.0
 google-generativeai==0.8.3
 pydantic>=2.5.3
 stellar-sdk>=12.0.0
+pytest
+pytest-asyncio
+httpx
 PyYAML==6.0.2
 numpy==2.2.4
 redis>=5.0.0,<6.0.0

tests/test_chat_streaming.py

@@ -0,0 +1,70 @@
+import pytest
+from fastapi.testclient import TestClient
+
+import main
+
+
+class FakeChunk:
+    def __init__(self, text: str):
+        self.text = text
+
+
+class FakeChatSession:
+    def __init__(self):
+        self.history = []
+
+    async def send_message_async(self, message, stream=True):
+        self.history.append({"role": "user", "text": message})
+
+        async def generator():
+            yield FakeChunk("Hello")
+            yield FakeChunk(" world")
+
+        return generator()
+
+
+class FailingChatSession(FakeChatSession):
+    async def send_message_async(self, message, stream=True):
+        self.history.append({"role": "user", "text": message})
+
+        async def generator():
+            yield FakeChunk("Hello")
+            raise RuntimeError("upstream failed")
+
+        return generator()
+
+
+@pytest.fixture(autouse=True)
+def reset_app_state(monkeypatch):
+    main.sessions.clear()
+    monkeypatch.setattr(main, "GEMINI_API_KEY", "fake-key")
+    monkeypatch.setattr(main, "CITATION_VERIFY_MODE", "off")
+    monkeypatch.setattr(main, "get_model", lambda: None)
+
+
+def test_chat_stream_endpoint_emits_sse_events(monkeypatch):
+    chat_session = FakeChatSession()
+    monkeypatch.setattr(main, "get_model", lambda: type("Model", (), {"start_chat": lambda self, history=[]: chat_session})())
+
+    client = TestClient(main.app)
+    response = client.post("/chat/stream", json={"message": "Hello", "chat_id": "stream-test"})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+    assert "event: delta" in response.text
+    assert "Hello" in response.text
+    assert "event: done" in response.text
+    assert '"chat_id": "stream-test"' in response.text
+
+
+def test_chat_stream_endpoint_emits_error_event_on_upstream_failure(monkeypatch):
+    chat_session = FailingChatSession()
+    monkeypatch.setattr(main, "get_model", lambda: type("Model", (), {"start_chat": lambda self, history=[]: chat_session})())
+
+    client = TestClient(main.app)
+    response = client.post("/chat/stream", json={"message": "Hello", "chat_id": "broken-stream"})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+    assert "event: error" in response.text
+    assert "upstream failed" in response.text