feat: stream provider responses

This commit is contained in:
2026-07-30 19:38:53 +02:00
parent 9148193806
commit 00a5a6ce8c
12 changed files with 850 additions and 41 deletions
+49 -1
View File
@@ -4,7 +4,7 @@ from pathlib import Path
from javis.core.chat_service import ChatService
from javis.memory.sqlite_store import SQLiteSessionStore
from javis.providers.base import ChatMessage, LocalModelProvider
from javis.providers.base import ChatMessage, LocalModelProvider, ProviderUnavailableError
class RecordingProvider:
@@ -19,6 +19,19 @@ class RecordingProvider:
return f"Antwort {len(self.calls)}"
class StreamingProvider(RecordingProvider):
def __init__(self, *, fail_after_first: bool = False) -> None:
super().__init__()
self.fail_after_first = fail_after_first
def stream_chat(self, messages: list[ChatMessage]):
self.calls.append(messages)
yield "Teil "
if self.fail_after_first:
raise ProviderUnavailableError("Streamfehler")
yield "Antwort"
class ChatServiceTests(unittest.TestCase):
def setUp(self) -> None:
self.temporary_directory = tempfile.TemporaryDirectory()
@@ -52,6 +65,41 @@ class ChatServiceTests(unittest.TestCase):
],
)
def test_streaming_response_is_persisted_exactly_once_after_completion(self) -> None:
provider = StreamingProvider()
service = ChatService(self.service.store, provider)
session = service.new_session()
chunks = list(service.stream_send(session.id, "Hallo"))
self.assertEqual(chunks, ["Teil ", "Antwort"])
self.assertEqual(
[message.content for message in service.load_session(session.id).messages],
["Hallo", "Teil Antwort"],
)
self.assertEqual(service.load_session(session.id).session.message_count, 2)
def test_closed_stream_does_not_persist_partial_response(self) -> None:
provider = StreamingProvider()
service = ChatService(self.service.store, provider)
session = service.new_session()
stream = service.stream_send(session.id, "Hallo")
self.assertEqual(next(stream), "Teil ")
stream.close()
self.assertEqual(service.load_session(session.id).messages, [])
def test_provider_error_during_stream_does_not_persist_partial_response(self) -> None:
provider = StreamingProvider(fail_after_first=True)
service = ChatService(self.service.store, provider)
session = service.new_session()
with self.assertRaises(ProviderUnavailableError):
list(service.stream_send(session.id, "Hallo"))
self.assertEqual(service.load_session(session.id).messages, [])
if __name__ == "__main__":
unittest.main()
+60 -1
View File
@@ -19,6 +19,22 @@ class _FakeProvider:
def chat(self, messages: list[ChatMessage]) -> str:
return f"Echo: {messages[-1].content}"
def stream_chat(self, messages: list[ChatMessage]):
yield "Echo: "
yield messages[-1].content
class _MarkdownProvider(_FakeProvider):
def stream_chat(self, messages: list[ChatMessage]):
yield "## **Ant"
yield "wort**\n`SQLite`"
class _InterruptingProvider(_FakeProvider):
def stream_chat(self, messages: list[ChatMessage]):
yield "angefangene Antwort"
raise KeyboardInterrupt
class _FakeHybridProvider(_FakeProvider):
name = "hybrid"
@@ -59,10 +75,53 @@ class CliTests(unittest.TestCase):
)
self.assertEqual(result, 0)
self.assertTrue(any(line == "Javis: Echo: Hallo" for line in output))
self.assertIn("Javis: Echo: Hallo\n", "".join(output))
self.assertTrue(any("Nachrichten" in line for line in output))
self.assertEqual(output[-1], "Chat beendet.")
def test_streaming_cleans_terminal_markdown_but_persists_raw_answer(self) -> None:
with tempfile.TemporaryDirectory() as directory:
store = SQLiteSessionStore(Path(directory) / "sessions.sqlite3")
service = ChatService(store, _MarkdownProvider())
inputs = iter(["Frage", "/exit"])
output: list[str] = []
streamed: list[str] = []
result = run_chat(
service,
input_fn=lambda _prompt: next(inputs),
output=output.append,
stream_output=streamed.append,
)
session = service.list_sessions()[0]
messages = service.load_session(session.id).messages
self.assertEqual(result, 0)
self.assertEqual("".join(streamed), "Javis: Antwort\nSQLite\n")
self.assertEqual(messages[-1].content, "## **Antwort**\n`SQLite`")
def test_keyboard_interrupt_drops_partial_answer_and_returns_to_prompt(self) -> None:
with tempfile.TemporaryDirectory() as directory:
store = SQLiteSessionStore(Path(directory) / "sessions.sqlite3")
service = ChatService(store, _InterruptingProvider())
inputs = iter(["Frage", "/exit"])
output: list[str] = []
streamed: list[str] = []
result = run_chat(
service,
input_fn=lambda _prompt: next(inputs),
output=output.append,
stream_output=streamed.append,
)
session = service.list_sessions()[0]
messages = service.load_session(session.id).messages
self.assertEqual(result, 0)
self.assertEqual(messages, [])
self.assertTrue(any("nicht gespeichert" in line for line in output))
self.assertEqual(output[-1], "Chat beendet.")
def test_provider_privacy_and_status_commands_are_sanitized(self) -> None:
with tempfile.TemporaryDirectory() as directory:
store = SQLiteSessionStore(Path(directory) / "sessions.sqlite3")
+59
View File
@@ -23,6 +23,7 @@ class _FakeModels:
def __init__(self, result: object) -> None:
self.result = result
self.call: dict[str, Any] | None = None
self.stream_call: dict[str, Any] | None = None
def generate_content(self, **kwargs: object) -> object:
self.call = kwargs
@@ -30,6 +31,14 @@ class _FakeModels:
raise self.result
return self.result
def generate_content_stream(self, **kwargs: object):
self.stream_call = kwargs
if isinstance(self.result, BaseException):
raise self.result
if isinstance(self.result, list):
return iter(self.result)
return iter((self.result,))
class _FakeClient:
def __init__(self, result: object) -> None:
@@ -56,6 +65,23 @@ def _response(text: str = "Cloud-Antwort") -> SimpleNamespace:
)
def _stream_response(
visible_text: str,
*,
thought_text: str | None = None,
with_usage: bool = False,
) -> SimpleNamespace:
parts = []
if thought_text:
parts.append(SimpleNamespace(text=thought_text, thought=True))
parts.append(SimpleNamespace(text=visible_text, thought=False))
usage = SimpleNamespace(prompt_token_count=12, candidates_token_count=7) if with_usage else None
return SimpleNamespace(
candidates=[SimpleNamespace(content=SimpleNamespace(parts=parts))],
usage_metadata=usage,
)
class GeminiProviderTests(unittest.TestCase):
def _provider(
self, result: object, **overrides: object
@@ -99,6 +125,39 @@ class GeminiProviderTests(unittest.TestCase):
self.assertEqual([content.role for content in call["contents"]], ["user", "model", "user"])
self.assertEqual(call["config"].max_output_tokens, 1024)
self.assertIsNone(call["config"].tools)
self.assertFalse(call["config"].thinking_config.include_thoughts)
def test_streaming_yields_only_visible_text_and_tracks_usage(self) -> None:
provider, factory = self._provider(
[
_stream_response("Cloud ", thought_text="verstecktes Denken"),
_stream_response("Antwort", with_usage=True),
]
)
chunks = list(provider.stream_chat([ChatMessage("user", "Was ist SQLite?")]))
self.assertEqual(chunks, ["Cloud ", "Antwort"])
self.assertNotIn("verstecktes Denken", "".join(chunks))
self.assertEqual(provider.last_usage.input_tokens, 12)
self.assertEqual(provider.last_usage.output_tokens, 7)
self.assertEqual(
factory.client.models.stream_call["model"],
"gemini-3.6-flash",
)
def test_streaming_error_is_sanitized(self) -> None:
provider, _ = self._provider(
errors.ClientError(
429,
{"message": "quota test-key-not-a-real-secret"},
)
)
with self.assertRaises(ProviderRateLimitError) as raised:
list(provider.stream_chat([ChatMessage("user", "Was ist Python?")]))
self.assertNotIn("test-key-not-a-real-secret", str(raised.exception))
def test_missing_key_stops_before_client_creation(self) -> None:
factory = _ClientFactory(_response())
+65 -7
View File
@@ -4,22 +4,39 @@ import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import ClassVar
from javis.providers.base import ChatMessage
from javis.providers.base import ChatMessage, ResponseAbortedError
from javis.providers.ollama import OllamaProvider
class _OllamaHandler(BaseHTTPRequestHandler):
request_payload: ClassVar[dict[str, object]] = {}
abort_stream: ClassVar[bool] = False
def do_POST(self) -> None:
length = int(self.headers["Content-Length"])
type(self).request_payload = json.loads(self.rfile.read(length))
body = json.dumps(
{
"message": {"role": "assistant", "content": "Lokale Antwort"},
"done": True,
}
).encode()
if type(self).request_payload["stream"]:
chunks = [
{
"message": {"role": "assistant", "content": "Lokale "},
"done": False,
}
]
if not type(self).abort_stream:
chunks.append(
{
"message": {"role": "assistant", "content": "Antwort"},
"done": True,
}
)
body = b"".join(json.dumps(chunk).encode() + b"\n" for chunk in chunks)
else:
body = json.dumps(
{
"message": {"role": "assistant", "content": "Lokale Antwort"},
"done": True,
}
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
@@ -39,6 +56,9 @@ class _OllamaHandler(BaseHTTPRequestHandler):
class OllamaProviderTests(unittest.TestCase):
def setUp(self) -> None:
_OllamaHandler.abort_stream = False
def test_provider_uses_local_chat_endpoint(self) -> None:
server = ThreadingHTTPServer(("127.0.0.1", 0), _OllamaHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
@@ -62,6 +82,44 @@ class OllamaProviderTests(unittest.TestCase):
self.assertTrue(status.reachable)
self.assertTrue(status.model_available)
def test_provider_streams_visible_text_without_thinking(self) -> None:
server = ThreadingHTTPServer(("127.0.0.1", 0), _OllamaHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
provider = OllamaProvider(
"test-model",
f"http://127.0.0.1:{server.server_port}",
2,
)
chunks = list(provider.stream_chat([ChatMessage("user", "Hallo")]))
finally:
server.shutdown()
server.server_close()
thread.join()
self.assertEqual(chunks, ["Lokale ", "Antwort"])
self.assertTrue(_OllamaHandler.request_payload["stream"])
self.assertFalse(_OllamaHandler.request_payload["think"])
def test_provider_rejects_incomplete_stream(self) -> None:
_OllamaHandler.abort_stream = True
server = ThreadingHTTPServer(("127.0.0.1", 0), _OllamaHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
provider = OllamaProvider(
"test-model",
f"http://127.0.0.1:{server.server_port}",
2,
)
with self.assertRaises(ResponseAbortedError):
list(provider.stream_chat([ChatMessage("user", "Hallo")]))
finally:
server.shutdown()
server.server_close()
thread.join()
if __name__ == "__main__":
unittest.main()
+74
View File
@@ -16,6 +16,7 @@ from javis.providers.base import (
ProviderRateLimitError,
ProviderTimeoutError,
ProviderUsage,
ResponseAbortedError,
)
from javis.security.privacy import CloudPolicy, PrivacyRouter
@@ -28,11 +29,15 @@ class _RecordingProvider:
*,
answer: str = "Antwort",
error: Exception | None = None,
stream_chunks: tuple[str, ...] | None = None,
stream_error: Exception | None = None,
) -> None:
self.name = name
self.model = model
self.answer = answer
self.error = error
self.stream_chunks = stream_chunks
self.stream_error = stream_error
self.calls: list[list[ChatMessage]] = []
self.last_usage = ProviderUsage(11, 5)
@@ -42,6 +47,17 @@ class _RecordingProvider:
raise self.error
return self.answer
def stream_chat(self, messages: list[ChatMessage]):
self.calls.append(messages)
if self.stream_chunks is None:
if self.error:
raise self.error
yield self.answer
else:
yield from self.stream_chunks
if self.stream_error:
raise self.stream_error
class HybridProviderTests(unittest.TestCase):
def setUp(self) -> None:
@@ -283,6 +299,64 @@ class HybridProviderTests(unittest.TestCase):
self.assertFalse(router.last_route.cloud_suppressed_by_local_mode)
self.assertTrue(router.last_route.technical_fallback)
def test_allowed_cloud_response_streams_visible_chunks(self) -> None:
self.cloud.stream_chunks = ("Cloud ", "Stream")
router = self._router()
chunks = list(router.stream_chat([ChatMessage("user", "Wie funktioniert SQLite?")]))
self.assertEqual(chunks, ["Cloud ", "Stream"])
self.assertEqual(router.last_route.provider, "gemini")
self.assertFalse(router.last_route.fallback)
def test_cloud_stream_error_before_output_falls_back_locally(self) -> None:
self.cloud.stream_chunks = ()
self.cloud.stream_error = CloudNetworkError("offline")
self.local.stream_chunks = ("Lokaler ", "Ersatz")
router = self._router()
chunks = list(router.stream_chat([ChatMessage("user", "Wie funktioniert Python?")]))
self.assertEqual(chunks, ["Lokaler ", "Ersatz"])
self.assertTrue(router.last_route.fallback)
self.assertTrue(router.last_route.technical_fallback)
self.assertIn("lokale Antwort", self.notices[-1])
def test_cloud_stream_error_after_output_never_appends_local_answer(self) -> None:
self.cloud.stream_chunks = ("Teilantwort",)
self.cloud.stream_error = CloudNetworkError("offline")
self.local.stream_chunks = ("Lokaler Ersatz",)
router = self._router()
stream = router.stream_chat([ChatMessage("user", "Wie funktioniert Python?")])
self.assertEqual(next(stream), "Teilantwort")
with self.assertRaises(ResponseAbortedError):
next(stream)
self.assertFalse(self.local.calls)
self.assertFalse(router.last_route.fallback)
self.assertTrue(router.last_route.technical_fallback)
def test_medical_never_streams_locally_before_cloud_construction(self) -> None:
constructed = False
self.local.stream_chunks = ("Lokal",)
def cloud_factory() -> _RecordingProvider:
nonlocal constructed
constructed = True
return self.cloud
router = self._router(cloud_provider_factory=cloud_factory, mode="gemini")
chunks = list(
router.stream_chat([ChatMessage("user", "Welche Diagnose passt zu meinen Schmerzen?")])
)
self.assertEqual(chunks, ["Lokal"])
self.assertFalse(constructed)
self.assertFalse(self.approvals)
self.assertEqual(router.last_route.privacy_policy, CloudPolicy.NEVER)
if __name__ == "__main__":
unittest.main()