from queue import SimpleQueue from typing import Final import httpx import pytest from litellm.proxy.lens.models import ( Claim, Execution, ExecutionContent, ModelRequest, ModelResult, Result, Sample, TracePart, ) from litellm.proxy.lens.state import queue_job from litellm.proxy.lens.worker import LensWorker, failure_message from tests.unit.proxy.lens.test_state import NOW, lens @pytest.mark.asyncio @pytest.mark.parametrize("failure", (429, 502, 503, 504, "timeout", 402, 409, 401)) async def test_model_retries_transient_failures_but_not_budget_or_revocation(failure: int | str) -> None: attempts: Final = SimpleQueue[str]() delays: Final = SimpleQueue[float]() expected: Final = ModelResult(content='{"observations":[]}', cost=0.01) def handle(request: httpx.Request) -> httpx.Response: attempts.put(request.url.path) if attempts.qsize() == 1: if failure == "timeout": raise httpx.ReadTimeout("upstream timeout", request=request) assert isinstance(failure, int) return httpx.Response(failure) return httpx.Response(200, json=expected.model_dump()) async def sleep(delay: float) -> None: delays.put(delay) async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: worker: Final = LensWorker(client, sleep=sleep) if failure in (402, 409, 401): with pytest.raises(httpx.HTTPStatusError): await worker.model_request("/model", ModelRequest(purpose="extract", prompt="review")) assert attempts.qsize() == 1 and delays.empty() else: assert await worker.model_request("/model", ModelRequest(purpose="extract", prompt="review")) == expected assert attempts.qsize() == 2 assert delays.get_nowait() == 1 and delays.empty() @pytest.mark.asyncio async def test_transient_retries_are_bounded() -> None: attempts: Final = SimpleQueue[str]() delays: Final = SimpleQueue[float]() def handle(request: httpx.Request) -> httpx.Response: attempts.put(request.url.path) return httpx.Response(503) async def sleep(delay: float) -> None: delays.put(delay) async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: with pytest.raises(httpx.HTTPStatusError): await LensWorker(client, sleep=sleep).model_request( "/model", ModelRequest(purpose="extract", prompt="review") ) assert attempts.qsize() == 3 assert tuple(delays.get_nowait() for _ in range(delays.qsize())) == (1, 2) @pytest.mark.asyncio async def test_idle_worker_does_not_start_an_analysis() -> None: def handle(request: httpx.Request) -> httpx.Response: assert request.url.path == "/lens/worker/claim" return httpx.Response(200, content="null") async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: assert await LensWorker(client).run_once() is False @pytest.mark.asyncio @pytest.mark.parametrize("model_status", (200, 402, 503)) async def test_worker_reads_claimed_activity_and_reports_analysis_or_failure(model_status: int) -> None: claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) execution: Final = Execution( id="run", source="traces", trace_id="trace", team_id="alpha", name="review", start_time="", span_count=1 ) sample: Final = Sample(executions=(execution,), eligible=1) content: Final = ExecutionContent( execution=execution, parts=(TracePart(execution_id="run", span_id="span", name="lead", kind="agent", content="Completed"),), ) saved: Final = SimpleQueue[Result]() def handle(request: httpx.Request) -> httpx.Response: match request.url.path: case "/lens/worker/claim": return httpx.Response(200, json=claim.model_dump(mode="json")) case "/lens/worker/lens/job/sample": return httpx.Response(200, json=sample.model_dump(mode="json")) case "/lens/worker/lens/job/content": assert request.url.params["execution_id"] == execution.id return httpx.Response(200, json=content.model_dump(mode="json")) case "/lens/worker/lens/job/model": return httpx.Response( model_status, json=ModelResult(content='{"observations":[],"cannot_assess":false}', cost=0.01).model_dump(), ) case "/lens/worker/lens/job/progress": return httpx.Response(200, json=True) case "/lens/worker/lens/job/result": saved.put(Result.model_validate_json(request.content)) return httpx.Response(200, json=True) case _: pytest.fail(f"Unexpected analyzer request: {request.url.path}") async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: assert await LensWorker(client).run_once() is True result: Final = saved.get_nowait() assert saved.empty() if model_status == 200: assert result.error == "" assert result.coverage.screened == 1 assert result.coverage.unassessable == 0 elif model_status == 402: assert "HTTP 402" in result.error and "remaining budget" in result.error else: assert result.error.startswith("Model request failed (HTTP 503).") @pytest.mark.parametrize("status", (400, 401, 402, 403, 404, 409, 429, 503)) def test_failure_reports_action_and_status_without_private_response_content(status: int) -> None: request: Final = httpx.Request( "POST", "https://private-host.test/lens/worker/private-lens/private-run/model?token=secret" ) response: Final = httpx.Response(status, request=request, text="private trace content and key") error: Final = httpx.HTTPStatusError("private exception details", request=request, response=response) message: Final = failure_message(error) assert message.startswith(f"Model request failed (HTTP {status}).") assert "private" not in message and "secret" not in message @pytest.mark.parametrize( "route,action", (("sample", "Reading trace data"), ("content", "Reading trace data"), ("result", "Saving results")) ) def test_failure_identifies_the_failing_worker_operation(route: str, action: str) -> None: request: Final = httpx.Request("GET", f"https://proxy.test/lens/worker/lens/job/{route}") response: Final = httpx.Response(503, request=request) error: Final = httpx.HTTPStatusError("private body", request=request, response=response) assert failure_message(error).startswith(f"{action} failed (HTTP 503).") def test_connection_timeout_and_invalid_response_have_distinct_private_diagnostics() -> None: assert "connect to the proxy" in failure_message(httpx.ConnectError("private hostname")) assert "timed out" in failure_message(httpx.ReadTimeout("private prompt")) assert "structured JSON" in failure_message(ValueError("private model response"))