spb/company-atlas

1"""LLM gateway: OpenAI-compatible JSON mode (respx-mocked), invalid-JSON repair path, 503/429 backoff, health, prompt loading."""2from __future__ import annotations34import json56import httpx7import pytest8import respx9from pydantic import BaseModel1011from companyatlas.services.llm.gateway import (12    LLMError,13    LLMValidationError,14    OpenAICompatibleProvider,15    extract_json_text,16    parse_json_object,17)18from companyatlas.services.llm.prompts import available_prompts, load_prompt19from companyatlas.services.llm.schemas import ChangeClassification, EventSummary, LegalDiffSummary2021BASE = "https://llm.test/v1"222324class Probe(BaseModel):25    ok: bool26    n: int272829def _completion(content: str, *, prompt_tokens: int = 50, completion_tokens: int = 20) -> dict:30    return {"id": "x", "object": "chat.completion", "model": "qwen3-4b-instruct-2507-4bit",31            "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],32            "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens}}333435def _provider(**kw) -> OpenAICompatibleProvider:  # type: ignore[no-untyped-def]36    sleeps: list[float] = []3738    async def fake_sleep(s: float) -> None:39        sleeps.append(s)4041    p = OpenAICompatibleProvider(base_url=BASE, api_key="k", timeout_s=5, sleep=fake_sleep, **kw)42    p.sleeps = sleeps  # type: ignore[attr-defined]43    return p444546@respx.mock47async def test_complete_json_happy_path():48    route = respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion('{"ok": true, "n": 3}')))49    p = _provider()50    res = await p.complete_json("small", "sys", "user", Probe, max_tokens=50)51    assert res.data == Probe(ok=True, n=3)52    assert res.model == "qwen3-4b-instruct-2507-4bit" and res.request_tokens == 50 and res.response_tokens == 20 and res.attempts == 153    sent = json.loads(route.calls[0].request.content)54    assert sent["response_format"] == {"type": "json_object"} and sent["model"] == "qwen3-4b-instruct-2507-4bit"55    assert "JSON schema" in sent["messages"][0]["content"] and sent["messages"][1]["content"] == "user"56    assert route.calls[0].request.headers["Authorization"] == "Bearer k"57    await p.close()585960@respx.mock61async def test_invalid_json_triggers_one_repair_round_trip():62    route = respx.post(f"{BASE}/chat/completions")63    route.side_effect = [httpx.Response(200, json=_completion("Sure! Here you go: {\"ok\": true, \"n\": \"three\"}")),64                         httpx.Response(200, json=_completion("```json\n{\"ok\": true, \"n\": 3}\n```"))]65    p = _provider()66    res = await p.complete_json("small", "sys", "user", Probe)67    assert res.data.n == 3 and res.attempts == 2 and res.repaired is True68    assert res.request_tokens == 100                                      # both calls accounted69    repair_msgs = json.loads(route.calls[1].request.content)["messages"]70    assert repair_msgs[-1]["role"] == "user" and "not valid" in repair_msgs[-1]["content"]71    await p.close()727374@respx.mock75async def test_still_invalid_after_repair_raises_validation_error():76    respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion("no json at all")))77    p = _provider()78    with pytest.raises(LLMValidationError):79        await p.complete_json("small", "sys", "user", Probe)80    await p.close()818283@respx.mock84async def test_backoff_on_503_then_success():85    route = respx.post(f"{BASE}/chat/completions")86    route.side_effect = [httpx.Response(503, text="loading model"), httpx.Response(429, text="busy", headers={"Retry-After": "7"}),87                         httpx.Response(200, json=_completion('{"ok": true, "n": 1}'))]88    p = _provider(max_tries=6, backoff_initial_s=15, backoff_max_s=120)89    res = await p.complete_json("small", "sys", "user", Probe)90    assert res.data.n == 1 and route.call_count == 391    assert p.sleeps == [15.0, 7.0]                                       # exponential backoff, then Retry-After honoured92    await p.close()939495@respx.mock96async def test_backoff_exhausted_raises_retryable_error():97    respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(503, text="still loading"))98    p = _provider(max_tries=3, backoff_initial_s=15, backoff_max_s=120)99    with pytest.raises(LLMError) as exc:100        await p.complete_json("small", "sys", "user", Probe)101    assert exc.value.retryable and exc.value.status == 503102    assert p.sleeps == [15.0, 30.0]103    await p.close()104105106@respx.mock107async def test_non_retryable_error_and_health():108    respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(401, text="bad key"))109    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"}]}))110    p = _provider()111    with pytest.raises(LLMError) as exc:112        await p.complete_json("small", "sys", "user", Probe)113    assert not exc.value.retryable and exc.value.status == 401 and p.sleeps == []114    h = await p.health()115    assert h.ok and "qwen3-4b-instruct-2507-4bit" in h.models116    await p.close()117118119async def test_not_configured():120    p = OpenAICompatibleProvider(base_url="", api_key="")121    from companyatlas.services.llm.gateway import LLMNotConfigured122123    with pytest.raises(LLMNotConfigured):124        await p.complete_json("small", "s", "u", Probe)125    assert (await p.health()).ok is False126127128def test_json_extraction_tolerates_reasoning_and_fences():129    assert extract_json_text("<think>hmm</think>\n```json\n{\"a\": 1}\n```") == '{"a": 1}'130    assert extract_json_text('prefix {"a": {"b": 2}} suffix') == '{"a": {"b": 2}}'131    with pytest.raises(ValueError):132        extract_json_text("nothing here")133    with pytest.raises(ValueError):134        parse_json_object("[1, 2]", Probe)135136137def test_schemas_enforce_taxonomy_and_wording():138    c = ChangeClassification(event_subtype="price increase", importance=0.7, confidence=0.8, title="Pro plan now $59", tags=["Pricing", "pricing", "x"])139    assert c.event_subtype == "PRICE_INCREASE" and c.tags == ["pricing", "x"]140    assert ChangeClassification(event_subtype="MADE_UP", importance=0, confidence=0, title="abc").event_subtype == "OTHER"141    with pytest.raises(ValueError):142        EventSummary(summary="The company laid off 72 employees")143    legal = LegalDiffSummary(summary="Terms: 2 sections changed.", materiality="huge", sections_changed=[{"section": "7", "change": "notice 30 days"}])144    assert legal.materiality == "unclear" and legal.sections_changed[0].section == "7"145    with pytest.raises(ValueError):146        ChangeClassification(event_subtype="OTHER", importance=1.5, confidence=0.5, title="x" * 10)147148149def test_prompt_files_load_by_name_and_version():150    avail = available_prompts()151    for task in ("change-classifier", "event-summarizer", "legal-diff", "industry-tagger", "ask-router"):152        assert "v1" in avail[task]153        p = load_prompt(task, "v1")154        assert p.ref == f"{task}/v1" and len(p.system) > 200 and not p.system.startswith("<!--")155        if task in ("change-classifier", "event-summarizer", "legal-diff"):156            assert "laid off" in p.system.lower() and "fired" in p.system.lower()      # content prompts state the wording ban157    assert load_prompt("legal-diff").version == "v1"158    with pytest.raises(FileNotFoundError):159        load_prompt("does-not-exist")160