"""LLM gateway: OpenAI-compatible JSON mode (respx-mocked), invalid-JSON repair path, 503/429 backoff, health, prompt loading.""" from __future__ import annotations import json import httpx import pytest import respx from pydantic import BaseModel from companyatlas.services.llm.gateway import ( LLMError, LLMValidationError, OpenAICompatibleProvider, extract_json_text, parse_json_object, ) from companyatlas.services.llm.prompts import available_prompts, load_prompt from companyatlas.services.llm.schemas import ChangeClassification, EventSummary, LegalDiffSummary BASE = "https://llm.test/v1" class Probe(BaseModel): ok: bool n: int def _completion(content: str, *, prompt_tokens: int = 50, completion_tokens: int = 20) -> dict: return {"id": "x", "object": "chat.completion", "model": "qwen3-4b-instruct-2507-4bit", "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens}} def _provider(**kw) -> OpenAICompatibleProvider: # type: ignore[no-untyped-def] sleeps: list[float] = [] async def fake_sleep(s: float) -> None: sleeps.append(s) p = OpenAICompatibleProvider(base_url=BASE, api_key="k", timeout_s=5, sleep=fake_sleep, **kw) p.sleeps = sleeps # type: ignore[attr-defined] return p @respx.mock async def test_complete_json_happy_path(): route = respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion('{"ok": true, "n": 3}'))) p = _provider() res = await p.complete_json("small", "sys", "user", Probe, max_tokens=50) assert res.data == Probe(ok=True, n=3) assert res.model == "qwen3-4b-instruct-2507-4bit" and res.request_tokens == 50 and res.response_tokens == 20 and res.attempts == 1 sent = json.loads(route.calls[0].request.content) assert sent["response_format"] == {"type": "json_object"} and sent["model"] == "qwen3-4b-instruct-2507-4bit" assert "JSON schema" in sent["messages"][0]["content"] and sent["messages"][1]["content"] == "user" assert route.calls[0].request.headers["Authorization"] == "Bearer k" await p.close() @respx.mock async def test_invalid_json_triggers_one_repair_round_trip(): route = respx.post(f"{BASE}/chat/completions") route.side_effect = [httpx.Response(200, json=_completion("Sure! Here you go: {\"ok\": true, \"n\": \"three\"}")), httpx.Response(200, json=_completion("```json\n{\"ok\": true, \"n\": 3}\n```"))] p = _provider() res = await p.complete_json("small", "sys", "user", Probe) assert res.data.n == 3 and res.attempts == 2 and res.repaired is True assert res.request_tokens == 100 # both calls accounted repair_msgs = json.loads(route.calls[1].request.content)["messages"] assert repair_msgs[-1]["role"] == "user" and "not valid" in repair_msgs[-1]["content"] await p.close() @respx.mock async def test_still_invalid_after_repair_raises_validation_error(): respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion("no json at all"))) p = _provider() with pytest.raises(LLMValidationError): await p.complete_json("small", "sys", "user", Probe) await p.close() @respx.mock async def test_backoff_on_503_then_success(): route = respx.post(f"{BASE}/chat/completions") route.side_effect = [httpx.Response(503, text="loading model"), httpx.Response(429, text="busy", headers={"Retry-After": "7"}), httpx.Response(200, json=_completion('{"ok": true, "n": 1}'))] p = _provider(max_tries=6, backoff_initial_s=15, backoff_max_s=120) res = await p.complete_json("small", "sys", "user", Probe) assert res.data.n == 1 and route.call_count == 3 assert p.sleeps == [15.0, 7.0] # exponential backoff, then Retry-After honoured await p.close() @respx.mock async def test_backoff_exhausted_raises_retryable_error(): respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(503, text="still loading")) p = _provider(max_tries=3, backoff_initial_s=15, backoff_max_s=120) with pytest.raises(LLMError) as exc: await p.complete_json("small", "sys", "user", Probe) assert exc.value.retryable and exc.value.status == 503 assert p.sleeps == [15.0, 30.0] await p.close() @respx.mock async def test_non_retryable_error_and_health(): respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(401, text="bad key")) respx.get(f"{BASE}/models").mock(return_value=httpx.Response(200, json={"data": [{"id": "qwen3-4b-instruct-2507-4bit"}, {"id": "qwen3.6-35b-a3b-4bit"}]})) p = _provider() with pytest.raises(LLMError) as exc: await p.complete_json("small", "sys", "user", Probe) assert not exc.value.retryable and exc.value.status == 401 and p.sleeps == [] h = await p.health() assert h.ok and "qwen3-4b-instruct-2507-4bit" in h.models await p.close() async def test_not_configured(): p = OpenAICompatibleProvider(base_url="", api_key="") from companyatlas.services.llm.gateway import LLMNotConfigured with pytest.raises(LLMNotConfigured): await p.complete_json("small", "s", "u", Probe) assert (await p.health()).ok is False def test_json_extraction_tolerates_reasoning_and_fences(): assert extract_json_text("hmm\n```json\n{\"a\": 1}\n```") == '{"a": 1}' assert extract_json_text('prefix {"a": {"b": 2}} suffix') == '{"a": {"b": 2}}' with pytest.raises(ValueError): extract_json_text("nothing here") with pytest.raises(ValueError): parse_json_object("[1, 2]", Probe) def test_schemas_enforce_taxonomy_and_wording(): c = ChangeClassification(event_subtype="price increase", importance=0.7, confidence=0.8, title="Pro plan now $59", tags=["Pricing", "pricing", "x"]) assert c.event_subtype == "PRICE_INCREASE" and c.tags == ["pricing", "x"] assert ChangeClassification(event_subtype="MADE_UP", importance=0, confidence=0, title="abc").event_subtype == "OTHER" with pytest.raises(ValueError): EventSummary(summary="The company laid off 72 employees") legal = LegalDiffSummary(summary="Terms: 2 sections changed.", materiality="huge", sections_changed=[{"section": "7", "change": "notice 30 days"}]) assert legal.materiality == "unclear" and legal.sections_changed[0].section == "7" with pytest.raises(ValueError): ChangeClassification(event_subtype="OTHER", importance=1.5, confidence=0.5, title="x" * 10) def test_prompt_files_load_by_name_and_version(): avail = available_prompts() for task in ("change-classifier", "event-summarizer", "legal-diff", "industry-tagger", "ask-router"): assert "v1" in avail[task] p = load_prompt(task, "v1") assert p.ref == f"{task}/v1" and len(p.system) > 200 and not p.system.startswith("