diff --git a/digital_mate/llm/client.py b/digital_mate/llm/client.py index 2394f0f..e2c6d15 100644 --- a/digital_mate/llm/client.py +++ b/digital_mate/llm/client.py @@ -127,6 +127,43 @@ def _parse_content(response_json: dict[str, Any]) -> str: raise LLMError("LLM returned empty response.") return content.strip() + @staticmethod + def _parse_response(response: httpx.Response) -> dict[str, Any]: + """Normalize JSON or SSE responses into a chat completion object.""" + content_type = response.headers.get("content-type", "").lower() + if "text/event-stream" not in content_type: + return response.json() + + content_parts: list[str] = [] + usage: dict[str, Any] | None = None + for line in response.text.splitlines(): + if not line.startswith("data: "): + continue + data = line[6:].strip() + if not data or data == "[DONE]": + continue + try: + chunk = json.loads(data) + except json.JSONDecodeError: + continue + if isinstance(chunk.get("usage"), dict): + usage = chunk["usage"] + choices = chunk.get("choices", []) + if not choices: + continue + choice = choices[0] + message = choice.get("delta") or choice.get("message") or {} + content = message.get("content") + if content: + content_parts.append(content) + + result: dict[str, Any] = { + "choices": [{"message": {"content": "".join(content_parts)}}] + } + if usage is not None: + result["usage"] = usage + return result + @staticmethod def _strip_code_fences(text: str) -> str: """Strip markdown code fences from JSON text if present.""" @@ -174,7 +211,7 @@ async def chat( try: response = await self._client.post("/chat/completions", json=body) response.raise_for_status() - response_json = response.json() + response_json = self._parse_response(response) metrics.record_usage(response_json.get("usage")) return self._parse_content(response_json) @@ -477,7 +514,7 @@ async def chat_with_image( try: response = await self._client.post("/chat/completions", json=body) response.raise_for_status() - response_json = response.json() + response_json = self._parse_response(response) metrics.record_usage(response_json.get("usage")) return self._parse_content(response_json) @@ -679,7 +716,7 @@ async def chat_json( try: response = await self._client.post("/chat/completions", json=body) response.raise_for_status() - response_json = response.json() + response_json = self._parse_response(response) metrics.record_usage(response_json.get("usage")) content = self._parse_content(response_json) diff --git a/tests/test_image_input.py b/tests/test_image_input.py index db08283..e6e78f3 100644 --- a/tests/test_image_input.py +++ b/tests/test_image_input.py @@ -41,6 +41,7 @@ def _make_httpx_response(content: str) -> MagicMock: """Build a mock httpx.Response returning the given content.""" resp = MagicMock(spec=httpx.Response) resp.status_code = 200 + resp.headers = {} resp.json.return_value = { "choices": [{"message": {"content": content}}] } diff --git a/tests/test_llm_client.py b/tests/test_llm_client.py index d5060af..1cd1513 100644 --- a/tests/test_llm_client.py +++ b/tests/test_llm_client.py @@ -42,6 +42,7 @@ def _make_httpx_response( """Build a mock httpx.Response returning the given content.""" resp = MagicMock(spec=httpx.Response) resp.status_code = 200 + resp.headers = {} resp.json.return_value = { "choices": [{"message": {"content": content}}], **({"usage": usage} if usage else {}), @@ -54,6 +55,7 @@ def _make_httpx_response_none_content() -> MagicMock: """Build a mock httpx.Response with None content.""" resp = MagicMock(spec=httpx.Response) resp.status_code = 200 + resp.headers = {} resp.json.return_value = { "choices": [{"message": {"content": None}}] } @@ -61,6 +63,25 @@ def _make_httpx_response_none_content() -> MagicMock: return resp +def _make_buffered_sse_response( + contents: list[str], usage: dict[str, int] | None = None +) -> httpx.Response: + """Build a buffered SSE response returned by stream-forcing gateways.""" + lines = [ + "data: " + json.dumps({"choices": [{"delta": {"content": content}}]}) + for content in contents + ] + if usage is not None: + lines.append("data: " + json.dumps({"choices": [], "usage": usage})) + lines.append("data: [DONE]") + return httpx.Response( + 200, + headers={"content-type": "text/event-stream; charset=utf-8"}, + text="\n\n".join(lines), + request=httpx.Request("POST", "http://test.api/v1/chat/completions"), + ) + + def _make_httpx_status_error(status_code: int = 400) -> httpx.HTTPStatusError: """Build an httpx.HTTPStatusError for the given status code.""" request = httpx.Request("POST", "http://test.api/v1/chat/completions") @@ -107,6 +128,21 @@ async def test_chat_records_provider_token_usage(llm_client: LLMClient) -> None: assert snapshot["tokens"]["total"] == 10 +@pytest.mark.asyncio +async def test_chat_accepts_buffered_sse_response(llm_client: LLMClient) -> None: + """chat() joins deltas when a gateway forces SSE for non-stream calls.""" + llm_client._client.post = AsyncMock( + return_value=_make_buffered_sse_response( + ["Hello", " world"], + {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4}, + ) + ) + + result = await llm_client.chat([{"role": "user", "content": "hi"}]) + + assert result == "Hello world" + + @pytest.mark.asyncio async def test_chat_retries_on_timeout(llm_client: LLMClient) -> None: """chat() retries on httpx.TimeoutException with exponential backoff.""" @@ -244,6 +280,24 @@ async def test_chat_json_success(llm_client: LLMClient) -> None: assert result == {"pillar": "content", "action": "caption"} +@pytest.mark.asyncio +async def test_chat_json_accepts_buffered_sse_response( + llm_client: LLMClient, +) -> None: + """chat_json() parses fenced JSON assembled from SSE deltas.""" + llm_client._client.post = AsyncMock( + return_value=_make_buffered_sse_response( + ["```json\n", '{"pillar":"content",', '"action":"caption"}', "\n```"] + ) + ) + + result = await llm_client.chat_json( + [{"role": "user", "content": "classify"}] + ) + + assert result == {"pillar": "content", "action": "caption"} + + @pytest.mark.asyncio async def test_chat_json_retries_on_invalid_json(llm_client: LLMClient) -> None: """chat_json() retries on JSONDecodeError, then succeeds."""