358 lines
11 KiB
Python
358 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from src.services.ai_providers import ollama as ollama_module
|
|
from src.services.ai_providers.ollama import (
|
|
OllamaGcpBProvider,
|
|
OllamaLocalProvider,
|
|
OllamaProvider,
|
|
_safe_ollama_error_code,
|
|
)
|
|
|
|
|
|
class _FakeRegistry:
|
|
def get_model(self, provider: str, use_case: str) -> str:
|
|
return "qwen3:14b"
|
|
|
|
def get_provider_options(self, provider: str) -> dict[str, Any]:
|
|
return {"num_predict": 32, "temperature": 0.1, "top_p": 0.9}
|
|
|
|
|
|
class _FakeResponse:
|
|
status_code = 200
|
|
|
|
def raise_for_status(self) -> None:
|
|
return None
|
|
|
|
def json(self) -> dict[str, Any]:
|
|
return {
|
|
"response": '{"summary":"ok"}',
|
|
"eval_count": 4,
|
|
"prompt_eval_count": 3,
|
|
}
|
|
|
|
|
|
class _FakeClient:
|
|
def __init__(self) -> None:
|
|
self.posted_urls: list[str] = []
|
|
self.posted_payloads: list[dict[str, Any]] = []
|
|
self.checked_urls: list[str] = []
|
|
|
|
async def post(self, url: str, **kwargs: Any) -> _FakeResponse:
|
|
self.posted_urls.append(url)
|
|
self.posted_payloads.append(kwargs.get("json", {}))
|
|
return _FakeResponse()
|
|
|
|
async def get(self, url: str, **kwargs: Any) -> _FakeResponse:
|
|
self.checked_urls.append(url)
|
|
return _FakeResponse()
|
|
|
|
|
|
class _FailingClient:
|
|
def __init__(self, error: BaseException) -> None:
|
|
self.error = error
|
|
|
|
async def post(self, _url: str, **_kwargs: Any) -> _FakeResponse:
|
|
raise self.error
|
|
|
|
|
|
def _synthetic_context(prompt: str, **overrides: Any) -> dict[str, Any]:
|
|
return {
|
|
"data_classification": "sanitized",
|
|
"trace_id": "trace-test",
|
|
"run_id": "run-test",
|
|
"work_item_id": "AIA-SRE-013",
|
|
"synthetic_validation": True,
|
|
"synthetic_prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
|
"cloud_transport_receipts": [
|
|
{"provider": "ollama_gcp_a", "receipt_id": "verified-a"},
|
|
{"provider": "ollama_gcp_b", "receipt_id": "verified-b"},
|
|
],
|
|
"infrastructure_remediation_write_allowed": False,
|
|
**overrides,
|
|
}
|
|
|
|
|
|
async def _injected_transport_verifier(**kwargs: Any) -> str | None:
|
|
receipts = (kwargs.get("context") or {}).get("cloud_transport_receipts") or []
|
|
if not any(item.get("provider") == kwargs["provider"] for item in receipts):
|
|
return "gcp_cloud_transport_receipt_missing"
|
|
return None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error", "expected"),
|
|
[
|
|
(
|
|
httpx.RemoteProtocolError(
|
|
"Server disconnected at http://private.invalid/?key=secret-like"
|
|
),
|
|
"provider_transport_error",
|
|
),
|
|
(
|
|
httpx.ReadTimeout(
|
|
"timeout at http://private.invalid/?key=secret-like",
|
|
request=httpx.Request("POST", "http://private.invalid"),
|
|
),
|
|
"provider_timeout",
|
|
),
|
|
(
|
|
httpx.HTTPStatusError(
|
|
"body=secret-like",
|
|
request=httpx.Request("POST", "http://private.invalid"),
|
|
response=httpx.Response(503),
|
|
),
|
|
"http_503",
|
|
),
|
|
(
|
|
RuntimeError("http://private.invalid/?key=secret-like"),
|
|
"unhandled_RuntimeError",
|
|
),
|
|
],
|
|
)
|
|
def test_safe_ollama_error_code_never_returns_raw_error_text(
|
|
error: BaseException,
|
|
expected: str,
|
|
) -> None:
|
|
code = _safe_ollama_error_code(error)
|
|
|
|
assert code == expected
|
|
assert "private.invalid" not in code
|
|
assert "secret-like" not in code
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_local_transport_disconnect_returns_safe_code(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
|
|
monkeypatch.setattr(
|
|
ollama_module.settings,
|
|
"OLLAMA_FALLBACK_URL",
|
|
"http://local-111:11434",
|
|
)
|
|
provider = OllamaLocalProvider()
|
|
client = _FailingClient(
|
|
httpx.RemoteProtocolError(
|
|
"Server disconnected at http://local-111:11434/?key=secret-like"
|
|
)
|
|
)
|
|
|
|
async def _get_client() -> _FailingClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze("sanitized synthetic diagnosis")
|
|
|
|
assert result.success is False
|
|
assert result.provider == "ollama_local"
|
|
assert result.error == "provider_transport_error"
|
|
assert "local-111" not in result.error
|
|
assert "secret-like" not in result.error
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_gcp_b_analyze_uses_secondary_url(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(
|
|
ollama_module.settings,
|
|
"OLLAMA_SECONDARY_URL",
|
|
"http://secondary:11436",
|
|
)
|
|
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
|
|
|
|
client = _FakeClient()
|
|
provider = OllamaGcpBProvider(
|
|
transport_receipt_verifier=_injected_transport_verifier
|
|
)
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze(
|
|
"diagnose",
|
|
context=_synthetic_context("diagnose", task_type="diagnose"),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "ollama_gcp_b"
|
|
assert client.posted_urls == ["http://secondary:11436/api/generate"]
|
|
assert client.posted_payloads[0]["model"] == "qwen3:14b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_gcp_a_allows_heavy_diagnose_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
|
|
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
|
|
|
|
client = _FakeClient()
|
|
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze(
|
|
"diagnose",
|
|
context=_synthetic_context("diagnose", task_type="diagnose"),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert client.posted_urls == ["http://primary:11435/api/generate"]
|
|
assert client.posted_payloads[0]["model"] == "qwen3:14b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_gcp_a_coerces_non_diagnosis_heavy_model_to_health_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_HEALTH_CHECK_MODEL", "gemma3:4b")
|
|
|
|
client = _FakeClient()
|
|
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze(
|
|
"background summary",
|
|
context=_synthetic_context("background summary", task_type="summary"),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert client.posted_urls == ["http://primary:11435/api/generate"]
|
|
assert client.posted_payloads[0]["model"] == "gemma3:4b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_gcp_a_can_explicitly_allow_heavy_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
|
|
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
|
|
|
|
client = _FakeClient()
|
|
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze(
|
|
"deep diagnose",
|
|
context=_synthetic_context(
|
|
"deep diagnose",
|
|
task_type="diagnose",
|
|
allow_gcp_heavy_model=True,
|
|
),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert client.posted_payloads[0]["model"] == "qwen3:14b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gcp_ollama_rejects_unreceipted_prompt_before_network_call(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(
|
|
ollama_module.settings,
|
|
"OLLAMA_SECONDARY_URL",
|
|
"http://secondary:11436",
|
|
)
|
|
client = _FakeClient()
|
|
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
result = await provider.analyze(
|
|
"raw operational log",
|
|
context={"data_classification": "sanitized"},
|
|
)
|
|
|
|
assert result.success is False
|
|
assert result.error == "gcp_cloud_transport_receipt_missing"
|
|
assert result.audit_metadata["generation_attempted"] is False
|
|
assert client.posted_urls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gcp_ollama_blocks_tool_loop_before_tool_or_network_execution(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
|
|
monkeypatch.setattr(
|
|
ollama_module.settings,
|
|
"OLLAMA_SECONDARY_URL",
|
|
"http://secondary:11436",
|
|
)
|
|
client = _FakeClient()
|
|
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
|
|
tool_executor = AsyncMock()
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
prompt = "sanitized investigation"
|
|
|
|
result = await provider.analyze_with_tools(
|
|
prompt,
|
|
available_tools=[object()], # type: ignore[list-item]
|
|
tool_executor=tool_executor,
|
|
context=_synthetic_context(prompt, task_type="diagnose"),
|
|
)
|
|
|
|
assert result.success is False
|
|
assert result.error == "gcp_tool_loop_blocked_unverified_tool_results"
|
|
assert result.audit_metadata == {
|
|
"generation_attempted": False,
|
|
"tool_execution_attempted": False,
|
|
"transport_boundary": "public_http_degraded",
|
|
}
|
|
tool_executor.assert_not_awaited()
|
|
assert client.posted_urls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ollama_base_provider_health_uses_endpoint_hook(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
client = _FakeClient()
|
|
provider = OllamaProvider()
|
|
monkeypatch.setattr(provider, "_endpoint_url", lambda: "http://primary:11435")
|
|
|
|
async def _get_client() -> _FakeClient:
|
|
return client
|
|
|
|
monkeypatch.setattr(provider, "_get_client", _get_client)
|
|
|
|
assert await provider.health_check() is True
|
|
assert client.checked_urls == ["http://primary:11435/api/tags"]
|