diff --git a/litellm/entrypoint.py b/litellm/entrypoint.py index e8d68842..6993dc89 100644 --- a/litellm/entrypoint.py +++ b/litellm/entrypoint.py @@ -69,6 +69,7 @@ class RobustDatetime(original_datetime): """A datetime subclass that handles flexible date format parsing in strptime.""" @classmethod def strptime(cls, date_str: str, fmt: str) -> original_datetime: + """Flexible strptime implementation that handles various ISO-like formats.""" if not isinstance(date_str, str): return original_datetime.strptime(date_str, fmt) diff --git a/router/main.py b/router/main.py index ca9072ba..35844949 100644 --- a/router/main.py +++ b/router/main.py @@ -1,3 +1,4 @@ +"""Main FastAPI application for the LLM Triage & Fallback Gateway.""" import os import aiofiles import re @@ -250,9 +251,11 @@ class ValkeyCooldownPersistence: """Persistence provider mapping Valkey/Redis client synchronization to the global handlers.""" async def sync(self) -> None: + """Synchronize cooldowns from Valkey to local memory.""" await sync_cooldowns_from_valkey() async def save(self) -> None: + """Persist local memory cooldowns to Valkey.""" await save_cooldowns_to_valkey() @@ -790,6 +793,7 @@ async def _register_ollama_models_in_db(master_key: str): ] def _load_yaml(p): + """Helper to load a YAML file safely.""" with open(p, "r", encoding="utf-8") as f: return yaml.safe_load(f) @@ -1015,6 +1019,14 @@ async def classify_request( When langfuse_trace_id is provided, the classifier HTTP call is wrapped in a child observation (span) so latency and output appear as a nested span in Langfuse traces. + + Args: + prompt: The user prompt to classify. + bypass_cache: If True, skip the in-memory TTL cache. + langfuse_trace_id: Optional trace ID to associate with the classification span. + + Returns: + A tuple containing (decision, latency_ms, cache_hit, raw_output). """ global triage_cache, stats @@ -1312,6 +1324,7 @@ def detect_active_tool(body: dict) -> str: @dataclass class ToolUsageRecord: + """Data class representing a single tool usage record for metrics tracking.""" tool_name: str prompt_tokens: int completion_tokens: int @@ -1802,7 +1815,17 @@ async def proxy_models(): @app.post("/v1/chat/completions") async def chat_completions(request: Request): - """Handle incoming OpenAI-compatible chat completions requests and route them dynamically based on triage logic.""" + """Handle incoming OpenAI-compatible chat completions requests. + + Routes requests dynamically based on triage logic, handling cascading fallbacks, + caching, and premium proxying (agy/ollama). + + Args: + request: The incoming FastAPI Request object. + + Returns: + A StreamingResponse or JSONResponse containing the model completion. + """ global stats start_time = time.time() @@ -2278,6 +2301,7 @@ async def agy_stream_generator(): logger.info(f"Ollama route: proxying to LiteLLM as model={target_model}") async def execute_proxy(model_name: str): + """Executes a proxy request to a backend model.""" # Resolve backend connection parameters backend_conf = backends.get(model_name) if not backend_conf: @@ -4071,6 +4095,7 @@ class AnnotationItem(BaseModel): @field_validator("tier") @classmethod def validate_tier(cls, v): + """Validate the tier field of an AnnotationItem.""" if v is None: return v if isinstance(v, int): @@ -4084,10 +4109,12 @@ def validate_tier(cls, v): return v class AnnotationPayload(RootModel): + """Pydantic model representing a payload of multiple annotations.""" root: Dict[str, AnnotationItem] @model_validator(mode="after") def validate_payload(self) -> "AnnotationPayload": + """Validate the entire annotation payload for size and key constraints.""" data = self.root if len(data) > 1000: raise ValueError("Payload size limit exceeded: maximum of 1000 annotations allowed per request.") @@ -4113,6 +4140,7 @@ def validate_payload(self) -> "AnnotationPayload": async def _read_annotations_async(path) -> dict: + """Read annotations from disk asynchronously with caching.""" import copy # Do not swallow OSError if file doesn't exist to preserve original behavior. diff --git a/router/memory_mcp.py b/router/memory_mcp.py index 02b15d11..abb367ec 100755 --- a/router/memory_mcp.py +++ b/router/memory_mcp.py @@ -242,6 +242,7 @@ async def handle_remove_memory_category(args: dict) -> str: deleted_count = 0 async def delete_item(client, entry, sem): + """Delete a single memory item from the backend.""" nonlocal deleted_count async with sem: key = entry["key"] @@ -316,11 +317,13 @@ async def handle_remove_specific_memory(args: dict) -> str: # --------------------------------------------------------------------------- def log(msg: str): + """Log a message to stderr for MCP diagnostics.""" sys.stderr.write(f"[memory-mcp] {msg}\n") sys.stderr.flush() async def handle_request(req: dict) -> dict | None: + """Handle an incoming JSON-RPC request from the MCP client.""" method = req.get("method") params = req.get("params", {}) @@ -478,6 +481,7 @@ async def handle_request(req: dict) -> dict | None: # --------------------------------------------------------------------------- async def main_loop(): + """Main execution loop for the MCP server, reading from stdin.""" log("LiteLLM Memory MCP Bridge v2 started (PostgreSQL-backed).") for line in sys.stdin: line = line.strip() diff --git a/scripts/benchmark_classifier.py b/scripts/benchmark_classifier.py index 69bc71e8..ed8ce2ed 100644 --- a/scripts/benchmark_classifier.py +++ b/scripts/benchmark_classifier.py @@ -61,6 +61,7 @@ def classify(prompt): confusion = defaultdict(Counter) # confusion[expected][predicted] def process_item(item): + """Process a single dataset item and return expected/predicted labels.""" try: if not isinstance(item, dict): raise TypeError("Item is not a dictionary") @@ -82,6 +83,7 @@ def process_item(item): next_start_time = [time.monotonic()] def process_item_with_rate_limit(index_and_item): + """Helper to process an item with a simplified rate limit delay.""" i, item = index_and_item sleep_delay = 0.0 with rate_lock: diff --git a/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py index 0737fa01..c27290ab 100644 --- a/scripts/benchmark_tokens.py +++ b/scripts/benchmark_tokens.py @@ -1,3 +1,4 @@ +"""Benchmark token estimation logic against ground truth examples.""" import sys import os from pathlib import Path diff --git a/scripts/get_pr_status.py b/scripts/get_pr_status.py index 223b3463..5c131d70 100644 --- a/scripts/get_pr_status.py +++ b/scripts/get_pr_status.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +"""Utility to query GitHub PR status and check for required approvals.""" import subprocess import json import sys @@ -54,6 +55,7 @@ def get_pr_status(pr_id: str = "") -> None: def main(): + """CLI entrypoint for PR status checking.""" pr_id = sys.argv[1] if len(sys.argv) > 1 else "" get_pr_status(pr_id) diff --git a/scripts/host_agy_daemon.py b/scripts/host_agy_daemon.py index 9233ea65..5948b74c 100755 --- a/scripts/host_agy_daemon.py +++ b/scripts/host_agy_daemon.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +"""HTTP daemon to bridge router requests to the host-side agy CLI.""" import asyncio import json import os @@ -11,6 +12,7 @@ CACHE_FILE = os.path.expanduser("~/.gemini/antigravity-cli/cache/last_conversations.json") def get_last_conversation_id(): + """Retrieve the last active conversation ID from the agy cache.""" try: if os.path.exists(CACHE_FILE): with open(CACHE_FILE, "r") as f: @@ -22,7 +24,9 @@ def get_last_conversation_id(): return None class AgyDaemonHandler(BaseHTTPRequestHandler): + """HTTP request handler for agy execution requests.""" def do_POST(self): + """Handle POST requests to execute agy commands.""" if self.path != "/run": self.send_response(404) self.end_headers() @@ -51,6 +55,7 @@ def do_POST(self): asyncio.set_event_loop(loop) async def run_stream(): + """Asynchronously execute agy and stream output via PTY.""" import pty env = os.environ.copy() @@ -84,6 +89,7 @@ async def run_stream(): loop_ref = asyncio.get_running_loop() def read_bytes(): + """Read raw bytes from the PTY master file descriptor.""" try: return os.read(master_fd, 1024) except OSError: @@ -136,6 +142,7 @@ def read_bytes(): asyncio.set_event_loop(loop) async def run(): + """Asynchronously execute agy and capture full output.""" env = os.environ.copy() if model_override: env["CASCADE_DEFAULT_MODEL_OVERRIDE"] = model_override @@ -216,10 +223,12 @@ async def run(): self.wfile.write(response_bytes) def log_message(self, format, *args): + """Override to silence standard HTTP logging.""" # Silence HTTP log outputs in standard output to keep service clean pass def run_server(): + """Start the ThreadingHTTPServer on the configured port.""" server = ThreadingHTTPServer(('127.0.0.1', PORT), AgyDaemonHandler) print(f"🚀 Host agy Daemon running on http://127.0.0.1:{PORT}") try: diff --git a/scripts/verification/verification_helpers.py b/scripts/verification/verification_helpers.py index a3a1c839..56fe9113 100644 --- a/scripts/verification/verification_helpers.py +++ b/scripts/verification/verification_helpers.py @@ -1,3 +1,4 @@ +"""Shared verification helpers for cooldown and routing integration tests.""" # Shared verification helpers for cooldown and routing tests try: from scripts.chat_helpers import parse_chat_response @@ -25,6 +26,7 @@ def load_litellm_key(workspace_dir: str) -> str: return litellm_key def get_triage_request_count(metrics_url: str = "http://localhost:5000/metrics") -> int: + """Parse Prometheus metrics to retrieve the total triage request count.""" try: response = httpx.get(metrics_url, timeout=5.0) response.raise_for_status() diff --git a/scripts/verification/verify_canonical_endpoints.py b/scripts/verification/verify_canonical_endpoints.py index c0c6ce62..23be879f 100644 --- a/scripts/verification/verify_canonical_endpoints.py +++ b/scripts/verification/verify_canonical_endpoints.py @@ -32,6 +32,7 @@ def load_env(dev: bool = False) -> dict: loaded_files = [] def _parse(path: Path): + """Helper to read and parse an .env file.""" if not path.exists(): print(f" ⚠ env file not found: {path}") return @@ -74,6 +75,7 @@ def _parse(path: Path): def check(label: str, ok: bool, detail: str = "") -> bool: + """Helper to log a verification check result.""" mark = "✓" if ok else "✗" extra = f" — {detail}" if detail else "" print(f" {mark} {label}{extra}") @@ -428,6 +430,7 @@ def test_canonical_urls(cfg: dict) -> tuple[int, int, int]: def main(): + """Main entrypoint for canonical endpoint verification.""" parser = argparse.ArgumentParser(description="Verify canonical endpoints") parser.add_argument( "--dev", action="store_true", help="Test dev environment (dev-router-pod)"