litellm/scripts/generate_lens_contract.py
moe-berri e9cfba2c17
feat(lens): isolate ingestion and investigations in a Rust service (#45148)
* feat(lens): isolate trace storage and investigation in a Rust service

* fix(lens): include Rust sources in the image build context

* feat(lens): wire service setup, scoped delivery receipts and lease attempts

* fix(lens): complete service routing and reject stale investigation results

* fix(lens): retry key propagation and validate isolated Compose setup

* fix(lens): seed through isolated ingestion and preserve upstream queue fixes

* chore: sync schema.prisma copies from root

* fix(lens): bind nullable due timestamps as text for Prisma

* chore(ui): remove stale lint suppressions

* fix(lens): address CI failures and review findings

* refactor(lens): remove retired Python worker and run evaluations in Rust

* fix(lens): reuse control connections and satisfy review checks

* test(lens): install and upgrade both Helm charts on Kubernetes

* test(lens): run connection reuse coverage as an integration test

* fix(ui): upgrade Next.js to 16.3.8 security release

* fix(lens): fence stale attempts and preserve reviewed evidence

* Revert "fix(ui): upgrade Next.js to 16.3.8 security release"

This reverts commit 2f79a51b25.

* fix(lens): stop failed investigations and stream history excerpts

* test(lens): cover model tool and result contracts

* test(lens): fix retired routes and reuse installation build artifacts

* test(lens): use portable grep in Helm installation smoke

* test(lens): wait for migrations before forwarding Helm services

* fix(lens): keep failed evidence reads retryable

* fix(lens): preserve sandbox output during process exit

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-10-07 17:18:58 -07:00

110 lines
2.9 KiB
Python

import argparse
import json
from itertools import chain
from pathlib import Path
from typing import Final
from pydantic import BaseModel, JsonValue
from litellm.proxy.lens.agent_contract import (
Candidate,
Checkpoint,
Clusters,
EvidenceReply,
EvidenceRequest,
FindingGroups,
Findings,
PythonAgentTurn,
PythonRequest,
)
from litellm.proxy.lens.models import (
Claim,
ExecutionContent,
Extraction,
ModelRequest,
ModelResult,
Progress,
Result,
Sample,
)
from litellm.proxy.lens.release import PROTOCOL_VERSION
MODELS: Final[tuple[type[BaseModel], ...]] = (
Claim,
ExecutionContent,
Extraction,
ModelRequest,
ModelResult,
Progress,
Result,
Sample,
Candidate,
Clusters,
Findings,
EvidenceRequest,
PythonRequest,
EvidenceReply,
PythonAgentTurn[Extraction],
PythonAgentTurn[Findings],
Checkpoint,
FindingGroups,
)
TARGET: Final = Path(__file__).resolve().parents[1] / "litellm-rust/crates/lens/contract.json"
def draft_seven(value: JsonValue, names: bool = False) -> JsonValue:
if isinstance(value, list):
return [draft_seven(item) for item in value]
if isinstance(value, dict):
fields: Final = {
"items" if name == "prefixItems" and not names else name: draft_seven(
item, not names and name in ("properties", "definitions", "patternProperties")
)
for name, item in value.items()
if names or (name != "title" and not (name == "default" and item is None))
}
return fields
return value
def contract() -> str:
schemas: Final = tuple(model.model_json_schema(ref_template="#/definitions/{model}") for model in MODELS)
definitions: Final = {
**dict(chain.from_iterable(document.get("$defs", {}).items() for document in schemas)),
**{
model.__name__: {key: value for key, value in schema.items() if key != "$defs"}
for model, schema in zip(MODELS, schemas, strict=True)
},
}
return (
json.dumps(
draft_seven(
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "LensProtocol",
"type": "object",
"definitions": definitions,
"x-lens-protocol-version": PROTOCOL_VERSION,
}
),
indent=2,
sort_keys=True,
)
+ "\n"
)
def main() -> None:
parser: Final = argparse.ArgumentParser()
parser.add_argument("--check", action="store_true")
args: Final = parser.parse_args()
generated: Final = contract()
if args.check:
if TARGET.read_text() != generated:
raise SystemExit("Lens contracts changed; run python scripts/generate_lens_contract.py")
return
TARGET.write_text(generated)
if __name__ == "__main__":
main()