8d8a863300
packages/artifacts: - render_cv_pdf(profile, sections) -> bytes: data-driven CV PDF generation using fpdf2 with bundled DejaVuSans TTF for unicode (Swedish chars tested). Jinja2 template for layout data prep, adapted from build_cv.py approach. - render_cover_letter(text, profile) -> bytes: simple cover letter PDF. - hash_bytes(b) -> str: sha256 hex digest. - next_version(existing) -> int: version numbering helper. - 16 tests, all passing: PDF validity, Swedish characters, hash stability, cover letter rendering, hash correctness, version logic. packages/llm-gateway: - Async-first Gateway class with provider config from env. - Mock mode default when no API key env present (deterministic canned outputs per task name, defined in mock.py). - Telemetry sink injectable (async or sync callable, receives TelemetryRow). - Budget guard raises BudgetExceeded BEFORE any provider call is made. - Retry policy: max 2 retries on 429/5xx, then fallback provider for STRONG tasks only. CHEAP tasks never use fallback (paid provider protection). - Schema validation via jsonschema; SchemaValidationError on mismatch. - Task class routing: CHEAP (score, extract, cv_assist) vs STRONG (critique, cl_critique, research). Model routing per task class. - Paid provider detection heuristic; warns on paid fallback config. - 31 tests, all passing: mock determinism, schema pass/fail, budget guard (mock + real mode), telemetry sink (async/sync/none), config from env, provider calls with mocked HTTP (retry, fallback, no-fallback-for-cheap). docker-compose.yml: - postgres:16 service, user/pass/db = jobhunt, host port 5433->5432, named volume jobhunt_pgdata, healthcheck. .env.example: - DATABASE_URL, LLM provider config (primary + fallback), task budgets, API and web settings.
491 lines
No EOL
17 KiB
Python
491 lines
No EOL
17 KiB
Python
"""Tests for the LLM gateway package."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from llm_gateway.config import GatewayConfig, ProviderConfig, TaskClass
|
|
from llm_gateway.exceptions import BudgetExceeded, SchemaValidationError
|
|
from llm_gateway.gateway import Gateway, TelemetryRow, run_task
|
|
from llm_gateway.mock import get_mock_output, MOCK_OUTPUTS
|
|
|
|
|
|
# ---- Fixtures ----
|
|
|
|
|
|
def mock_config(**overrides: Any) -> GatewayConfig:
|
|
"""Build a config in mock mode (no API key)."""
|
|
primary = ProviderConfig(
|
|
name="primary",
|
|
base_url="https://mock.example.com/v1",
|
|
api_key="",
|
|
model="glm-5.2",
|
|
)
|
|
defaults: dict[str, Any] = {
|
|
"primary": primary,
|
|
"fallback": None,
|
|
"cheap_model": "glm-5.2",
|
|
"strong_model": "glm-5.2",
|
|
"budgets": {"score": 2000, "extract": 4000, "default": 4000},
|
|
"max_retries": 2,
|
|
}
|
|
defaults.update(overrides)
|
|
return GatewayConfig(**defaults)
|
|
|
|
|
|
def real_config(**overrides: Any) -> GatewayConfig:
|
|
"""Build a config with a fake API key (non-mock mode)."""
|
|
primary = ProviderConfig(
|
|
name="primary",
|
|
base_url="https://mock.example.com/v1",
|
|
api_key="fake-key-1234",
|
|
model="glm-5.2",
|
|
)
|
|
defaults: dict[str, Any] = {
|
|
"primary": primary,
|
|
"fallback": None,
|
|
"cheap_model": "glm-5.2",
|
|
"strong_model": "glm-5.2",
|
|
"budgets": {"score": 2000, "extract": 4000, "default": 4000},
|
|
"max_retries": 2,
|
|
}
|
|
defaults.update(overrides)
|
|
return GatewayConfig(**defaults)
|
|
|
|
|
|
# ---- Mock mode tests ----
|
|
|
|
|
|
class TestMockMode:
|
|
async def test_mock_returns_deterministic_output(self) -> None:
|
|
"""Mock mode returns deterministic canned outputs per task."""
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
result_a = await gw.run_task("score", "Score this job")
|
|
result_b = await gw.run_task("score", "Score this job")
|
|
assert result_a == result_b
|
|
assert result_a["score"] == 75
|
|
await gw.aclose()
|
|
|
|
async def test_mock_different_tasks_different_output(self) -> None:
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
score_result = await gw.run_task("score", "prompt")
|
|
extract_result = await gw.run_task("extract", "prompt")
|
|
assert score_result != extract_result
|
|
assert "score" in score_result
|
|
assert "company" in extract_result
|
|
await gw.aclose()
|
|
|
|
async def test_mock_unknown_task(self) -> None:
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
result = await gw.run_task("unknown_task", "prompt")
|
|
assert result["result"] == "mock output"
|
|
assert result["task"] == "unknown_task"
|
|
await gw.aclose()
|
|
|
|
async def test_mock_mode_property(self) -> None:
|
|
config = mock_config()
|
|
assert config.mock_mode is True
|
|
|
|
async def test_real_mode_not_mock(self) -> None:
|
|
config = real_config()
|
|
assert config.mock_mode is False
|
|
|
|
|
|
# ---- Schema validation tests ----
|
|
|
|
|
|
class TestSchemaValidation:
|
|
async def test_schema_passes(self) -> None:
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"score": {"type": "number"},
|
|
"rationale": {"type": "object"},
|
|
},
|
|
"required": ["score"],
|
|
}
|
|
result = await gw.run_task("score", "prompt", schema=schema)
|
|
assert "score" in result
|
|
await gw.aclose()
|
|
|
|
async def test_schema_fails(self) -> None:
|
|
"""Schema validation failure should raise SchemaValidationError."""
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
# The mock output for 'score' has score=75 (number). We require a string,
|
|
# which should fail validation.
|
|
bad_schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"score": {"type": "string"},
|
|
},
|
|
"required": ["score"],
|
|
}
|
|
with pytest.raises(SchemaValidationError):
|
|
await gw.run_task("score", "prompt", schema=bad_schema)
|
|
await gw.aclose()
|
|
|
|
async def test_schema_missing_required_field(self) -> None:
|
|
config = mock_config()
|
|
gw = Gateway(config)
|
|
schema = {
|
|
"type": "object",
|
|
"required": ["nonexistent_field"],
|
|
}
|
|
with pytest.raises(SchemaValidationError):
|
|
await gw.run_task("score", "prompt", schema=schema)
|
|
await gw.aclose()
|
|
|
|
|
|
# ---- Budget guard tests ----
|
|
|
|
|
|
class TestBudgetGuard:
|
|
async def test_budget_exceeded_raises_before_call(self) -> None:
|
|
"""Over-budget prompt should raise BudgetExceeded before any call."""
|
|
config = mock_config(budgets={"score": 10})
|
|
gw = Gateway(config)
|
|
# 10 token budget, ~4 chars/token, so >40 chars should exceed.
|
|
long_prompt = "x" * 100
|
|
with pytest.raises(BudgetExceeded):
|
|
await gw.run_task("score", long_prompt)
|
|
await gw.aclose()
|
|
|
|
async def test_budget_within_limit_does_not_raise(self) -> None:
|
|
config = mock_config(budgets={"score": 10000})
|
|
gw = Gateway(config)
|
|
result = await gw.run_task("score", "short prompt")
|
|
assert result["score"] == 75
|
|
await gw.aclose()
|
|
|
|
async def test_budget_guard_in_real_mode(self) -> None:
|
|
"""Budget guard must raise before call even in real (non-mock) mode."""
|
|
config = real_config(budgets={"score": 10})
|
|
gw = Gateway(config)
|
|
with pytest.raises(BudgetExceeded):
|
|
await gw.run_task("score", "x" * 100)
|
|
await gw.aclose()
|
|
|
|
async def test_default_budget_fallback(self) -> None:
|
|
"""Unknown task should use default budget."""
|
|
config = mock_config(budgets={"score": 2000, "default": 100})
|
|
gw = Gateway(config)
|
|
# Unknown task uses default=100, so >400 chars exceeds.
|
|
with pytest.raises(BudgetExceeded):
|
|
await gw.run_task("unknown_task", "x" * 500)
|
|
await gw.aclose()
|
|
|
|
|
|
# ---- Telemetry tests ----
|
|
|
|
|
|
class TestTelemetry:
|
|
async def test_telemetry_sink_called_in_mock_mode(self) -> None:
|
|
sink_calls: list[TelemetryRow] = []
|
|
|
|
async def sink(row: TelemetryRow) -> None:
|
|
sink_calls.append(row)
|
|
|
|
config = mock_config()
|
|
gw = Gateway(config, telemetry_sink=sink)
|
|
await gw.run_task("score", "test prompt")
|
|
assert len(sink_calls) == 1
|
|
assert sink_calls[0].task == "score"
|
|
assert sink_calls[0].mock is True
|
|
await gw.aclose()
|
|
|
|
async def test_telemetry_sink_sync_callable(self) -> None:
|
|
"""Sync sinks should also work (no await needed)."""
|
|
sink_calls: list[TelemetryRow] = []
|
|
|
|
def sync_sink(row: TelemetryRow) -> None:
|
|
sink_calls.append(row)
|
|
|
|
config = mock_config()
|
|
gw = Gateway(config, telemetry_sink=sync_sink)
|
|
await gw.run_task("score", "test prompt")
|
|
assert len(sink_calls) == 1
|
|
await gw.aclose()
|
|
|
|
async def test_no_sink_no_error(self) -> None:
|
|
config = mock_config()
|
|
gw = Gateway(config, telemetry_sink=None)
|
|
result = await gw.run_task("score", "prompt")
|
|
assert result["score"] == 75
|
|
await gw.aclose()
|
|
|
|
async def test_telemetry_row_to_dict(self) -> None:
|
|
row = TelemetryRow(task="score", model="glm-5.2", provider="mock")
|
|
d = row.to_dict()
|
|
assert d["task"] == "score"
|
|
assert d["model"] == "glm-5.2"
|
|
assert d["provider"] == "mock"
|
|
assert "id" in d
|
|
|
|
|
|
# ---- Config tests ----
|
|
|
|
|
|
class TestGatewayConfig:
|
|
def test_from_env_mock_mode(self) -> None:
|
|
env = {}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.mock_mode is True
|
|
assert config.primary.model == "glm-5.2"
|
|
|
|
def test_from_env_real_mode(self) -> None:
|
|
env = {
|
|
"LLM_PRIMARY_KEY": "test-key",
|
|
"LLM_PRIMARY_BASE_URL": "https://api.example.com/v1",
|
|
"LLM_PRIMARY_MODEL": "custom-model",
|
|
}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.mock_mode is False
|
|
assert config.primary.api_key == "test-key"
|
|
assert config.primary.model == "custom-model"
|
|
|
|
def test_from_env_ollama_key(self) -> None:
|
|
env = {"OLLAMA_API_KEY": "ollama-key-123"}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.mock_mode is False
|
|
assert config.primary.api_key == "ollama-key-123"
|
|
|
|
def test_from_env_budgets(self) -> None:
|
|
env = {"LLM_BUDGET_SCORE": "500"}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.budgets["score"] == 500
|
|
|
|
def test_task_class_routing(self) -> None:
|
|
config = mock_config()
|
|
assert config.get_task_class("score") == TaskClass.CHEAP
|
|
assert config.get_task_class("extract") == TaskClass.CHEAP
|
|
assert config.get_task_class("cv_assist") == TaskClass.CHEAP
|
|
assert config.get_task_class("critique") == TaskClass.STRONG
|
|
assert config.get_task_class("cl_critique") == TaskClass.STRONG
|
|
|
|
def test_get_model_routing(self) -> None:
|
|
config = mock_config(cheap_model="cheap-model", strong_model="strong-model")
|
|
assert config.get_model("score") == "cheap-model"
|
|
assert config.get_model("critique") == "strong-model"
|
|
|
|
def test_paid_fallback_detection(self) -> None:
|
|
"""Config should detect paid provider URLs."""
|
|
env = {
|
|
"LLM_PRIMARY_KEY": "key",
|
|
"LLM_FALLBACK_BASE_URL": "https://api.openai.com/v1",
|
|
"LLM_FALLBACK_KEY": "fb-key",
|
|
}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.fallback is not None
|
|
assert config.fallback.is_paid is True
|
|
|
|
def test_paid_fallback_not_used_for_cheap(self) -> None:
|
|
"""The gateway must not route cheap tasks to paid fallback.
|
|
|
|
We verify this by checking _select_provider returns primary for cheap.
|
|
"""
|
|
primary = ProviderConfig(
|
|
name="primary", base_url="https://a.com/v1", api_key="k", model="m"
|
|
)
|
|
fallback = ProviderConfig(
|
|
name="fallback",
|
|
base_url="https://api.openai.com/v1",
|
|
api_key="k2",
|
|
model="m2",
|
|
is_paid=True,
|
|
)
|
|
config = GatewayConfig(primary=primary, fallback=fallback)
|
|
gw = Gateway(config)
|
|
provider = gw._select_provider("score")
|
|
assert provider.name == "primary"
|
|
# The _select_provider method enforces this by always returning primary
|
|
# for cheap tasks, never the paid fallback.
|
|
|
|
def test_non_paid_fallback(self) -> None:
|
|
"""Non-paid (e.g. ollama) fallback is fine."""
|
|
env = {
|
|
"LLM_PRIMARY_KEY": "key",
|
|
"LLM_FALLBACK_BASE_URL": "https://api.ollama-cloud.com/v1",
|
|
"LLM_FALLBACK_KEY": "fb-key",
|
|
}
|
|
config = GatewayConfig.from_env(env=env)
|
|
assert config.fallback is not None
|
|
assert config.fallback.is_paid is False
|
|
|
|
|
|
# ---- Provider call tests (with mocked HTTP) ----
|
|
|
|
|
|
class TestProviderCalls:
|
|
async def test_real_mode_calls_provider(self) -> None:
|
|
"""In non-mock mode, the gateway should make an HTTP call."""
|
|
config = real_config()
|
|
mock_response_data = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps({"score": 85, "rationale": {"ok": True}})
|
|
}
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 50, "completion_tokens": 30},
|
|
"model": "glm-5.2",
|
|
}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, json=mock_response_data)
|
|
|
|
client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://mock.example.com",
|
|
)
|
|
gw = Gateway(config, http_client=client)
|
|
result = await gw.run_task("score", "Score this job")
|
|
assert result["score"] == 85
|
|
assert "_usage" not in result # _usage should be popped
|
|
await client.aclose()
|
|
|
|
async def test_retry_on_429(self) -> None:
|
|
"""Gateway should retry on 429 then succeed."""
|
|
config = real_config(max_retries=2)
|
|
call_count = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count < 2:
|
|
return httpx.Response(429, json={"error": "rate limited"})
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"choices": [
|
|
{"message": {"content": json.dumps({"score": 50})}}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5},
|
|
},
|
|
)
|
|
|
|
client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://mock.example.com",
|
|
)
|
|
gw = Gateway(config, http_client=client)
|
|
result = await gw.run_task("score", "test")
|
|
assert result["score"] == 50
|
|
assert call_count == 2
|
|
await client.aclose()
|
|
|
|
async def test_retry_exhausted_raises(self) -> None:
|
|
"""After all retries, ProviderError should be raised."""
|
|
config = real_config(max_retries=1)
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(500, json={"error": "server error"})
|
|
|
|
client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://mock.example.com",
|
|
)
|
|
from llm_gateway.exceptions import ProviderError
|
|
|
|
gw = Gateway(config, http_client=client)
|
|
with pytest.raises(ProviderError):
|
|
await gw.run_task("critique", "test")
|
|
await client.aclose()
|
|
|
|
async def test_fallback_used_for_strong_task(self) -> None:
|
|
"""Strong tasks should fall back when primary fails."""
|
|
primary = ProviderConfig(
|
|
name="primary",
|
|
base_url="https://primary.example.com/v1",
|
|
api_key="pk",
|
|
model="glm-5.2",
|
|
)
|
|
fallback = ProviderConfig(
|
|
name="fallback",
|
|
base_url="https://fallback.example.com/v1",
|
|
api_key="fk",
|
|
model="glm-5.2",
|
|
is_paid=False,
|
|
)
|
|
config = GatewayConfig(
|
|
primary=primary,
|
|
fallback=fallback,
|
|
budgets={"critique": 4000, "default": 4000},
|
|
max_retries=0,
|
|
)
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
if "primary.example.com" in str(request.url):
|
|
return httpx.Response(500, json={"error": "primary down"})
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"choices": [
|
|
{"message": {"content": json.dumps({"comments": []})}}
|
|
],
|
|
"usage": {"prompt_tokens": 20, "completion_tokens": 10},
|
|
"model": "glm-5.2",
|
|
},
|
|
)
|
|
|
|
client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
)
|
|
gw = Gateway(config, http_client=client)
|
|
result = await gw.run_task("critique", "review this")
|
|
assert result == {"comments": []}
|
|
await client.aclose()
|
|
|
|
async def test_no_fallback_for_cheap_task(self) -> None:
|
|
"""Cheap tasks must not use fallback even when primary fails."""
|
|
primary = ProviderConfig(
|
|
name="primary",
|
|
base_url="https://primary.example.com/v1",
|
|
api_key="pk",
|
|
model="glm-5.2",
|
|
)
|
|
fallback = ProviderConfig(
|
|
name="fallback",
|
|
base_url="https://fallback.example.com/v1",
|
|
api_key="fk",
|
|
model="glm-5.2",
|
|
)
|
|
config = GatewayConfig(
|
|
primary=primary,
|
|
fallback=fallback,
|
|
budgets={"score": 4000, "default": 4000},
|
|
max_retries=0,
|
|
)
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(500, json={"error": "down"})
|
|
|
|
client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
)
|
|
from llm_gateway.exceptions import ProviderError
|
|
|
|
gw = Gateway(config, http_client=client)
|
|
with pytest.raises(ProviderError):
|
|
await gw.run_task("score", "score this")
|
|
await client.aclose()
|
|
|
|
|
|
# ---- Convenience function test ----
|
|
|
|
|
|
class TestRunTaskFunction:
|
|
async def test_run_task_convenience_mock(self) -> None:
|
|
result = await run_task("score", "test", config=mock_config())
|
|
assert result["score"] == 75 |