267 lines
8.7 KiB
Python
267 lines
8.7 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from src.services import ai_router as ai_router_module
|
|
from src.services.ai_providers import ollama as ollama_module
|
|
from src.services.ai_providers.interfaces import AIResult
|
|
from src.services.ai_providers.ollama import _gcp_cloud_input_block_reason
|
|
from src.services.ai_router import AIProviderRegistry, AIRouterExecutor
|
|
|
|
|
|
def _chat_context(prompt: str) -> dict[str, Any]:
|
|
trace_id = "sre-chat-trace:test"
|
|
run_id = "sre-chat-run:test"
|
|
work_item_id = "SRE-CHAT-ADVISORY"
|
|
return {
|
|
"trace_id": trace_id,
|
|
"run_id": run_id,
|
|
"work_item_id": work_item_id,
|
|
"intent_hint": "chat",
|
|
"task_type": "chat",
|
|
"data_classification": "sanitized",
|
|
"provider_timeout_seconds": 35,
|
|
"infrastructure_remediation_write_allowed": False,
|
|
"executor_invocation_allowed": False,
|
|
"agent99_dispatch_allowed": False,
|
|
"cross_domain_execution_allowed": False,
|
|
"sre_chat_sanitization_receipt": {
|
|
"schema_version": "sre_chat_sanitization_receipt_v1",
|
|
"status": "verified",
|
|
"sanitizer": "src.services.sanitization_service.sanitize",
|
|
"source_label": "awoooi_sre_group_chat",
|
|
"data_classification": "sanitized",
|
|
"prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(),
|
|
"trace_id": trace_id,
|
|
"run_id": run_id,
|
|
"work_item_id": work_item_id,
|
|
"raw_input_forwarded": False,
|
|
"raw_log_payload_forwarded": False,
|
|
"secret_value_exposed": False,
|
|
"private_network_value_exposed": False,
|
|
"infrastructure_remediation_write_allowed": False,
|
|
"provider_timeout_seconds": 35,
|
|
},
|
|
}
|
|
|
|
|
|
class _NoCacheRedis:
|
|
async def get(self, _key: str) -> None:
|
|
return None
|
|
|
|
async def set(self, _key: str, _value: str, ex: int | None = None) -> None:
|
|
del ex
|
|
|
|
|
|
class _Provider:
|
|
is_enabled = True
|
|
capabilities = {"chat"}
|
|
|
|
def __init__(self, name: str, *, success: bool) -> None:
|
|
self.name = name
|
|
self.success = success
|
|
self.calls = 0
|
|
self.privacy_level = "local" if name == "ollama_local" else "cloud"
|
|
|
|
async def analyze(
|
|
self,
|
|
_prompt: str,
|
|
_context: dict[str, Any] | None = None,
|
|
) -> AIResult:
|
|
self.calls += 1
|
|
return AIResult(
|
|
raw_response=f"{self.name}-result" if self.success else "",
|
|
success=self.success,
|
|
provider=self.name,
|
|
error=None if self.success else "forced_failure",
|
|
)
|
|
|
|
|
|
def _executor(*providers: _Provider) -> AIRouterExecutor:
|
|
registry = AIProviderRegistry()
|
|
for provider in providers:
|
|
registry.register(provider)
|
|
return AIRouterExecutor(registry)
|
|
|
|
|
|
def test_ai_router_accepts_prompt_bound_real_chat_receipt_without_synthetic() -> None:
|
|
prompt = "sanitized real SRE conversation"
|
|
context = _chat_context(prompt)
|
|
|
|
(
|
|
routed_prompt,
|
|
routed_context,
|
|
reason,
|
|
) = AIRouterExecutor._paid_alert_execution_inputs(prompt, context)
|
|
|
|
assert reason is None
|
|
assert routed_prompt == prompt
|
|
assert routed_context == context
|
|
assert "synthetic_validation" not in context
|
|
assert "synthetic_prompt_sha256" not in context
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("mutation", "expected_reason"),
|
|
[
|
|
("missing", "cloud_sanitization_receipt_missing"),
|
|
("prompt_hash", "sre_chat_sanitized_prompt_mismatch"),
|
|
("status", "sre_chat_sanitization_receipt_invalid"),
|
|
("run_binding", "sre_chat_sanitization_receipt_invalid"),
|
|
("write_boundary", "sre_chat_sanitization_receipt_invalid"),
|
|
("private_exposure", "sre_chat_sanitization_receipt_invalid"),
|
|
("timeout", "sre_chat_sanitization_receipt_invalid"),
|
|
],
|
|
)
|
|
def test_ai_router_real_chat_receipt_missing_or_tampered_fails_closed(
|
|
mutation: str,
|
|
expected_reason: str,
|
|
) -> None:
|
|
prompt = "sanitized real SRE conversation"
|
|
context = deepcopy(_chat_context(prompt))
|
|
receipt = context["sre_chat_sanitization_receipt"]
|
|
if mutation == "missing":
|
|
context.pop("sre_chat_sanitization_receipt")
|
|
elif mutation == "prompt_hash":
|
|
receipt["prompt_sha256"] = "0" * 64
|
|
elif mutation == "status":
|
|
receipt["status"] = "caller_claimed"
|
|
elif mutation == "run_binding":
|
|
receipt["run_id"] = "another-run"
|
|
elif mutation == "write_boundary":
|
|
context["executor_invocation_allowed"] = True
|
|
elif mutation == "private_exposure":
|
|
receipt["private_network_value_exposed"] = True
|
|
elif mutation == "timeout":
|
|
context["provider_timeout_seconds"] = 300
|
|
|
|
(
|
|
routed_prompt,
|
|
routed_context,
|
|
reason,
|
|
) = AIRouterExecutor._paid_alert_execution_inputs(prompt, context)
|
|
|
|
assert routed_prompt is None
|
|
assert routed_context is None
|
|
assert reason == expected_reason
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ai_router_executor_calls_gcp_for_verified_real_chat_receipt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
prompt = "sanitized real SRE conversation"
|
|
gcp_a = _Provider("ollama_gcp_a", success=True)
|
|
local = _Provider("ollama_local", success=True)
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr(
|
|
"src.core.redis_client.get_redis",
|
|
lambda: _NoCacheRedis(),
|
|
)
|
|
|
|
result = await _executor(gcp_a, local).execute(
|
|
prompt,
|
|
provider_order=["ollama_gcp_a", "ollama_gcp_b", "ollama_local"],
|
|
context=_chat_context(prompt),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "ollama_gcp_a"
|
|
assert gcp_a.calls == 1
|
|
assert local.calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("mutation", ["missing", "prompt_hash"])
|
|
async def test_ai_router_executor_never_calls_gcp_with_invalid_real_chat_receipt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mutation: str,
|
|
) -> None:
|
|
prompt = "sanitized real SRE conversation"
|
|
context = deepcopy(_chat_context(prompt))
|
|
if mutation == "missing":
|
|
context.pop("sre_chat_sanitization_receipt")
|
|
else:
|
|
context["sre_chat_sanitization_receipt"]["prompt_sha256"] = "0" * 64
|
|
gcp_a = _Provider("ollama_gcp_a", success=True)
|
|
gcp_b = _Provider("ollama_gcp_b", success=True)
|
|
local = _Provider("ollama_local", success=True)
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr(
|
|
"src.core.redis_client.get_redis",
|
|
lambda: _NoCacheRedis(),
|
|
)
|
|
|
|
result = await _executor(gcp_a, gcp_b, local).execute(
|
|
prompt,
|
|
provider_order=["ollama_gcp_a", "ollama_gcp_b", "ollama_local"],
|
|
context=context,
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "ollama_local"
|
|
assert gcp_a.calls == 0
|
|
assert gcp_b.calls == 0
|
|
assert local.calls == 1
|
|
|
|
|
|
async def _verified_transport(**_kwargs: Any) -> str | None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gcp_ollama_verifier_accepts_prompt_bound_real_chat_receipt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
endpoint = "http://gcp-a-chat.test:11435"
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", endpoint)
|
|
prompt = "sanitized real SRE conversation"
|
|
|
|
reason = await _gcp_cloud_input_block_reason(
|
|
prompt=prompt,
|
|
context=_chat_context(prompt),
|
|
endpoint_url=endpoint,
|
|
receipt_verifier=_verified_transport,
|
|
)
|
|
|
|
assert reason is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("mutation", "expected_reason"),
|
|
[
|
|
("missing", "gcp_cloud_sanitization_receipt_missing"),
|
|
("prompt_hash", "gcp_sre_chat_sanitized_prompt_mismatch"),
|
|
("run_binding", "gcp_sre_chat_sanitization_receipt_invalid"),
|
|
],
|
|
)
|
|
async def test_gcp_ollama_verifier_rejects_missing_or_tampered_chat_receipt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mutation: str,
|
|
expected_reason: str,
|
|
) -> None:
|
|
endpoint = "http://gcp-a-chat.test:11435"
|
|
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", endpoint)
|
|
prompt = "sanitized real SRE conversation"
|
|
context = deepcopy(_chat_context(prompt))
|
|
if mutation == "missing":
|
|
context.pop("sre_chat_sanitization_receipt")
|
|
elif mutation == "prompt_hash":
|
|
context["sre_chat_sanitization_receipt"]["prompt_sha256"] = "0" * 64
|
|
elif mutation == "run_binding":
|
|
context["sre_chat_sanitization_receipt"]["run_id"] = "another-run"
|
|
|
|
reason = await _gcp_cloud_input_block_reason(
|
|
prompt=prompt,
|
|
context=context,
|
|
endpoint_url=endpoint,
|
|
receipt_verifier=_verified_transport,
|
|
)
|
|
|
|
assert reason == expected_reason
|