Celes Renata
a72f336ad1
Multi-stage evidence-grounded inference architecture replacing the monolithic 9B model extraction pipeline. CPU-first specialist services handle routine extraction while the 9B vLLM model is preserved for semantic adjudication of ambiguous cases. Key components: - Capability-aware inference gateway (OpenAI-compatible + Ollama) - Endpoint registry with DB migrations and REST API - Sentence-aware document segmenter (property tests) - Deterministic financial parsing with offset integrity - Symbol resolution with ambiguity detection - Specialist service (GLiNER2, dynamic batching, K8s deployment) - Company-specific sentiment (FinBERT, calibration) - Retrieval-based novelty and duplicate detection - Confidence calibration pipeline - Deterministic routing engine (property tests) - 9B adjudication layer with VRAM gating - Stock-specific impact model (features, labels, baseline, trained) - Pipeline orchestrator (state machine, queues, leases, feature flags) - Bounded parallelism (async workers, semaphore, load shedding) - Observability (tracing, metrics, alerts) - Compatibility adapter (v3→v2 golden mapping tests) - Shadow/canary promotion framework - Active learning and fine-tuning pipeline Test results: 1,161 tests pass, ruff lint clean. All 282 spec tasks completed.
496 lines
17 KiB
Python
496 lines
17 KiB
Python
"""Tests for the InferenceGateway facade, lineage recording, and adapters.
|
|
|
|
Covers:
|
|
- Gateway creates correct client type per protocol
|
|
- Gateway reuses clients for same endpoint
|
|
- Target refresh invalidates cached client
|
|
- Lineage recording captures all required fields
|
|
- Extraction adapter returns lineage metadata
|
|
- Unknown protocols fail closed
|
|
|
|
Requirements: 2.1, 2.6, 2.12, 13.6
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from services.shared.inference.gateway import InferenceGateway
|
|
from services.shared.inference.lineage import (
|
|
build_lineage_from_result,
|
|
lineage_to_persistence_dict,
|
|
)
|
|
from services.shared.inference.models import (
|
|
ChatMessage,
|
|
InferenceResult,
|
|
InferenceTarget,
|
|
ModelLineage,
|
|
ProviderCapabilities,
|
|
StructuredGenerationRequest,
|
|
TokenUsage,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_target(
|
|
protocol: str = "openai_chat",
|
|
endpoint_id: uuid.UUID | None = None,
|
|
deployment_id: uuid.UUID | None = None,
|
|
model: str = "test-model",
|
|
) -> InferenceTarget:
|
|
"""Create a minimal InferenceTarget for testing."""
|
|
return InferenceTarget(
|
|
endpoint_id=endpoint_id or uuid.uuid4(),
|
|
deployment_id=deployment_id or uuid.uuid4(),
|
|
protocol=protocol,
|
|
base_url="http://localhost:8000",
|
|
model=model,
|
|
capabilities=ProviderCapabilities(
|
|
chat_completions=True,
|
|
json_schema=True,
|
|
usage=True,
|
|
),
|
|
)
|
|
|
|
|
|
def _make_request() -> StructuredGenerationRequest:
|
|
"""Create a minimal StructuredGenerationRequest."""
|
|
return StructuredGenerationRequest(
|
|
messages=[ChatMessage(role="user", content="hello")],
|
|
max_output_tokens=256,
|
|
)
|
|
|
|
|
|
def _make_inference_result(
|
|
endpoint_id: uuid.UUID | None = None,
|
|
deployment_id: uuid.UUID | None = None,
|
|
model: str = "test-model",
|
|
protocol: str = "openai_chat",
|
|
) -> InferenceResult:
|
|
"""Create a typical InferenceResult."""
|
|
return InferenceResult(
|
|
content='{"answer": 42}',
|
|
parsed={"answer": 42},
|
|
endpoint_id=endpoint_id or uuid.uuid4(),
|
|
deployment_id=deployment_id or uuid.uuid4(),
|
|
model=model,
|
|
protocol=protocol,
|
|
structured_mode="json_schema",
|
|
latency_ms=150,
|
|
usage=TokenUsage(input_tokens=50, output_tokens=20, total_tokens=70),
|
|
request_id="req-001",
|
|
retries=0,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Gateway: correct client type per protocol
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGatewayClientCreation:
|
|
"""Gateway creates the correct client type based on protocol."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_creates_openai_client_for_openai_chat(self) -> None:
|
|
"""openai_chat protocol creates OpenAICompatibleClient."""
|
|
gateway = InferenceGateway()
|
|
target = _make_target(protocol="openai_chat")
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate",
|
|
new_callable=AsyncMock,
|
|
return_value=_make_inference_result(
|
|
endpoint_id=target.endpoint_id,
|
|
deployment_id=target.deployment_id,
|
|
),
|
|
):
|
|
result = await gateway.generate(target, _make_request())
|
|
assert result.protocol == "openai_chat"
|
|
|
|
await gateway.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_creates_ollama_client_for_ollama_native(self) -> None:
|
|
"""ollama_native protocol creates OllamaNativeClient."""
|
|
gateway = InferenceGateway()
|
|
target = _make_target(protocol="ollama_native")
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.ollama_native.OllamaNativeClient.generate",
|
|
new_callable=AsyncMock,
|
|
return_value=_make_inference_result(
|
|
endpoint_id=target.endpoint_id,
|
|
deployment_id=target.deployment_id,
|
|
protocol="ollama_native",
|
|
),
|
|
):
|
|
result = await gateway.generate(target, _make_request())
|
|
assert result.protocol == "ollama_native"
|
|
|
|
await gateway.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unknown_protocol_fails_closed(self) -> None:
|
|
"""Unknown protocol raises ValueError — never silently routes to Ollama."""
|
|
gateway = InferenceGateway()
|
|
# Use a type-ignore here since we're intentionally passing an invalid protocol
|
|
target = InferenceTarget(
|
|
endpoint_id=uuid.uuid4(),
|
|
deployment_id=uuid.uuid4(),
|
|
protocol="unknown_protocol", # type: ignore[arg-type]
|
|
base_url="http://localhost:8000",
|
|
model="test",
|
|
capabilities=ProviderCapabilities(),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Unknown inference protocol"):
|
|
await gateway.generate(target, _make_request())
|
|
|
|
await gateway.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Gateway: client reuse
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGatewayClientReuse:
|
|
"""Gateway reuses clients for the same endpoint."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reuses_client_for_same_endpoint(self) -> None:
|
|
"""Repeated calls with the same target reuse the cached client."""
|
|
gateway = InferenceGateway()
|
|
endpoint_id = uuid.uuid4()
|
|
target = _make_target(endpoint_id=endpoint_id)
|
|
|
|
mock_result = _make_inference_result(endpoint_id=endpoint_id)
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_result,
|
|
):
|
|
await gateway.generate(target, _make_request())
|
|
await gateway.generate(target, _make_request())
|
|
|
|
# Only one client should be cached
|
|
assert len(gateway._clients) == 1
|
|
assert endpoint_id in gateway._clients
|
|
|
|
await gateway.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_creates_separate_clients_for_different_endpoints(self) -> None:
|
|
"""Different endpoint IDs get separate cached clients."""
|
|
gateway = InferenceGateway()
|
|
target_a = _make_target(endpoint_id=uuid.uuid4())
|
|
target_b = _make_target(endpoint_id=uuid.uuid4())
|
|
|
|
mock_result_a = _make_inference_result(endpoint_id=target_a.endpoint_id)
|
|
mock_result_b = _make_inference_result(endpoint_id=target_b.endpoint_id)
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate",
|
|
new_callable=AsyncMock,
|
|
side_effect=[mock_result_a, mock_result_b],
|
|
):
|
|
await gateway.generate(target_a, _make_request())
|
|
await gateway.generate(target_b, _make_request())
|
|
|
|
assert len(gateway._clients) == 2
|
|
assert target_a.endpoint_id in gateway._clients
|
|
assert target_b.endpoint_id in gateway._clients
|
|
|
|
await gateway.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Gateway: target refresh
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGatewayTargetRefresh:
|
|
"""Target refresh invalidates cached client."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_invalidates_cached_client(self) -> None:
|
|
"""After refresh, the next call creates a new client."""
|
|
gateway = InferenceGateway()
|
|
endpoint_id = uuid.uuid4()
|
|
target = _make_target(endpoint_id=endpoint_id)
|
|
mock_result = _make_inference_result(endpoint_id=endpoint_id)
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_result,
|
|
), patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.close",
|
|
new_callable=AsyncMock,
|
|
) as mock_close:
|
|
await gateway.generate(target, _make_request())
|
|
assert endpoint_id in gateway._clients
|
|
|
|
await gateway.refresh_target(endpoint_id)
|
|
assert endpoint_id not in gateway._clients
|
|
mock_close.assert_called_once()
|
|
|
|
await gateway.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_nonexistent_endpoint_is_safe(self) -> None:
|
|
"""Refreshing an endpoint that isn't cached does nothing."""
|
|
gateway = InferenceGateway()
|
|
# Should not raise
|
|
await gateway.refresh_target(uuid.uuid4())
|
|
await gateway.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lineage recording
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLineageRecording:
|
|
"""Lineage recording captures all required fields."""
|
|
|
|
def test_build_lineage_from_result_captures_all_fields(self) -> None:
|
|
"""All lineage fields are extracted from InferenceResult."""
|
|
eid = uuid.uuid4()
|
|
did = uuid.uuid4()
|
|
result = InferenceResult(
|
|
content="test",
|
|
endpoint_id=eid,
|
|
deployment_id=did,
|
|
model="qwen-9b",
|
|
protocol="openai_chat",
|
|
structured_mode="json_schema",
|
|
latency_ms=300,
|
|
request_id="req-abc",
|
|
retries=1,
|
|
)
|
|
|
|
lineage = build_lineage_from_result(result, trace_id="trace-123")
|
|
|
|
assert lineage.endpoint_id == eid
|
|
assert lineage.deployment_id == did
|
|
assert lineage.model == "qwen-9b"
|
|
assert lineage.protocol == "openai_chat"
|
|
assert lineage.structured_mode == "json_schema"
|
|
assert lineage.request_id == "req-abc"
|
|
assert lineage.latency_ms == 300
|
|
assert lineage.retries == 1
|
|
assert lineage.trace_id == "trace-123"
|
|
|
|
def test_lineage_to_persistence_dict_maps_protocol(self) -> None:
|
|
"""Protocol is mapped to human-friendly model_provider for persistence."""
|
|
lineage = ModelLineage(
|
|
endpoint_id=uuid.uuid4(),
|
|
deployment_id=uuid.uuid4(),
|
|
model="test-model",
|
|
protocol="openai_chat",
|
|
structured_mode="json_schema",
|
|
request_id="req-1",
|
|
latency_ms=100,
|
|
retries=0,
|
|
trace_id="t-1",
|
|
)
|
|
|
|
d = lineage_to_persistence_dict(lineage)
|
|
|
|
assert d["model_provider"] == "openai_compatible"
|
|
assert d["model_name"] == "test-model"
|
|
assert d["protocol"] == "openai_chat"
|
|
assert d["endpoint_id"] is not None
|
|
assert d["deployment_id"] is not None
|
|
assert d["structured_mode"] == "json_schema"
|
|
assert d["request_id"] == "req-1"
|
|
assert d["latency_ms"] == 100
|
|
assert d["retries"] == 0
|
|
assert d["trace_id"] == "t-1"
|
|
|
|
def test_lineage_ollama_protocol_maps_to_ollama_provider(self) -> None:
|
|
"""ollama_native protocol maps to 'ollama' provider."""
|
|
lineage = ModelLineage(
|
|
model="qwen-9b",
|
|
protocol="ollama_native",
|
|
)
|
|
d = lineage_to_persistence_dict(lineage)
|
|
assert d["model_provider"] == "ollama"
|
|
|
|
def test_lineage_specialist_protocol_maps_to_specialist_provider(self) -> None:
|
|
"""specialist_http protocol maps to 'specialist' provider."""
|
|
lineage = ModelLineage(
|
|
model="gliner2-large",
|
|
protocol="specialist_http",
|
|
)
|
|
d = lineage_to_persistence_dict(lineage)
|
|
assert d["model_provider"] == "specialist"
|
|
|
|
def test_lineage_serialization(self) -> None:
|
|
"""ModelLineage serializes all fields via model_dump()."""
|
|
lineage = ModelLineage(
|
|
endpoint_id=uuid.uuid4(),
|
|
deployment_id=uuid.uuid4(),
|
|
model="test",
|
|
protocol="openai_chat",
|
|
structured_mode="json_object",
|
|
request_id="r-1",
|
|
latency_ms=42,
|
|
retries=2,
|
|
trace_id="t-abc",
|
|
)
|
|
data = lineage.model_dump()
|
|
assert data["model"] == "test"
|
|
assert data["protocol"] == "openai_chat"
|
|
assert data["trace_id"] == "t-abc"
|
|
assert data["latency_ms"] == 42
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Extraction adapter: lineage metadata
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractionAdapterLineage:
|
|
"""Extraction adapter returns lineage metadata."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_document_returns_lineage(self) -> None:
|
|
"""extract_document bundles lineage with the extraction response."""
|
|
from services.extractor.inference_adapter import extract_document
|
|
|
|
gateway = InferenceGateway()
|
|
endpoint_id = uuid.uuid4()
|
|
deployment_id = uuid.uuid4()
|
|
target = _make_target(
|
|
endpoint_id=endpoint_id,
|
|
deployment_id=deployment_id,
|
|
)
|
|
|
|
# Mock the gateway to return a valid extraction JSON
|
|
extraction_json = '{"summary":"test","companies":[],"macro_themes":[],"novelty_score":0.5,"confidence":0.8,"extraction_warnings":[]}'
|
|
mock_result = InferenceResult(
|
|
content=extraction_json,
|
|
endpoint_id=endpoint_id,
|
|
deployment_id=deployment_id,
|
|
model="test-model",
|
|
protocol="openai_chat",
|
|
structured_mode="json_schema",
|
|
latency_ms=200,
|
|
usage=TokenUsage(input_tokens=100, output_tokens=50),
|
|
request_id="req-ext-1",
|
|
retries=0,
|
|
)
|
|
|
|
with patch.object(
|
|
gateway,
|
|
"generate",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_result,
|
|
):
|
|
result = await extract_document(
|
|
gateway=gateway,
|
|
target=target,
|
|
document_text="Test document text for extraction.",
|
|
document_id="doc-123",
|
|
)
|
|
|
|
# Verify lineage is populated
|
|
assert result.lineage is not None
|
|
assert result.lineage.endpoint_id == endpoint_id
|
|
assert result.lineage.deployment_id == deployment_id
|
|
assert result.lineage.model == "test-model"
|
|
assert result.lineage.protocol == "openai_chat"
|
|
assert result.lineage.request_id == "req-ext-1"
|
|
assert result.lineage.latency_ms == 200
|
|
|
|
await gateway.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_document_lineage_on_failure(self) -> None:
|
|
"""Lineage is still captured even when extraction fails."""
|
|
from services.extractor.inference_adapter import extract_document
|
|
|
|
gateway = InferenceGateway()
|
|
endpoint_id = uuid.uuid4()
|
|
deployment_id = uuid.uuid4()
|
|
target = _make_target(
|
|
endpoint_id=endpoint_id,
|
|
deployment_id=deployment_id,
|
|
)
|
|
|
|
# Mock the gateway to return an error result
|
|
mock_result = InferenceResult(
|
|
content="",
|
|
endpoint_id=endpoint_id,
|
|
deployment_id=deployment_id,
|
|
model="test-model",
|
|
protocol="openai_chat",
|
|
structured_mode="json_schema",
|
|
latency_ms=5000,
|
|
request_id="req-timeout",
|
|
retries=3,
|
|
error="Request timed out",
|
|
error_category="timeout",
|
|
)
|
|
|
|
with patch.object(
|
|
gateway,
|
|
"generate",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_result,
|
|
):
|
|
result = await extract_document(
|
|
gateway=gateway,
|
|
target=target,
|
|
document_text="Some text",
|
|
document_id="doc-fail",
|
|
max_retries=0, # no retries for test speed
|
|
)
|
|
|
|
# Extraction failed but lineage is still captured
|
|
assert not result.response.success
|
|
assert result.lineage.endpoint_id == endpoint_id
|
|
assert result.lineage.model == "test-model"
|
|
assert result.lineage.protocol == "openai_chat"
|
|
|
|
await gateway.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Gateway: active_endpoints property
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGatewayProperties:
|
|
"""Gateway exposes useful state for monitoring."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_active_endpoints_tracks_cached_clients(self) -> None:
|
|
"""active_endpoints shows all cached endpoint IDs."""
|
|
gateway = InferenceGateway()
|
|
eid = uuid.uuid4()
|
|
target = _make_target(endpoint_id=eid)
|
|
mock_result = _make_inference_result(endpoint_id=eid)
|
|
|
|
with patch(
|
|
"services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_result,
|
|
):
|
|
await gateway.generate(target, _make_request())
|
|
|
|
assert eid in gateway.active_endpoints
|
|
assert len(gateway.active_endpoints) == 1
|
|
|
|
await gateway.close()
|
|
assert len(gateway.active_endpoints) == 0
|