test-int/semantic/litellm_live_harness.py
Rama main · réplica de lectura del repositorio original.
"""Opt-in live LiteLLM embedding evaluation harness."""
from __future__ import annotations
import argparse
import asyncio
import json
import math
import os
import sys
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
from rich.console import Console
from rich.table import Table
from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider
RELATED_DOCUMENT = "OAuth login refresh tokens keep an authenticated web session active."
DISTRACTOR_DOCUMENT = "A sourdough starter ferments flour and water before bread baking."
QUERY_TEXT = "authentication login token flow"
@dataclass(frozen=True)
class LiteLLMLiveCase:
"""A real LiteLLM embedding model to exercise end-to-end."""
name: str
model: str
dimensions: int
api_key_env: str | None = None
api_base: str | None = None
document_input_type: str | None = None
query_input_type: str | None = None
forward_dimensions: bool | None = None
@dataclass(frozen=True)
class LiteLLMLiveResult:
"""Measured result for a live LiteLLM embedding case."""
name: str
model: str
dimensions: int
api_key_env: str | None
document_input_type: str | None
query_input_type: str | None
forward_dimensions: bool | None
related_score: float
distractor_score: float
min_norm: float
max_norm: float
embed_documents_latency_ms: float
embed_query_latency_ms: float
total_latency_ms: float
type ProviderFactory = Callable[..., Any]
def _required_string(case_data: Mapping[str, Any], key: str) -> str:
value = case_data.get(key)
if not isinstance(value, str) or not value:
raise ValueError(f"LiteLLM live case must include non-empty string field {key!r}")
return value
def _optional_string(case_data: Mapping[str, Any], key: str) -> str | None:
value = case_data.get(key)
if value is None:
return None
if not isinstance(value, str) or not value:
raise ValueError(f"LiteLLM live case field {key!r} must be a non-empty string")
return value
def _optional_bool(case_data: Mapping[str, Any], key: str) -> bool | None:
value = case_data.get(key)
if value is None:
return None
if not isinstance(value, bool):
raise ValueError(f"LiteLLM live case field {key!r} must be a boolean")
return value
def load_custom_cases(raw: str | None) -> list[LiteLLMLiveCase]:
"""Load additional live model cases from JSON."""
if not raw:
return []
values = json.loads(raw)
if not isinstance(values, list):
raise ValueError("LiteLLM live cases JSON must be an array")
cases: list[LiteLLMLiveCase] = []
for value in values:
if not isinstance(value, dict):
raise ValueError("Each LiteLLM live case must be a JSON object")
case_data: dict[str, Any] = value
dimensions = case_data.get("dimensions")
if type(dimensions) is not int or dimensions <= 0:
raise ValueError("LiteLLM live case must include positive integer field 'dimensions'")
cases.append(
LiteLLMLiveCase(
name=_required_string(case_data, "name"),
model=_required_string(case_data, "model"),
dimensions=dimensions,
api_key_env=_optional_string(case_data, "api_key_env"),
api_base=_optional_string(case_data, "api_base"),
document_input_type=_optional_string(case_data, "document_input_type"),
query_input_type=_optional_string(case_data, "query_input_type"),
forward_dimensions=_optional_bool(case_data, "forward_dimensions"),
)
)
return cases
def configured_cases(
environ: Mapping[str, str] | None = None,
*,
custom_cases_raw: str | None = None,
) -> list[LiteLLMLiveCase]:
"""Return built-in and user-supplied live cases whose credentials are available."""
env = os.environ if environ is None else environ
cases: list[LiteLLMLiveCase] = []
if env.get("OPENAI_API_KEY"):
cases.append(
LiteLLMLiveCase(
name="openai-text-embedding-3-small",
model="openai/text-embedding-3-small",
dimensions=1536,
api_key_env="OPENAI_API_KEY",
)
)
if env.get("COHERE_API_KEY"):
cases.append(
LiteLLMLiveCase(
name="cohere-embed-english-v3",
model="cohere/embed-english-v3.0",
dimensions=1024,
api_key_env="COHERE_API_KEY",
)
)
raw = (
custom_cases_raw
if custom_cases_raw is not None
else env.get("BASIC_MEMORY_TEST_LITELLM_CASES")
)
cases.extend(load_custom_cases(raw))
return cases
def cosine(a: list[float], b: list[float]) -> float:
"""Compute cosine similarity for live ranking sanity checks."""
dot = sum(x * y for x, y in zip(a, b, strict=True))
norm_a = vector_norm(a)
norm_b = vector_norm(b)
if norm_a == 0 or norm_b == 0:
return 0.0
return dot / (norm_a * norm_b)
def vector_norm(vector: list[float]) -> float:
"""Return the Euclidean norm of a vector."""
return math.sqrt(sum(value * value for value in vector))
def assert_valid_vector(vector: list[float], dimensions: int) -> float:
"""Assert provider output is a usable normalized vector and return its norm."""
assert len(vector) == dimensions
assert all(math.isfinite(value) for value in vector)
norm = vector_norm(vector)
assert math.isclose(norm, 1.0, abs_tol=1e-6)
return norm
async def evaluate_case(
case: LiteLLMLiveCase,
*,
environ: Mapping[str, str] | None = None,
provider_factory: ProviderFactory = LiteLLMEmbeddingProvider,
) -> LiteLLMLiveResult:
"""Run one live LiteLLM case and return measured ranking/vector metrics."""
env = os.environ if environ is None else environ
api_key = env.get(case.api_key_env) if case.api_key_env else None
provider = provider_factory(
model_name=case.model,
dimensions=case.dimensions,
batch_size=2,
api_key=api_key,
api_base=case.api_base,
timeout=60.0,
document_input_type=case.document_input_type,
query_input_type=case.query_input_type,
forward_dimensions=case.forward_dimensions,
)
start = time.perf_counter()
documents_start = time.perf_counter()
vectors = await provider.embed_documents([RELATED_DOCUMENT, DISTRACTOR_DOCUMENT])
documents_elapsed = time.perf_counter() - documents_start
query_start = time.perf_counter()
query_vector = await provider.embed_query(QUERY_TEXT)
query_elapsed = time.perf_counter() - query_start
total_elapsed = time.perf_counter() - start
assert len(vectors) == 2
norms = [assert_valid_vector(vector, case.dimensions) for vector in [*vectors, query_vector]]
related_score = cosine(query_vector, vectors[0])
distractor_score = cosine(query_vector, vectors[1])
assert related_score > distractor_score, (
f"{case.name} ranked the related document at {related_score:.4f}, "
f"not above distractor {distractor_score:.4f}"
)
return LiteLLMLiveResult(
name=case.name,
model=case.model,
dimensions=case.dimensions,
api_key_env=case.api_key_env,
document_input_type=case.document_input_type,
query_input_type=case.query_input_type,
forward_dimensions=case.forward_dimensions,
related_score=related_score,
distractor_score=distractor_score,
min_norm=min(norms),
max_norm=max(norms),
embed_documents_latency_ms=documents_elapsed * 1000,
embed_query_latency_ms=query_elapsed * 1000,
total_latency_ms=total_elapsed * 1000,
)
async def evaluate_cases(
cases: Sequence[LiteLLMLiveCase],
*,
environ: Mapping[str, str] | None = None,
provider_factory: ProviderFactory = LiteLLMEmbeddingProvider,
) -> tuple[list[LiteLLMLiveResult], list[tuple[LiteLLMLiveCase, str]]]:
"""Evaluate cases sequentially and collect all failures instead of failing fast."""
env = os.environ if environ is None else environ
results: list[LiteLLMLiveResult] = []
failures: list[tuple[LiteLLMLiveCase, str]] = []
for case in cases:
if case.api_key_env and not env.get(case.api_key_env):
failures.append((case, f"missing required env var {case.api_key_env}"))
continue
try:
results.append(
await evaluate_case(case, environ=env, provider_factory=provider_factory)
)
except Exception as exc: # pragma: no cover - exercised by live provider failures
failures.append((case, f"{type(exc).__name__}: {exc}"))
return results, failures
def build_results_table(
results: Sequence[LiteLLMLiveResult],
failures: Sequence[tuple[LiteLLMLiveCase, str]],
) -> Table:
"""Build a rich report table for live LiteLLM results."""
table = Table(title="LiteLLM Live Embedding Evaluation", show_lines=False)
table.add_column("Status", no_wrap=True)
table.add_column("Case", style="cyan", no_wrap=True)
table.add_column("Model", style="magenta")
table.add_column("Dims", justify="right")
table.add_column("Roles", no_wrap=True)
table.add_column("Forward dims", no_wrap=True)
table.add_column("Related", justify="right")
table.add_column("Distractor", justify="right")
table.add_column("Norm", justify="right")
table.add_column("Total ms", justify="right")
for result in results:
roles = _roles_label(result.document_input_type, result.query_input_type)
table.add_row(
"[green]PASS[/green]",
result.name,
result.model,
str(result.dimensions),
roles,
_bool_label(result.forward_dimensions),
f"{result.related_score:.4f}",
f"{result.distractor_score:.4f}",
f"{result.min_norm:.4f}-{result.max_norm:.4f}",
f"{result.total_latency_ms:.1f}",
)
for case, reason in failures:
roles = _roles_label(case.document_input_type, case.query_input_type)
table.add_row(
"[red]FAIL[/red]",
case.name,
case.model,
str(case.dimensions),
roles,
_bool_label(case.forward_dimensions),
"-",
"-",
"-",
reason,
)
return table
def _roles_label(document_input_type: str | None, query_input_type: str | None) -> str:
if document_input_type or query_input_type:
return f"{document_input_type or '-'} / {query_input_type or '-'}"
return "auto"
def _bool_label(value: bool | None) -> str:
if value is True:
return "true"
if value is False:
return "false"
return "auto"
def _load_cases_arg(args: argparse.Namespace) -> str | None:
if args.cases_json and args.cases_file:
raise ValueError("Use --cases-json or --cases-file, not both")
if args.cases_json:
return str(args.cases_json)
if args.cases_file:
return Path(args.cases_file).read_text(encoding="utf-8")
return None
async def _async_main(args: argparse.Namespace, environ: Mapping[str, str]) -> int:
console = Console()
# Trigger: this command performs real network calls and can spend API quota.
# Why: keeping the same explicit opt-in guard as pytest prevents accidental
# live calls when someone discovers the harness through `just --list`.
# Outcome: humans get a clear command to run before any provider is called.
if environ.get("BASIC_MEMORY_RUN_LITELLM_INTEGRATION") != "1":
console.print(
"[red]Set BASIC_MEMORY_RUN_LITELLM_INTEGRATION=1 to run live LiteLLM "
"provider checks.[/red]"
)
return 2
custom_cases_raw = _load_cases_arg(args)
cases = configured_cases(environ, custom_cases_raw=custom_cases_raw)
if not cases:
console.print(
"[yellow]No LiteLLM live cases configured.[/yellow]\n"
"Set OPENAI_API_KEY, COHERE_API_KEY, or provide custom cases with "
"BASIC_MEMORY_TEST_LITELLM_CASES / --cases-file."
)
return 2
results, failures = await evaluate_cases(cases, environ=environ)
if args.json:
payload = {
"results": [asdict(result) for result in results],
"failures": [{"case": asdict(case), "reason": reason} for case, reason in failures],
}
json.dump(payload, sys.stdout, indent=2)
sys.stdout.write("\n")
else:
console.print(build_results_table(results, failures))
return 1 if failures else 0
def main(argv: Sequence[str] | None = None) -> int:
"""CLI entrypoint for live LiteLLM evaluation."""
parser = argparse.ArgumentParser(description="Run opt-in live LiteLLM embedding checks")
parser.add_argument(
"--cases-json",
help="JSON array of custom LiteLLM cases; overrides BASIC_MEMORY_TEST_LITELLM_CASES",
)
parser.add_argument(
"--cases-file",
help="Path to a JSON file containing custom LiteLLM cases",
)
parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
args = parser.parse_args(argv)
try:
return asyncio.run(_async_main(args, os.environ))
except ValueError as exc:
Console(stderr=True).print(f"[red]{exc}[/red]")
return 2
if __name__ == "__main__":
raise SystemExit(main())