feat: stream provider responses
This commit is contained in:
@@ -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
@@ -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")
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user