diff --git a/litellm/proxy/lens/analysis.py b/litellm/proxy/lens/analysis.py index 7286cec9ba8..914b265977c 100644 --- a/litellm/proxy/lens/analysis.py +++ b/litellm/proxy/lens/analysis.py @@ -1,11 +1,13 @@ import asyncio import json +import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable from contextlib import aclosing +from datetime import datetime, timezone from functools import reduce from itertools import chain, islice from types import MappingProxyType -from typing import Final, Literal, TypeAlias, TypeVar +from typing import Final, Literal, Protocol, TypeAlias, TypeVar from pydantic import Field, TypeAdapter, ValidationError @@ -16,10 +18,14 @@ from .models import ( Execution, ExecutionContent, FindingDraft, + InFlight, ModelRequest, ModelResult, Record, Result, + Review, + ReviewSpan, + ReviewVerdict, RunAssessment, Sample, TracePart, @@ -38,6 +44,7 @@ class Observation(Record): class Extraction(Record): observations: tuple[Observation, ...] = () cannot_assess: bool = False + reasoning: str = Field(default="", max_length=800) class SpanRead(Record): @@ -84,6 +91,8 @@ class Examined(Record): partial: bool cannot_assess: bool error: str = "" + reasoning: str = "" + shown: tuple[TracePart, ...] = () class Investigation(Record): @@ -94,7 +103,17 @@ class Investigation(Record): ModelCall: TypeAlias = Callable[[ModelRequest], Awaitable[ModelResult]] ReadContent: TypeAlias = Callable[[str, str, int], Awaitable[ExecutionContent]] -ReportProgress: TypeAlias = Callable[[str, Coverage], Awaitable[None]] + + +class ReportProgress(Protocol): + def __call__( + self, + stage: str, + coverage: Coverage, + review: Review | None = None, + reading: tuple[InFlight, ...] | None = None, + /, + ) -> Awaitable[None]: ... ResponseT = TypeVar("ResponseT", bound=Record) @@ -326,7 +345,9 @@ async def extract_stored( request: Final = ModelRequest(purpose="extract", prompt=prompt) if must_decide: final: Final = await structured_response(request, Extraction, model) - return TraceReview(observations=final.observations, cannot_assess=final.cannot_assess) + return TraceReview( + observations=final.observations, cannot_assess=final.cannot_assess, reasoning=final.reasoning + ) return await structured_response(request, TraceReview, model) response: TraceReview @@ -375,6 +396,7 @@ async def extract_stored( parts=evidence, partial=page.partial or page.next_cursor is not None or bool(response.reads) or invalid_observations, cannot_assess=not span_count or response.cannot_assess or bool(response.reads) or invalid_observations, + reasoning=response.reasoning, ) reviews: Final = tuple([await examine(catalog) for catalog in store.catalogs(root_count)]) @@ -383,12 +405,52 @@ async def extract_stored( retained: Final = tuple( p for p in chain.from_iterable(r.parts for r in reviews) if p.span_id in cited or not p.parent_span_id ) + leading: Final = MappingProxyType( + { + p.span_id: p + for p in (*((first_root,) if first_root else ()), *(p for p in store.parts() if p.span_id in cited)) + } + ) + shown: Final = islice(chain(leading.values(), (p for p in store.parts() if p.span_id not in leading)), 8) return Examined( execution=execution, observations=observations, parts=tuple(dict.fromkeys((*retained, *((first_root,) if first_root else ())))), partial=any(r.partial for r in reviews), cannot_assess=not reviews or all(r.cannot_assess for r in reviews), + reasoning=" ".join(r.reasoning for r in reviews if r.reasoning), + shown=tuple(p.model_copy(update=MappingProxyType({"content": overview_content(p, root_count)})) for p in shown), + ) + + +def review_of(examined: Examined, model: str, duration_ms: int, at: datetime) -> Review: + execution: Final = examined.execution + cited: Final = frozenset( + (e.execution_id, e.span_id) for e in chain.from_iterable(o.evidence for o in examined.observations) + ) + return Review( + execution_id=execution.id, + trace_id=execution.trace_id, + agent=execution.service or execution.name, + name=execution.name, + spans=tuple( + ReviewSpan( + span_id=p.span_id, + name=p.name[:120], + kind=p.kind[:40], + preview=p.content[:240], + cited=(p.execution_id, p.span_id) in cited, + ) + for p in examined.shown[:8] + ), + reasoning=examined.reasoning[:800], + verdicts=tuple( + ReviewVerdict(check_id=o.check_id, kind=o.kind, summary=o.summary[:300]) for o in examined.observations + ), + cannot_assess=examined.cannot_assess, + model=model, + duration_ms=max(duration_ms, 0), + at=at, ) @@ -630,8 +692,29 @@ async def analyze_sample( ) ) + def original(identity: str) -> str: + return originals[identity].id + + async def progress_original( + stage: str, coverage: Coverage, review: Review | None = None, reading: tuple[InFlight, ...] | None = None, / + ) -> None: + await progress( + stage, + coverage, + review and review.model_copy(update=MappingProxyType({"execution_id": original(review.execution_id)})), + None + if reading is None + else tuple( + r.model_copy(update=MappingProxyType({"execution_id": original(r.execution_id)})) for r in reading + ), + ) + result: Final = await _analyze_sample( - claim, sample.model_copy(update=MappingProxyType({"executions": executions})), read_alias, model, progress + claim, + sample.model_copy(update=MappingProxyType({"executions": executions})), + read_alias, + model, + progress_original, ) return result.model_copy( update=MappingProxyType( @@ -851,16 +934,36 @@ async def merge_candidates( async def examine_executions( claim: Claim, sample: Sample, read: ReadContent, model: ModelCall, progress: ReportProgress ) -> AsyncIterator[Examined]: - async def examine(execution: Execution) -> Examined: - return await extract(claim, execution, read, model) + reading: tuple[InFlight, ...] = () # rebind-ok: the in-flight set changes as each read starts and finishes + screened = 0 # rebind-ok: counts finished reads for progress + reporting: Final = asyncio.Lock() + + async def report(change: Callable[[tuple[InFlight, ...]], tuple[InFlight, ...]], review: Review | None) -> None: + nonlocal reading + async with reporting: + reading = change(reading) + coverage: Final = Coverage(eligible=sample.eligible, selected=len(sample.executions), screened=screened) + await progress("Reading executions", coverage, review, reading) + + async def examine(execution: Execution) -> tuple[Examined, Review]: + entry: Final = InFlight( + execution_id=execution.id, + trace_id=execution.trace_id, + agent=execution.service or execution.name, + started_at=datetime.now(timezone.utc), + ) + await report(lambda current: (*current, entry), None) + started: Final = time.perf_counter() + examined: Final = await extract(claim, execution, read, model) + elapsed: Final = round((time.perf_counter() - started) * 1000) + return examined, review_of(examined, claim.job.settings.model, elapsed, datetime.now(timezone.utc)) await progress("Reading executions", Coverage(eligible=sample.eligible, selected=len(sample.executions))) - completed: Final = iter(range(1, len(sample.executions) + 1)) async with aclosing(concurrent_results(sample.executions, examine, claim.job.settings.concurrency)) as results: - async for item in results: - await progress( - "Reading executions", - Coverage(eligible=sample.eligible, selected=len(sample.executions), screened=next(completed)), + async for item, review in results: + screened += 1 + await report( + lambda current, done=item.execution.id: tuple(r for r in current if r.execution_id != done), review ) yield item diff --git a/litellm/proxy/lens/endpoints.py b/litellm/proxy/lens/endpoints.py index 12119f07962..b04d3616cc0 100644 --- a/litellm/proxy/lens/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -35,10 +35,10 @@ from litellm.proxy.lens.models import ( ModelResult, Progress, Result, + ReviewPage, RunRequest, Sample, Scope, - Step, WatchAllResult, WatchSkipped, Worker, @@ -48,16 +48,20 @@ from litellm.proxy.lens.release import PROTOCOL_VERSION, release_tag, worker_ima from litellm.proxy.lens.repository import LensRepository, WriterDatabase from litellm.proxy.lens.sources import ActivityAvailability, SourceReader, Storage, parse_execution from litellm.proxy.lens.state import ( - add_step, + apply_progress, can_access, + cancel_job, claim_job, current_job, + end_job, merge_finding, next_scan_start, queue_job, replace_job, + reviews_after, scheduled_window, snapshot_finding, + summarized, ) from litellm.proxy.tracing_runtime import provide_storage @@ -200,7 +204,7 @@ async def validate_workers(settings: LensSettings, scope: Scope) -> None: async def list_lenses(auth: Auth, storage: StorageDep) -> LensList: scope: Final = user_scope(auth) return LensList( - lenses=tuple(e for e in await repository().lenses() if can_access(scope, e.scope)), + lenses=tuple(summarized(e) for e in await repository().lenses() if can_access(scope, e.scope)), workers=tuple(w for w in await repository().workers() if can_access(scope, w.scope)), tracing_enabled=storage is not None, ) @@ -330,7 +334,7 @@ async def run_lens(lens_id: str, body: RunRequest, auth: Auth) -> Lens: @router.get("/{lens_id}", response_model=Lens) async def read_lens(lens_id: str, auth: Auth) -> Lens: - return await get_lens(lens_id, user_scope(auth)) + return summarized(await get_lens(lens_id, user_scope(auth))) @router.get("/{lens_id}/runs", response_model=tuple[Job, ...]) @@ -351,23 +355,17 @@ async def read_run(lens_id: str, job_id: str, auth: Auth) -> Job: return job +@router.get("/{lens_id}/runs/{job_id}/reviews", response_model=ReviewPage) +async def read_reviews(lens_id: str, job_id: str, auth: Auth, after: int = Query(default=0, ge=0)) -> ReviewPage: + return reviews_after(await read_run(lens_id, job_id, auth), after) + + @router.post("/{lens_id}/cancel", response_model=Lens) async def cancel_lens(lens_id: str, auth: Auth) -> Lens: await get_lens(lens_id, user_scope(auth, write=True)) now: Final = datetime.now(timezone.utc) - def cancel(e: Lens) -> Lens: - job: Final = current_job(e) - if job is None: - return e - cancelled: Final = job.model_copy( - update=MappingProxyType({"status": "cancelled", "stage": "Cancelled", "finished_at": now}) - ) - return replace_job(e, cancelled).model_copy( - update=MappingProxyType({"next_run_at": now + timedelta(minutes=e.settings.interval_minutes)}) - ) - - return required(await repository().update(lens_id, cancel)) + return required(await repository().update(lens_id, lambda e: cancel_job(e, now))) @router.patch("/{lens_id}/findings/{finding_id}", response_model=Lens) @@ -505,15 +503,7 @@ async def progress(lens_id: str, job_id: str, body: Progress, worker: WorkerAuth job: Final = current_job(e) if job is None or job.id != job_id or job.worker_id != worker.id: return e - renewed: Final = job.model_copy( - update=MappingProxyType( - {"stage": body.stage, "coverage": body.coverage, "lease_until": now + timedelta(minutes=5)} - ) - ) - return replace_job( - e, - renewed if body.stage == job.stage else add_step(renewed, Step(at=now, kind="stage", label=body.stage)), - ) + return replace_job(e, apply_progress(job, body, now)) required(await repository().update(lens_id, renew)) await repository().heartbeat(worker.id, now.isoformat()) @@ -639,12 +629,9 @@ async def result(lens_id: str, job_id: str, body: Result, worker: WorkerAuth, st merged_ids: Final = frozenset(f.id for f in merged) return replace_job( e, - active.model_copy( + end_job(active, "failed" if body.error else "completed", now).model_copy( update=MappingProxyType( { - "status": "failed" if body.error else "completed", - "stage": "Failed" if body.error else "Complete", - "finished_at": now, "coverage": active.coverage if body.error else body.coverage, "error": body.error, "assessments": body.assessments, diff --git a/litellm/proxy/lens/models.py b/litellm/proxy/lens/models.py index 90c92cc7acd..cc414173c2c 100644 --- a/litellm/proxy/lens/models.py +++ b/litellm/proxy/lens/models.py @@ -207,6 +207,49 @@ class Step(Record): cost: float = 0 +MAX_REVIEWS = 60 + + +class ReviewSpan(Record): + span_id: str + name: str = Field(max_length=120) + kind: str = Field(max_length=40) + preview: str = Field(max_length=240) + cited: bool = False + + +class ReviewVerdict(Record): + check_id: str + kind: Literal["issue", "pattern"] + summary: str = Field(max_length=300) + + +class Review(Record): + execution_id: str + trace_id: str + agent: str + name: str + spans: tuple[ReviewSpan, ...] = Field(default=(), max_length=8) + reasoning: str = Field(default="", max_length=800) + verdicts: tuple[ReviewVerdict, ...] = () + cannot_assess: bool = False + model: str + duration_ms: int = Field(ge=0) + at: datetime + + +class ReviewPage(Record): + reviews: tuple[Review, ...] + reviewed: int + + +class InFlight(Record): + execution_id: str + trace_id: str + agent: str + started_at: datetime + + class Job(Record): id: str status: Literal["queued", "running", "completed", "failed", "cancelled"] = "queued" @@ -227,6 +270,9 @@ class Job(Record): findings: tuple[Finding, ...] | None = None assessments: tuple[RunAssessment, ...] = () steps: tuple[Step, ...] = () + reviews: tuple[Review, ...] = () + reviewed: int = 0 + reading: tuple[InFlight, ...] = () trigger: Literal["schedule", "manual"] = "schedule" @@ -307,6 +353,8 @@ class Claim(Record): class Progress(Record): stage: str = Field() coverage: Coverage = Coverage() + review: Review | None = None + reading: tuple[InFlight, ...] | None = None class Result(Record): diff --git a/litellm/proxy/lens/prompts/review.md b/litellm/proxy/lens/prompts/review.md index d1a5e590dfd..138f546574f 100644 --- a/litellm/proxy/lens/prompts/review.md +++ b/litellm/proxy/lens/prompts/review.md @@ -26,3 +26,5 @@ If you need more evidence, return reads; otherwise return reads=[] and your fina Carry forward still-valid earlier observations and remove disproved ones. cannot_assess means insufficient evidence to assess this run, not absence of an issue. Never manufacture an issue just to produce a result. +Set reasoning to 1-3 plain sentences: what the agent was asked, what happened, and why your observations follow, or why the run is fine. +Keep reasoning under 800 characters and do not quote any secrets or long trace text in it. diff --git a/litellm/proxy/lens/repository.py b/litellm/proxy/lens/repository.py index 6e1e2da112a..15d2cfc34e6 100644 --- a/litellm/proxy/lens/repository.py +++ b/litellm/proxy/lens/repository.py @@ -1,4 +1,6 @@ +import asyncio import json +import random from collections.abc import AsyncIterator, Awaitable, Callable from types import MappingProxyType from typing import Final, Protocol @@ -19,11 +21,14 @@ class Row(BaseModel): _ROWS: Final = TypeAdapter(tuple[Row, ...]) +UPDATE_ATTEMPTS: Final = 40 +UPDATE_BACKOFF_SECONDS: Final = 0.02 class LensRepository: - def __init__(self, db: Database) -> None: + def __init__(self, db: Database, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None: self.db: Final = db + self.sleep: Final = sleep async def lenses(self) -> tuple[Lens, ...]: rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Lens" ORDER BY id')) @@ -47,12 +52,18 @@ class LensRepository: return lens async def update( - self, lens_id: str, transform: Callable[[Lens], Lens], attempts: int = 8, *, changed_only: bool = False + self, + lens_id: str, + transform: Callable[[Lens], Lens], + attempts: int = UPDATE_ATTEMPTS, + *, + changed_only: bool = False, ) -> Lens | None: - for _ in range(attempts): + for attempt in range(attempts): completed, updated = await self._try_update(lens_id, transform, changed_only) if completed: return updated + await self.sleep(random.uniform(0, UPDATE_BACKOFF_SECONDS * min(attempt + 1, 8))) return None async def _try_update( diff --git a/litellm/proxy/lens/state.py b/litellm/proxy/lens/state.py index 3a26078ec82..83cb20320c5 100644 --- a/litellm/proxy/lens/state.py +++ b/litellm/proxy/lens/state.py @@ -4,12 +4,17 @@ from types import MappingProxyType from typing import Final, Literal from litellm.proxy.lens.models import ( + MAX_REVIEWS, MAX_STEPS, Finding, FindingDraft, Job, Lens, LensSettings, + Progress, + Review, + ReviewPage, + Sample, Scope, Step, Worker, @@ -82,6 +87,45 @@ def add_step(job: Job, step: Step) -> Job: return job.model_copy(update=MappingProxyType({"steps": (*job.steps, step)[-MAX_STEPS:]})) +def end_job(job: Job, status: Literal["completed", "failed", "cancelled"], now: datetime) -> Job: + stage: Final = {"completed": "Complete", "failed": "Failed", "cancelled": "Cancelled"}[status] + return job.model_copy( + update=MappingProxyType({"status": status, "stage": stage, "finished_at": now, "reading": ()}) + ) + + +def cancel_job(lens: Lens, now: datetime) -> Lens: + job: Final = current_job(lens) + if job is None: + return lens + return replace_job(lens, end_job(job, "cancelled", now)).model_copy( + update=MappingProxyType({"next_run_at": now + timedelta(minutes=lens.settings.interval_minutes)}) + ) + + +def apply_progress(job: Job, progress: Progress, now: datetime) -> Job: + updates: Final = MappingProxyType( + { + "stage": progress.stage, + "coverage": progress.coverage, + "lease_until": now + timedelta(minutes=5), + "reading": job.reading if progress.reading is None else progress.reading, + } + ) + renewed: Final = add_review(job.model_copy(update=updates), progress.review) + if progress.stage == job.stage: + return renewed + return add_step(renewed, Step(at=now, kind="stage", label=progress.stage)) + + +def add_review(job: Job, review: Review | None) -> Job: + if review is None: + return job + return job.model_copy( + update=MappingProxyType({"reviews": (*job.reviews, review)[-MAX_REVIEWS:], "reviewed": job.reviewed + 1}) + ) + + def claim_job(lens: Lens, worker: Worker, now: datetime) -> Lens: job: Final = current_job(lens) if job is None or not can_access(worker.scope, lens.scope): @@ -91,15 +135,8 @@ def claim_job(lens: Lens, worker: Worker, now: datetime) -> Lens: if job.attempts >= 3: return replace_job( lens, - job.model_copy( - update=MappingProxyType( - { - "status": "failed", - "stage": "Failed", - "error": "Worker disconnected repeatedly", - "finished_at": now, - } - ) + end_job(job, "failed", now).model_copy( + update=MappingProxyType({"error": "Worker disconnected repeatedly"}) ), ).model_copy(update=MappingProxyType({"next_run_at": now + timedelta(minutes=lens.settings.interval_minutes)})) return replace_job( @@ -112,6 +149,9 @@ def claim_job(lens: Lens, worker: Worker, now: datetime) -> Lens: "worker_id": worker.id, "lease_until": now + timedelta(minutes=5), "attempts": job.attempts + 1, + "reviews": (), + "reviewed": 0, + "reading": (), } ) ), @@ -188,3 +228,22 @@ def snapshot_finding(lens: Lens, draft: FindingDraft, revision: int, now: dateti } ) ) + + +def without_attributes(sample: Sample) -> Sample: + executions: Final = tuple(e.model_copy(update=MappingProxyType({"metadata": ()})) for e in sample.executions) + return sample.model_copy(update=MappingProxyType({"executions": executions})) + + +def summarized_job(job: Job) -> Job: + sample: Final = without_attributes(job.sample) if job.sample else None + return job.model_copy(update=MappingProxyType({"reviews": (), "sample": sample})) + + +def summarized(lens: Lens) -> Lens: + return lens.model_copy(update=MappingProxyType({"jobs": tuple(summarized_job(job) for job in lens.jobs)})) + + +def reviews_after(job: Job, after: int) -> ReviewPage: + first_kept: Final = job.reviewed - len(job.reviews) + return ReviewPage(reviews=job.reviews[max(0, after - first_kept) :], reviewed=job.reviewed) diff --git a/litellm/proxy/lens/worker.py b/litellm/proxy/lens/worker.py index 06455dc1a9a..baeac6c096e 100644 --- a/litellm/proxy/lens/worker.py +++ b/litellm/proxy/lens/worker.py @@ -10,10 +10,25 @@ import httpx from pydantic import BaseModel, ConfigDict, ValidationError from .analysis import AnalysisResponseError, analyze_sample, validation_details -from .models import Claim, Coverage, ExecutionContent, ModelRequest, ModelResult, Progress, Result, Sample +from .models import ( + Claim, + Coverage, + ExecutionContent, + InFlight, + ModelRequest, + ModelResult, + Progress, + Result, + Review, + Sample, +) from .release import PROTOCOL_VERSION, release_tag logger: Final = logging.getLogger("litellm.lens.worker") +MODEL_RETRIES: Final = 4 +MODEL_RETRY_MAX_SECONDS: Final = 60.0 +SLOTS: Final = 3 +POLL_SECONDS: Final = 2.0 class ClaimedJobIdentity(BaseModel): @@ -37,6 +52,17 @@ class ModelErrorEnvelope(BaseModel): detail: PublicModelError +def retry_delay(error: httpx.TransportError | httpx.HTTPStatusError, attempt: int) -> float: + backoff: Final = float(min(2**attempt, MODEL_RETRY_MAX_SECONDS)) + if not isinstance(error, httpx.HTTPStatusError): + return backoff + requested: Final = error.response.headers.get("retry-after", "") + try: + return min(max(float(requested), backoff), MODEL_RETRY_MAX_SECONDS) + except ValueError: + return backoff + + def failure_message(error: Exception) -> str: if isinstance(error, AnalysisResponseError): return str(error) @@ -116,11 +142,23 @@ class LensWorker: 503, 504, ) - if not retryable or attempt >= 2: + if not retryable or attempt >= MODEL_RETRIES: raise - await self.sleep(2**attempt) + await self.sleep(retry_delay(exc, attempt)) return await self.model_request(path, body, attempt + 1) + async def serve(self, slots: int, poll_seconds: float) -> None: + await asyncio.gather(*(self.slot(poll_seconds) for _ in range(slots))) + + async def slot(self, poll_seconds: float) -> None: + while True: + try: + if await self.run_once(): + continue + except (httpx.HTTPError, ValueError) as exc: + logger.warning("Worker could not reach Lens (%s)", type(exc).__name__) + await self.sleep(poll_seconds) + async def report_unreadable_claim(self, identity: ClaimIdentity) -> None: failure: Final = await self.client.post( f"/lens/worker/{identity.lens_id}/{identity.job.id}/result", @@ -169,9 +207,16 @@ class LensWorker: result.raise_for_status() return ExecutionContent.model_validate(result.json()) - async def progress(stage: str, coverage: Coverage) -> None: + async def progress( + stage: str, + coverage: Coverage, + review: Review | None = None, + reading: tuple[InFlight, ...] | None = None, + /, + ) -> None: result: Final = await self.client.post( - prefix + "/progress", json=Progress(stage=stage, coverage=coverage).model_dump() + prefix + "/progress", + json=Progress(stage=stage, coverage=coverage, review=review, reading=reading).model_dump(mode="json"), ) result.raise_for_status() @@ -222,13 +267,7 @@ async def main() -> None: async with httpx.AsyncClient( base_url=url, headers=MappingProxyType({"Authorization": f"Bearer {token}"}), timeout=180 ) as client: - worker: Final = LensWorker(client) - while True: - try: - await worker.run_once() - except (httpx.HTTPError, ValueError) as exc: - logger.warning("Worker could not reach Lens (%s)", type(exc).__name__) - await asyncio.sleep(10) + await LensWorker(client).serve(SLOTS, POLL_SECONDS) if __name__ == "__main__": diff --git a/tests/unit/proxy/lens/test_analysis.py b/tests/unit/proxy/lens/test_analysis.py index 5231000e14b..832a40c0fd3 100644 --- a/tests/unit/proxy/lens/test_analysis.py +++ b/tests/unit/proxy/lens/test_analysis.py @@ -13,8 +13,10 @@ from litellm.proxy.lens.models import ( Evidence, Execution, ExecutionContent, + InFlight, ModelRequest, ModelResult, + Review, Sample, TracePart, ) @@ -66,8 +68,10 @@ async def test_parallel_review_shares_one_model_limit_and_cleans_up(outcome: str finally: exited.put(request.prompt) - async def progress(stage: str, coverage: Coverage) -> None: - if stage == "Reading executions": + async def progress( + stage: str, coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: + if stage == "Reading executions" and (_reading is None or _review is not None): counts.put(coverage.screened) claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) @@ -117,7 +121,9 @@ async def test_independent_investigations_overlap_and_report_completions() -> No async def read(_execution_id: str, _cursor: str, _offset: int) -> ExecutionContent: pytest.fail("Inconclusive decisions must not fetch evidence") - async def progress(stage: str, coverage: Coverage) -> None: + async def progress( + stage: str, coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: assert stage == "Checking original evidence" progress_counts.put(coverage.investigated) @@ -464,7 +470,9 @@ async def test_grouping_consolidates_prior_batches_and_reports_real_progress() - ) stages: Final = iter((0, 1)) - async def progress(stage: str, coverage: Coverage) -> None: + async def progress( + stage: str, coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: assert stage == "Grouping observations" assert coverage.grouping_batches == 2 assert coverage.grouped_batches == next(stages) @@ -558,7 +566,9 @@ async def test_thousands_of_matching_runs_keep_all_members_without_a_growing_mod cost=0, ) - async def progress(_stage: str, coverage: Coverage) -> None: + async def progress( + _stage: str, coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: counts.put(coverage.grouped_batches) batches: Final = observation_batches(observations) @@ -632,7 +642,9 @@ async def test_review_keeps_original_ids_in_per_run_assessments() -> None: async def model(_request: ModelRequest) -> ModelResult: return ModelResult(content='{"observations":[],"cannot_assess":false}', cost=0) - async def progress(_stage: str, _coverage: Coverage) -> None: + async def progress( + _stage: str, _coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: pass claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) @@ -891,7 +903,9 @@ async def test_final_registry_reconciles_patterns_split_across_pages() -> None: ) return ModelResult(content=Clusters(candidates=grouped).model_dump_json(), cost=0) - async def progress(_stage: str, _coverage: Coverage) -> None: + async def progress( + _stage: str, _coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: return None result: Final = await cluster_batches((observations,), model, progress, Coverage()) @@ -918,7 +932,9 @@ async def test_distinct_patterns_are_consolidated_in_batches_without_losing_runs payload: Final = json.loads(request.prompt) return ModelResult(content=json.dumps({"candidates": payload["candidates"]}), cost=0) - async def progress(_stage: str, _coverage: Coverage) -> None: + async def progress( + _stage: str, _coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: pass result: Final = await cluster_batches(observation_batches(observations), model, progress, Coverage()) @@ -950,7 +966,9 @@ async def test_invalid_candidate_response_preserves_other_findings_and_reports_i return ModelResult(content="not JSON", cost=0) return ModelResult(content=json.dumps({"action": "submit", "finding": finding("run").model_dump()}), cost=0) - async def progress(_stage: str, coverage: Coverage) -> None: + async def progress( + _stage: str, coverage: Coverage, _review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: counts.put(coverage.inconclusive) claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) @@ -1196,3 +1214,148 @@ async def test_investigator_can_read_all_evidence_pages_across_successive_span_b assert result.error == "" assert tuple(seen.get_nowait() for _ in range(seen.qsize())) == tuple(p.span_id for p in parts) assert tuple(read_cursors.get_nowait() for _ in range(read_cursors.qsize())) == ("", "span039") + + +def test_review_flags_only_spans_cited_by_this_runs_observations() -> None: + from litellm.proxy.lens.analysis import Observation, review_of + + execution: Final = Execution( + id="run1", + source="traces", + trace_id="trace-1", + team_id="", + name="task", + start_time="", + span_count=3, + service="bot", + ) + shown: Final = tuple( + TracePart(execution_id="run1", span_id=span, name=span, kind="tool", content=f"{span} output") + for span in ("root", "search", "answer") + ) + observation: Final = Observation( + check_id="retries", + summary="Search failed twice", + evidence=( + Evidence(execution_id="run1", span_id="search", quote="search output"), + Evidence(execution_id="other", span_id="answer", quote="answer output"), + ), + ) + examined: Final = Examined( + execution=execution, + observations=(observation,), + parts=shown, + partial=False, + cannot_assess=False, + reasoning="Asked to search; it retried without recovering.", + shown=shown, + ) + review: Final = review_of(examined, "cerebras/model", 42, NOW) + assert tuple((s.span_id, s.cited) for s in review.spans) == (("root", False), ("search", True), ("answer", False)) + assert (review.agent, review.trace_id, review.duration_ms) == ("bot", "trace-1", 42) + assert review.reasoning == examined.reasoning + assert tuple((v.check_id, v.summary) for v in review.verdicts) == (("retries", "Search failed twice"),) + + +@pytest.mark.asyncio +async def test_each_screened_run_reports_a_review_with_the_models_reasoning() -> None: + from litellm.proxy.lens.analysis import analyze_sample + + execution: Final = Execution( + id="opaque-original", source="traces", trace_id="trace", team_id="", name="task", start_time="", span_count=2 + ) + reasoning: Final = "The user asked for a refund; the tool timed out and the agent gave up." + reviews: Final = SimpleQueue[Review]() + + async def read(identity: str, _cursor: str, _offset: int) -> ExecutionContent: + return ExecutionContent( + execution=execution, + parts=( + TracePart(execution_id=identity, span_id="a-root", name="agent", kind="agent", content="Refund please"), + TracePart( + execution_id=identity, + span_id="b-tool", + parent_span_id="a-root", + name="refund", + kind="tool", + content="Tool timeout", + ), + ), + ) + + async def model(request: ModelRequest) -> ModelResult: + if request.purpose == "cluster": + return ModelResult(content='{"candidates":[]}', cost=0) + if request.purpose == "investigate": + return ModelResult(content='{"action":"inconclusive"}', cost=0) + return ModelResult( + content=json.dumps( + { + "reasoning": reasoning, + "observations": [ + { + "check_id": "retries", + "summary": "Gave up after a timeout", + "evidence": [{"execution_id": "r0", "span_id": "b-tool", "quote": "Tool timeout"}], + } + ], + } + ), + cost=0, + ) + + async def progress( + _stage: str, _coverage: Coverage, review: Review | None = None, _reading: tuple[InFlight, ...] | None = None, / + ) -> None: + if review is not None: + reviews.put(review) + + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) + await analyze_sample(claim, Sample(executions=(execution,), eligible=1), read, model, progress) + review: Final = reviews.get_nowait() + assert reviews.empty() + assert review.execution_id == execution.id + assert review.reasoning == reasoning + assert review.model == claim.job.settings.model + assert tuple((s.span_id, s.cited) for s in review.spans) == (("a-root", False), ("b-tool", True)) + assert tuple(v.summary for v in review.verdicts) == ("Gave up after a timeout",) + + +@pytest.mark.asyncio +async def test_a_run_is_reported_in_flight_under_its_original_id_until_its_review_arrives() -> None: + from litellm.proxy.lens.analysis import analyze_sample + + execution: Final = Execution( + id="opaque-original", source="traces", trace_id="trace", team_id="", name="task", start_time="", span_count=1 + ) + reports: Final = SimpleQueue[tuple[str | None, tuple[str, ...] | None]]() + + async def read(identity: str, _cursor: str, _offset: int) -> ExecutionContent: + return ExecutionContent( + execution=execution, + parts=(TracePart(execution_id=identity, span_id="s", name="agent", kind="agent", content="Hi"),), + ) + + async def model(request: ModelRequest) -> ModelResult: + if request.purpose == "cluster": + return ModelResult(content='{"candidates":[]}', cost=0) + return ModelResult(content='{"observations":[]}', cost=0) + + async def progress( + stage: str, _coverage: Coverage, review: Review | None = None, reading: tuple[InFlight, ...] | None = None, / + ) -> None: + if stage == "Reading executions": + reports.put( + ( + review and review.execution_id, + None if reading is None else tuple(f"{r.execution_id}:{r.trace_id}" for r in reading), + ) + ) + + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) + await analyze_sample(claim, Sample(executions=(execution,), eligible=1), read, model, progress) + assert tuple(reports.get_nowait() for _ in range(reports.qsize())) == ( + (None, None), + (None, ("opaque-original:trace",)), + ("opaque-original", ()), + ) diff --git a/tests/unit/proxy/lens/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py index 54dbaac1e42..ec09cce00f4 100644 --- a/tests/unit/proxy/lens/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -10,6 +10,7 @@ from litellm import Router from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.lens.endpoints import ( list_agents, + read_reviews, run_settings, run_window, user_scope, @@ -174,6 +175,15 @@ async def test_incompatible_worker_is_rejected_before_claiming_work( assert "Upgrade" in error.value.detail +@pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM, None)) +@pytest.mark.asyncio +async def test_regular_keys_cannot_poll_live_reviews(role: LitellmUserRoles | None) -> None: + auth: Final = UserAPIKeyAuth(user_role=role, team_id="team", token="hashed-test-key") + with pytest.raises(HTTPException) as error: + await read_reviews("lens", "job", auth) + assert error.value.status_code == 403 + + @pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM, None)) def test_regular_keys_cannot_read_lens_results(role: LitellmUserRoles | None) -> None: auth: Final = UserAPIKeyAuth(user_role=role, team_id="team", token="hashed-test-key") diff --git a/tests/unit/proxy/lens/test_repository.py b/tests/unit/proxy/lens/test_repository.py new file mode 100644 index 00000000000..40f2b83bc1b --- /dev/null +++ b/tests/unit/proxy/lens/test_repository.py @@ -0,0 +1,65 @@ +from datetime import datetime, timezone +from typing import Final + +import pytest + +from litellm.proxy.lens.models import Check, Lens, LensSettings, Scope +from litellm.proxy.lens.repository import UPDATE_ATTEMPTS, LensRepository, Row + +NOW: Final = datetime(2026, 1, 15, tzinfo=timezone.utc) +STORED: Final = Lens( + id="lens", + scope=Scope(team_id="alpha"), + settings=LensSettings( + name="Swarm", model="cerebras/gpt-oss-120b", checks=(Check(id="c", instruction="Find loops"),) + ), + created_at=NOW, + next_run_at=NOW, + budget_month=NOW.strftime("%Y-%m"), +) + + +class ContendedDatabase: + def __init__(self, losses: int) -> None: + self.losses: Final = losses + self.writes = 0 # rebind-ok: counts write attempts made under contention + + async def query_raw(self, query: str, *args: object) -> object: + if query.startswith("SELECT data FROM"): + return (Row(data=STORED.model_dump(mode="json")),) + self.writes += 1 + return (Row(data=1 if self.writes > self.losses else 0),) + + async def execute_raw(self, query: str, *args: object) -> int: + return 0 + + +async def no_wait(_: float) -> None: + return None + + +def renamed(lens: Lens) -> Lens: + return lens.model_copy(update={"settings": lens.settings.model_copy(update={"name": "Swarm (renamed)"})}) + + +@pytest.mark.asyncio +async def test_update_survives_the_contention_of_a_fast_model_writing_every_review() -> None: + db: Final = ContendedDatabase(losses=12) + updated: Final = await LensRepository(db, sleep=no_wait).update("lens", renamed) + assert updated is not None + assert updated.settings.name == "Swarm (renamed)" + assert db.writes == 13 + + +@pytest.mark.asyncio +async def test_update_backs_off_between_lost_writes_and_gives_up_after_the_limit() -> None: + waits: list[float] = [] # mutable-ok: records each backoff the repository requests + + async def record(seconds: float) -> None: + waits.append(seconds) + + db: Final = ContendedDatabase(losses=UPDATE_ATTEMPTS) + assert await LensRepository(db, sleep=record).update("lens", renamed) is None + assert db.writes == UPDATE_ATTEMPTS + assert len(waits) == UPDATE_ATTEMPTS + assert all(w >= 0 for w in waits) diff --git a/tests/unit/proxy/lens/test_state.py b/tests/unit/proxy/lens/test_state.py index f5ad36ecfeb..0066b40a633 100644 --- a/tests/unit/proxy/lens/test_state.py +++ b/tests/unit/proxy/lens/test_state.py @@ -1,31 +1,46 @@ from datetime import datetime, timedelta, timezone from functools import reduce -from typing import Final +from typing import Final, Literal import pytest from litellm.proxy.lens.models import ( + MAX_REVIEWS, MAX_STEPS, AgentTestCase, Check, Evidence, + Execution, FindingDraft, + InFlight, IssueBrief, + Job, Lens, LensSettings, + MetadataFilter, + Progress, + Review, + Sample, Scope, Step, Worker, ) from litellm.proxy.lens.state import ( + add_review, add_step, + apply_progress, can_access, + cancel_job, claim_job, current_job, + end_job, merge_finding, next_scan_start, queue_job, renew_budget, + replace_job, + reviews_after, + summarized, ) NOW: Final = datetime(2026, 1, 15, tzinfo=timezone.utc) @@ -345,3 +360,128 @@ def test_calendar_overflow_is_rejected_without_the_old_history_and_interval_caps assert getattr(accepted, field) == 100000 with pytest.raises(ValidationError, match="supported calendar range"): LensSettings.model_validate({**lens().settings.model_dump(), field: 10**30}) + + +def review(index: int) -> Review: + return Review( + execution_id=f"run-{index}", trace_id="t", agent="support", name="task", model="analysis", duration_ms=1, at=NOW + ) + + +def test_reviews_keep_the_newest_window_while_counting_every_review() -> None: + job: Final = queue_job(lens(), NOW, "job").jobs[0] + grown: Final = reduce(add_review, tuple(review(i) for i in range(MAX_REVIEWS + 3)), job) + assert grown.reviewed == MAX_REVIEWS + 3 + assert len(grown.reviews) == MAX_REVIEWS + assert grown.reviews[0].execution_id == "run-3" + assert grown.reviews[-1].execution_id == f"run-{MAX_REVIEWS + 2}" + + +def test_reclaimed_run_starts_its_review_history_over() -> None: + queued: Final = queue_job(lens(), NOW, "job") + first: Final = claim_job(queued, worker(), NOW) + reviewed: Final = replace_job(first, reduce(add_review, (review(0), review(1)), first.jobs[0])) + stalled: Final = reviewed.jobs[0].model_copy( + update={"reading": (InFlight(execution_id="run-2", trace_id="t", agent="support", started_at=NOW),)} + ) + reclaimed: Final = claim_job(replace_job(reviewed, stalled), worker(identity="other"), NOW + timedelta(minutes=6)) + job: Final = reclaimed.jobs[0] + assert job.worker_id == "other" + assert (job.reviews, job.reviewed, job.reading) == ((), 0, ()) + replayed: Final = reduce(add_review, (review(0), review(1)), job) + assert replayed.reviewed == len(replayed.reviews) == 2 + + +def test_progress_without_a_review_leaves_the_review_history_alone() -> None: + job: Final = add_review(queue_job(lens(), NOW, "job").jobs[0], review(0)) + assert add_review(job, None) == job + + +def reviewed_job() -> Job: + execution: Final = Execution( + id="run-0", + source="traces", + trace_id="t", + team_id="alpha", + name="task", + start_time="2026-01-15 00:00:00", + span_count=3, + service="support", + metadata=(MetadataFilter(key="gen_ai.agent.name", value="support"),), + ) + job: Final = ( + queue_job(lens(), NOW, "job") + .jobs[0] + .model_copy(update={"sample": Sample(executions=(execution,), eligible=4, selected=1)}) + ) + timed: Final = tuple(review(i).model_copy(update={"at": NOW + timedelta(seconds=i)}) for i in range(3)) + return reduce(add_review, timed, job) + + +def test_summary_drops_reviews_and_run_attributes_but_keeps_counts_and_run_identity() -> None: + job: Final = reviewed_job() + listed: Final = summarized(lens().model_copy(update={"jobs": (job,)})).jobs[0] + assert listed.reviews == () + assert listed.reviewed == job.reviewed == 3 + assert listed.sample is not None and job.sample is not None + assert listed.sample.executions[0].metadata == () + assert ( + listed.sample.executions[0].model_copy(update={"metadata": job.sample.executions[0].metadata}) + == (job.sample.executions[0]) + ) + assert listed.model_copy(update={"reviews": job.reviews, "sample": job.sample}) == job + + +def test_review_polling_returns_only_reviews_after_the_cursor_even_when_they_finished_out_of_order() -> None: + job: Final = reduce(add_review, (review(5).model_copy(update={"at": NOW - timedelta(hours=1)}),), reviewed_job()) + assert reviews_after(job, 0).reviews == job.reviews + assert [r.execution_id for r in reviews_after(job, 2).reviews] == ["run-2", "run-5"] + assert reviews_after(job, 4).reviews == () + assert reviews_after(job, 4).reviewed == 4 + + +def test_review_polling_after_the_window_moved_on_returns_what_is_still_kept() -> None: + job: Final = reduce(add_review, tuple(review(i) for i in range(MAX_REVIEWS + 10)), reviewed_job()) + page: Final = reviews_after(job, 5) + assert page.reviews == job.reviews + assert page.reviewed == MAX_REVIEWS + 13 + assert [r.execution_id for r in reviews_after(job, page.reviewed - 2).reviews] == [ + f"run-{MAX_REVIEWS + 8}", + f"run-{MAX_REVIEWS + 9}", + ] + + +def in_flight(execution: str) -> InFlight: + return InFlight(execution_id=execution, trace_id="t", agent="support", started_at=NOW) + + +def reading_job() -> Job: + running: Final = claim_job(queue_job(lens(), NOW, "job"), worker(), NOW).jobs[0] + return apply_progress(running, Progress(stage=running.stage, reading=(in_flight("a"), in_flight("b"))), NOW) + + +def test_progress_replaces_the_in_flight_runs_and_old_workers_leave_them_alone() -> None: + job: Final = reading_job() + assert [r.execution_id for r in job.reading] == ["a", "b"] + finished: Final = apply_progress(job, Progress(stage=job.stage, review=review(0), reading=(in_flight("b"),)), NOW) + assert [r.execution_id for r in finished.reading] == ["b"] + assert finished.reviewed == 1 + assert apply_progress(job, Progress(stage=job.stage, review=review(1)), NOW).reading == job.reading + assert apply_progress(job, Progress(stage=job.stage, reading=()), NOW).reading == () + + +@pytest.mark.parametrize("status", ("completed", "failed", "cancelled")) +def test_finished_jobs_stop_showing_runs_in_flight(status: Literal["completed", "failed", "cancelled"]) -> None: + ended: Final = end_job(reading_job(), status, NOW) + assert ended.status == status + assert ended.finished_at == NOW + assert ended.reading == () + + +def test_cancel_and_repeated_disconnects_clear_runs_in_flight() -> None: + reading: Final = replace_job(queue_job(lens(), NOW, "job"), reading_job()) + cancelled: Final = cancel_job(reading, NOW).jobs[0] + assert (cancelled.status, cancelled.reading) == ("cancelled", ()) + abandoned: Final = reading.model_copy(update={"jobs": (reading.jobs[0].model_copy(update={"attempts": 3}),)}) + expired: Final = claim_job(abandoned, worker(), NOW + timedelta(minutes=10)).jobs[0] + assert (expired.status, expired.reading) == ("failed", ()) diff --git a/tests/unit/proxy/lens/test_worker.py b/tests/unit/proxy/lens/test_worker.py index 7983aec8af4..4d656da1956 100644 --- a/tests/unit/proxy/lens/test_worker.py +++ b/tests/unit/proxy/lens/test_worker.py @@ -12,12 +12,19 @@ from litellm.proxy.lens.models import ( ExecutionContent, ModelRequest, ModelResult, + Progress, Result, Sample, TracePart, ) from litellm.proxy.lens.state import queue_job -from litellm.proxy.lens.worker import LensWorker, failure_message +from litellm.proxy.lens.worker import ( + MODEL_RETRIES, + MODEL_RETRY_MAX_SECONDS, + LensWorker, + failure_message, + retry_delay, +) from tests.unit.proxy.lens.test_state import NOW, lens @@ -69,8 +76,44 @@ async def test_transient_retries_are_bounded() -> None: 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) + assert attempts.qsize() == MODEL_RETRIES + 1 + assert tuple(delays.get_nowait() for _ in range(delays.qsize())) == tuple( + float(min(2**n, MODEL_RETRY_MAX_SECONDS)) for n in range(MODEL_RETRIES) + ) + + +@pytest.mark.asyncio +async def test_rate_limited_model_waits_as_long_as_the_provider_asks_then_completes() -> 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() <= 3: + return httpx.Response(429, headers={"retry-after": "30"}) + 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: + result: Final = await LensWorker(client, sleep=sleep).model_request( + "/model", ModelRequest(purpose="extract", prompt="review") + ) + assert result == expected + assert tuple(delays.get_nowait() for _ in range(delays.qsize())) == (30, 30, 30) + + +@pytest.mark.parametrize( + ("retry_after", "attempt", "expected"), + (("", 1, 2), ("5", 0, 5), ("1", 3, 8), ("9999", 0, MODEL_RETRY_MAX_SECONDS), ("soon", 2, 4)), +) +def test_retry_delay_prefers_the_providers_wait_within_bounds(retry_after: str, attempt: int, expected: float) -> None: + request: Final = httpx.Request("POST", "https://proxy.test/model") + headers: Final = {"retry-after": retry_after} if retry_after else {} + error: Final = httpx.HTTPStatusError("limited", request=request, response=httpx.Response(429, headers=headers)) + assert retry_delay(error, attempt) == expected @pytest.mark.asyncio @@ -419,6 +462,72 @@ async def test_transient_heartbeat_failure_recovers_without_cancelling_analysis( assert attempts.qsize() == 2 and saved.empty() +@pytest.mark.asyncio +async def test_worker_sends_each_runs_review_with_its_progress() -> 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="t", team_id="", name="task", start_time="", span_count=1 + ) + sent: Final = SimpleQueue[Progress]() + + def handle(request: httpx.Request) -> httpx.Response: + match request.url.path.rsplit("/", 1)[-1]: + case "claim": + return httpx.Response(200, json=claim.model_dump(mode="json")) + case "sample": + return httpx.Response(200, json=Sample(executions=(execution,), eligible=1).model_dump()) + case "content": + return httpx.Response( + 200, + json=ExecutionContent( + execution=execution, + parts=(TracePart(execution_id="run", span_id="s", name="step", kind="agent", content="Done"),), + ).model_dump(), + ) + case "model": + return httpx.Response( + 200, json={"content": '{"observations":[],"reasoning":"Finished the task."}', "cost": 0} + ) + case "progress": + sent.put(Progress.model_validate_json(request.content)) + return httpx.Response(200, json=True) + case "result": + return httpx.Response(200, json=True) + case _: + pytest.fail(f"Unexpected worker 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() + reviews: Final = tuple(p.review for p in (sent.get_nowait() for _ in range(sent.qsize())) if p.review) + assert tuple((r.execution_id, r.reasoning) for r in reviews) == (("run", "Finished the task."),) + + +@pytest.mark.asyncio +async def test_worker_runs_investigations_in_parallel_and_polls_quickly_when_idle() -> None: + claims: Final = SimpleQueue[str]() + running: Final = asyncio.Event() + waits: Final = SimpleQueue[float]() + + class Worker(LensWorker): + async def run_once(self) -> bool: + claims.put("claim") + if claims.qsize() <= 2: + if claims.qsize() == 2: + running.set() + await running.wait() + return True + raise asyncio.CancelledError + + async def sleep(delay: float) -> None: + waits.put(delay) + + async with httpx.AsyncClient(base_url="https://proxy.test") as client: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(Worker(client, sleep=sleep).serve(slots=2, poll_seconds=2), timeout=1) + assert running.is_set() + assert waits.empty() + + @pytest.mark.asyncio async def test_worker_announces_release_and_waits_on_incompatible_gateway( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture diff --git a/ui/litellm-dashboard/public/assets/logos/cerebras.svg b/ui/litellm-dashboard/public/assets/logos/cerebras.svg index 426f6430c23..5f2fdabe845 100644 --- a/ui/litellm-dashboard/public/assets/logos/cerebras.svg +++ b/ui/litellm-dashboard/public/assets/logos/cerebras.svg @@ -1,89 +1,89 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index cfa85817ea6..ccc8e7414ed 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -67,6 +67,16 @@ --animate-slot-slide-in: slot-slide-in 150ms cubic-bezier(0, 0, 0.2, 1) both; --animate-trace-drawer-in: trace-drawer-in 200ms cubic-bezier(0.25, 1, 0.5, 1) both; --animate-trace-drawer-out: trace-drawer-out 200ms cubic-bezier(0.4, 0, 1, 1) both; + --animate-lens-shimmer: lens-shimmer 1.1s cubic-bezier(0.4, 0, 0.2, 1) infinite; + + @keyframes lens-shimmer { + from { + transform: translateX(-100%); + } + to { + transform: translateX(250%); + } + } @keyframes trace-drawer-in { from { diff --git a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx index c68a1112c8b..4f3e372e333 100644 --- a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx @@ -1,103 +1,63 @@ "use client"; import { Tabs as TabsPrimitive } from "@base-ui/react/tabs"; -import { Activity, ScanSearch, Settings } from "lucide-react"; +import { Settings } from "lucide-react"; import { StatusDot } from "@/components/shared/StatusDot"; import { cn } from "@/lib/cva.config"; import type { InvestigationActivity } from "./model/status"; import { useWorkerConnected } from "./hooks/useWorkerConnected"; import type { LensList } from "./model/types"; -import { LENS_TABS, type LensTab } from "./route"; -import { frameCorner, frameTab } from "./ui/frame"; +import { LENS_TABS } from "./route"; -const MODE_ICONS = { traces: Activity, investigations: ScanSearch, settings: Settings } as const; - -const ACTIVITY_DOT: Record, { className: string; label: string }> = { - running: { className: "bg-info motion-safe:animate-pulse", label: "An investigation is running" }, - queued: { className: "bg-muted-foreground/60", label: "An investigation is queued" }, -}; - -function ActivityDot({ activity }: { activity: InvestigationActivity }) { - if (activity === "idle") return null; - return ( -