diff --git a/.env.dev b/.env.dev index 10e3f22b..9dc4b1d0 100644 --- a/.env.dev +++ b/.env.dev @@ -28,3 +28,4 @@ ROUTER_IMAGE="localhost/llm-routing-dev:latest" # Shared infrastructure (same llama classifier across dev & prod) LLAMA_CLASSIFIER_URL="https://x570.vendeuvre.lan/llm-routing/llama/v1" +LLAMA_SERVER_URL="https://x570.vendeuvre.lan/llm-routing/llama" diff --git a/router/config.yaml b/router/config.yaml index 4cb81e48..80cb812f 100644 --- a/router/config.yaml +++ b/router/config.yaml @@ -9,6 +9,8 @@ router: api_key: "os.environ/ROUTER_API_KEY" model: "qwen-4b-routing" +llama_server_url: "os.environ/LLAMA_SERVER_URL" + classification_rules: system_prompt: | Classify the coding task complexity. Output ONLY the tier name. diff --git a/router/main.py b/router/main.py index 35844949..c4dbe1a5 100644 --- a/router/main.py +++ b/router/main.py @@ -1,5 +1,6 @@ """Main FastAPI application for the LLM Triage & Fallback Gateway.""" import os +import uuid import aiofiles import re import sys @@ -12,7 +13,7 @@ import yaml import httpx import redis.asyncio as aioredis -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, nullcontext from dataclasses import dataclass from fastapi import FastAPI, Request, HTTPException, Response @@ -27,8 +28,12 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator, RootModel from typing import Dict, Optional, Union +try: + from langfuse import propagate_attributes # noqa: F401 +except ImportError: + propagate_attributes = None + LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or f"http://127.0.0.1:{os.getenv('LITELLM_PORT') or '4000'}").rstrip("/") -LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip("/") LANGFUSE_HOST = (os.getenv("LANGFUSE_HOST") or f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}").rstrip("/") GEMINI_OAUTH_CREDS_PATH = "/config/gemini_auth/oauth_creds.json" @@ -109,6 +114,25 @@ def get_http_client(): _classifier_client: httpx.AsyncClient | None = None +def _resolve_verify(env_var: str) -> bool | str: + """Resolve TLS verify setting from an environment variable. + + Returns: + False if unset, empty, or boolean-like false. + True if boolean-like true. + A string path for a CA bundle file. + """ + ca_bundle = os.getenv(env_var) + if ca_bundle is None: + return False + v = ca_bundle.strip() + if v.lower() in ("false", "0", "off", "no", "none", "null", "disabled", ""): + return False + if v.lower() in ("true", "1", "on", "yes"): + return True + return v + + def get_classifier_client(): """Return a singleton httpx client for classifier calls (internal self-signed TLS). @@ -118,23 +142,34 @@ def get_classifier_client(): """ global _classifier_client if _classifier_client is None: - ca_bundle = os.getenv("CLASSIFIER_CA_BUNDLE") - if ca_bundle is not None: - ca_bundle_stripped = ca_bundle.strip() - if ca_bundle_stripped.lower() in ("false", "0", "off", "no", "none", "null", "disabled", ""): - verify = False - elif ca_bundle_stripped.lower() in ("true", "1", "on", "yes"): - verify = True - else: - verify = ca_bundle_stripped - else: - verify = False _classifier_client = httpx.AsyncClient( - limits=_http_limits(), timeout=3600.0, verify=verify + limits=_http_limits(), + timeout=3600.0, + verify=_resolve_verify("CLASSIFIER_CA_BUNDLE"), ) return _classifier_client +_llama_client: httpx.AsyncClient | None = None + + +def get_llama_client(): + """Return a singleton httpx client for llama.cpp server calls (internal self-signed TLS). + + By default verify is disabled because the llama-server sits behind HAProxy with + a self-signed certificate on the internal network. Set LLAMA_CA_BUNDLE + to a PEM file path to enable TLS verification (e.g. for CI or staging). + """ + global _llama_client + if _llama_client is None: + _llama_client = httpx.AsyncClient( + limits=_http_limits(), + timeout=3600.0, + verify=_resolve_verify("LLAMA_CA_BUNDLE"), + ) + return _llama_client + + # Compiled regular expressions for token estimation heuristics WORD_RE = re.compile(r'[a-zA-Z0-9]+') NON_ASCII_RE = re.compile(r'[^\s\x00-\x7F]') @@ -293,6 +328,81 @@ def get_langfuse(): return _langfuse_client if _langfuse_client is not False else None +def _end_parent_obs(parent_obs, output=None, metadata=None) -> None: + """Safely finalize a Langfuse parent observation (SDK v4: update + end). + + Non-fatal — swallows all exceptions. + """ + if parent_obs is None: + return None + try: + update_kwargs = {} + if output is not None: + update_kwargs["output"] = output + if metadata is not None: + update_kwargs["metadata"] = metadata + if update_kwargs: + parent_obs.update(**update_kwargs) + parent_obs.end() + except Exception: + logger.debug("_end_parent_obs failed (non-fatal)", exc_info=True) + pass + return None + + +def _end_child_span(span, output=None, metadata=None) -> None: + """Safely finalize a Langfuse child span (SDK v4: update + end). + + Non-fatal — errors are never propagated. + """ + if span is None: + return + try: + update_kwargs = {} + if output is not None: + update_kwargs["output"] = output + if metadata is not None: + update_kwargs["metadata"] = metadata + if update_kwargs: + span.update(**update_kwargs) + span.end() + except Exception: + logger.debug("_end_child_span failed (non-fatal)", exc_info=True) + pass + + +def _close_prop_ctx(prop_ctx): + """Safely exit a propagate_attributes context manager if active. + + Non-fatal — swallows all exceptions. + Returns None after exit for idempotent cleanup. + """ + if prop_ctx is not None: + try: + prop_ctx.__exit__(None, None, None) + except Exception: + logger.debug("_close_prop_ctx failed (non-fatal)", exc_info=True) + pass + return None + + +def _make_prop_ctx(session_id, user_id): + """Create a propagate_attributes context manager if session/user propagation is active. + + Returns a context manager (entered by the caller) or None if + propagate_attributes is unavailable or no session/user data is provided. + DRY-consolidates the 4 duplicate condition+builder blocks across streaming + generators and the non-streaming init path. + """ + if not propagate_attributes or not (session_id or user_id): + return None + return propagate_attributes( + session_id=session_id or None, + user_id=user_id or None, + tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], + ) + + async def push_aggregate_scores(): """Push aggregate KPIs as Langfuse scores every 5 minutes.""" while True: @@ -415,6 +525,21 @@ async def push_aggregate_scores(): system_prompt = config.get("classification_rules", {}).get("system_prompt", "") backends = {b["name"]: b for b in config.get("backends", [])} +# --- Resolve llama_server_url from config (os.environ/ pattern) --- +_raw_llama_url = config.get("llama_server_url") or "http://127.0.0.1:8080" +if isinstance(_raw_llama_url, str) and _raw_llama_url.startswith("os.environ/"): + env_var = _raw_llama_url.split("/", 1)[1] + _raw_llama_url = os.environ.get(env_var, "") + if not _raw_llama_url: + if "pytest" in sys.modules: + _raw_llama_url = "http://127.0.0.1:8080" + else: + logger.warning( + "LLAMA_SERVER_URL env var not set, falling back to http://127.0.0.1:8080" + ) + _raw_llama_url = "http://127.0.0.1:8080" +LLAMA_SERVER_URL = str(_raw_llama_url).rstrip("/") + # Default colors for tool visualization badges and charts TOOL_COLORS = { "tree": "#34d399", # Green @@ -971,6 +1096,12 @@ async def lifespan(app: FastAPI): await _classifier_client.aclose() _classifier_client = None + # Close llama client + global _llama_client + if _llama_client is not None: + await _llama_client.aclose() + _llama_client = None + # Close Redis client global _redis_client if _redis_client is not None and _redis_client is not False: @@ -1012,6 +1143,16 @@ async def check_http_endpoint(url: str) -> bool: return False +async def _check_llama_health() -> bool: + """Check llama-server health using the llama client (verify=False for self-signed TLS).""" + try: + client = get_llama_client() + r = await client.get(f"{LLAMA_SERVER_URL}/health", timeout=3.0) + return r.status_code < 500 + except Exception: + return False + + async def classify_request( prompt: str, bypass_cache: bool = False, langfuse_trace_id: str | None = None ) -> tuple[str, float, bool, str]: @@ -1104,17 +1245,13 @@ async def classify_request( latency = (time.time() - start_time) * 1000.0 if response.status_code != 200: - if class_span_obj: - try: - class_span_obj.end( - output={ - "status": response.status_code, - "error": "classification_failed", - }, - metadata={"latency_ms": latency}, - ) - except Exception: - pass + _end_child_span(class_span_obj, + output={ + "status": response.status_code, + "error": "classification_failed", + }, + metadata={"latency_ms": latency}, + ) logger.error( f"Classification failed with status {response.status_code}: {response.text}" ) @@ -1141,14 +1278,10 @@ async def classify_request( decision = "agent-advanced-core" # Finalize classifier child span - if class_span_obj: - try: - class_span_obj.end( - output={"tier": decision, "raw": raw_result}, - metadata={"latency_ms": latency}, - ) - except Exception: - pass + _end_child_span(class_span_obj, + output={"tier": decision, "raw": raw_result}, + metadata={"latency_ms": latency}, + ) # Store in cache triage_cache[normalized_prompt] = (decision, time.time()) @@ -1448,7 +1581,7 @@ async def get_llamacpp_metrics() -> dict: """Fetches live model inventory and slot statistics from the local llama-server.""" result = {"models": [], "slots": [], "build": "unknown"} try: - client = get_http_client() + client = get_llama_client() # Fetch model list r = await client.get(f"{LLAMA_SERVER_URL}/v1/models", timeout=3.0) if r.status_code == 200: @@ -1878,199 +2011,261 @@ async def chat_completions(request: Request): client_model = body.get("model", "llm-routing-auto-free") + # Extract session_id and user_id for Langfuse tracing + _trace_session_id = ( + body.get("session_id") + or body.get("session") + or request.headers.get("x-session-id") + ) + if _trace_session_id: + _trace_session_id = str(_trace_session_id) + _trace_user_id = ( + body.get("user") + or request.headers.get("x-user-id") + ) + if _trace_user_id: + _trace_user_id = str(_trace_user_id) + # --- Langfuse parent trace: create early so child spans can reference it --- langfuse_trace_id = None parent_obs = None + _prop_ctx = None + _is_streaming = body.get("stream", False) lf = get_langfuse() if lf: try: langfuse_trace_id = lf.create_trace_id( - seed=f"triage_{stats['total_requests']}" + seed=str(uuid.uuid4()) ) + # Propagate session_id/user_id via Langfuse's native session mechanism. + # For non-streaming: enter here (same asyncio task, contextvars work). + # For streaming: each generator creates its own context in its own task + # because OpenTelemetry contextvars are task-isolated. + if not _is_streaming: + _prop_ctx = _make_prop_ctx(_trace_session_id, _trace_user_id) + if _prop_ctx is not None: + _prop_ctx.__enter__() parent_obs = lf.start_observation( trace_context={"trace_id": langfuse_trace_id}, name=f"triage-{client_model}", input=last_user_message[:200], level="DEFAULT", + metadata={ + "client_model": client_model, + "environment": os.getenv("ENVIRONMENT", "production"), + }, ) except Exception as e: logger.warning(f"Langfuse trace init failed (non-fatal): {e}") langfuse_trace_id = None parent_obs = None + if _prop_ctx: + _prop_ctx = _close_prop_ctx(_prop_ctx) - if client_model in AUTO_MODELS or client_model == "llm-routing-ollama": - # Full pipeline: classify → route to best tier - bypass_cache = request.headers.get("x-bypass-cache") == "true" - ( - target_model, - triage_latency, - was_cache_hit, - raw_classification, - ) = await classify_request( - last_user_message, - bypass_cache=bypass_cache, - langfuse_trace_id=langfuse_trace_id, - ) - logger.info(f"Triage decision (auto/gated): Routing to -> '{target_model}'") - elif client_model in DIRECT_TIERS: - # Direct routing: client knows what tier they want, skip classifier - target_model = client_model - triage_latency = 0.0 - was_cache_hit = False - raw_classification = f"direct ({client_model})" - logger.info( - f"Direct routing: Client requested '{client_model}', skipping classifier" - ) - else: - raise HTTPException( - status_code=400, - detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " - f"or one of: {', '.join(sorted(DIRECT_TIERS))}", - ) + try: + _non_streaming_finalized = False + if client_model in AUTO_MODELS or client_model == "llm-routing-ollama": + # Full pipeline: classify → route to best tier + bypass_cache = request.headers.get("x-bypass-cache") == "true" + ( + target_model, + triage_latency, + was_cache_hit, + raw_classification, + ) = await classify_request( + last_user_message, + bypass_cache=bypass_cache, + langfuse_trace_id=langfuse_trace_id, + ) + logger.info(f"Triage decision (auto/gated): Routing to -> '{target_model}'") + elif client_model in DIRECT_TIERS: + # Direct routing: client knows what tier they want, skip classifier + target_model = client_model + triage_latency = 0.0 + was_cache_hit = False + raw_classification = f"direct ({client_model})" + logger.info( + f"Direct routing: Client requested '{client_model}', skipping classifier" + ) + else: + # guard: end parent obs before raising + _end_parent_obs(parent_obs, + output={"error": f"Unknown model: {client_model}"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise HTTPException( + status_code=400, + detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " + f"or one of: {', '.join(sorted(DIRECT_TIERS))}", + ) - # Update in-memory statistics - stats["total_requests"] += 1 - stats["last_triage_decision"] = target_model - stats["total_triage_time_ms"] += triage_latency - stats["avg_triage_latency_ms"] = ( - stats["total_triage_time_ms"] / stats["total_requests"] - ) + # Update in-memory statistics + stats["total_requests"] += 1 + stats["last_triage_decision"] = target_model + stats["total_triage_time_ms"] += triage_latency + stats["avg_triage_latency_ms"] = ( + stats["total_triage_time_ms"] / stats["total_requests"] + ) - if target_model == "agent-simple-core": - stats["simple_requests"] = stats.get("simple_requests", 0) + 1 - elif target_model == "agent-medium-core": - stats["medium_requests"] = stats.get("medium_requests", 0) + 1 - elif target_model == "agent-complex-core": - stats["complex_requests"] = stats.get("complex_requests", 0) + 1 - elif target_model == "agent-reasoning-core": - stats["reasoning_requests"] = stats.get("reasoning_requests", 0) + 1 - elif target_model == "agent-advanced-core": - stats["advanced_requests"] = stats.get("advanced_requests", 0) + 1 - await save_persisted_stats() - - # Update the parent Langfuse observation with classification results - if parent_obs: - try: - parent_obs.update( - output={"tier": target_model, "raw": raw_classification}, - metadata={ - "triage_latency_ms": round(triage_latency, 2), - "cache_hit": was_cache_hit, - "total_requests": stats["total_requests"], - }, + if target_model == "agent-simple-core": + stats["simple_requests"] = stats.get("simple_requests", 0) + 1 + elif target_model == "agent-medium-core": + stats["medium_requests"] = stats.get("medium_requests", 0) + 1 + elif target_model == "agent-complex-core": + stats["complex_requests"] = stats.get("complex_requests", 0) + 1 + elif target_model == "agent-reasoning-core": + stats["reasoning_requests"] = stats.get("reasoning_requests", 0) + 1 + elif target_model == "agent-advanced-core": + stats["advanced_requests"] = stats.get("advanced_requests", 0) + 1 + await save_persisted_stats() + + # Update the parent Langfuse observation with classification results + if parent_obs: + try: + parent_obs.update( + output={"tier": target_model, "raw": raw_classification}, + metadata={ + "triage_latency_ms": round(triage_latency, 2), + "cache_hit": was_cache_hit, + "total_requests": stats["total_requests"], + }, + ) + except Exception as e: + logger.warning(f"Langfuse trace update failed (non-fatal): {e}") + + # --- PREMIUM PROXY ROUTES --- + # agy: triggered unconditionally for llm-routing-agy (direct). + # For AUTO models: only triggered when classifier picks agent-advanced-core + # or agent-reasoning-core. + # Reasoning tier → gemini-3.5-flash (single tier, low thinking) + # Advanced tier → gemini-3.5-flash → claude-opus-4.6 (full 2-tier chain) + # Proxied to host agy daemon on port 5005. + # ollama: triggered unconditionally for llm-routing-ollama (direct). + # For AUTO models: only triggered when classifier picks agent-advanced-core + # or agent-reasoning-core. + # Reasoning tier → deepseek-v4-flash (lighter, faster) + # Advanced tier → deepseek-v4-pro (full power) + # Proxied to LiteLLM as ollama-deepseek-v4-* — LiteLLM handles the + # native Ollama API call via its built-in ollama_chat provider. + # Classification gating (2026-06-16): auto models skip premium proxies entirely + # unless classified as advanced or reasoning, avoiding 4-minute agy timeouts on + # simple/medium/complex prompts that the fast OpenRouter free tier handles better. + + should_try_agy = ( + client_model == "llm-routing-agy" # direct — always try + or ( + client_model in ("llm-routing-auto-agy", "llm-routing-auto-agy-ollama") + and target_model in ("agent-advanced-core", "agent-reasoning-core") ) - except Exception as e: - logger.warning(f"Langfuse trace update failed (non-fatal): {e}") - - # --- PREMIUM PROXY ROUTES --- - # agy: triggered unconditionally for llm-routing-agy (direct). - # For AUTO models: only triggered when classifier picks agent-advanced-core - # or agent-reasoning-core. - # Reasoning tier → gemini-3.5-flash (single tier, low thinking) - # Advanced tier → gemini-3.5-flash → claude-opus-4.6 (full 2-tier chain) - # Proxied to host agy daemon on port 5005. - # ollama: triggered unconditionally for llm-routing-ollama (direct). - # For AUTO models: only triggered when classifier picks agent-advanced-core - # or agent-reasoning-core. - # Reasoning tier → deepseek-v4-flash (lighter, faster) - # Advanced tier → deepseek-v4-pro (full power) - # Proxied to LiteLLM as ollama-deepseek-v4-* — LiteLLM handles the - # native Ollama API call via its built-in ollama_chat provider. - # Classification gating (2026-06-16): auto models skip premium proxies entirely - # unless classified as advanced or reasoning, avoiding 4-minute agy timeouts on - # simple/medium/complex prompts that the fast OpenRouter free tier handles better. - - should_try_agy = ( - client_model == "llm-routing-agy" # direct — always try - or ( - client_model in ("llm-routing-auto-agy", "llm-routing-auto-agy-ollama") - and target_model in ("agent-advanced-core", "agent-reasoning-core") - ) - ) - should_try_ollama = ( - client_model - == "llm-routing-ollama" # always try (will map to flash for complex/below) - or ( - client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama") - and target_model - in ("agent-advanced-core", "agent-reasoning-core", "agent-complex-core") ) - ) - - # --- AGY PROXY --- - if should_try_agy: - agy_span_obj = None - try: - from agy_proxy import try_agy_proxy - - last_prompt = "" - for msg in reversed(messages): - if not isinstance(msg, dict): - continue - if msg.get("role") == "user": - content = msg.get("content") or "" - if isinstance(content, list): - content = "".join( - block.get("text") or "" - for block in content - if isinstance(block, dict) and block.get("type") == "text" - ) - last_prompt = str(content) - break - - session_id = ( - body.get("session_id") - or body.get("session") - or request.headers.get("x-session-id") + should_try_ollama = ( + client_model + == "llm-routing-ollama" # always try (will map to flash for complex/below) + or ( + client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama") + and target_model + in ("agent-advanced-core", "agent-reasoning-core", "agent-complex-core") ) - if session_id: - session_id = str(session_id) + ) - if last_prompt: - # --- Langfuse child span: agy proxy --- - if langfuse_trace_id: - lf_agy = get_langfuse() - if lf_agy: - try: - agy_span_obj = lf_agy.start_observation( - trace_context={"trace_id": langfuse_trace_id}, - name="agy-proxy", - input=last_prompt[:200], - metadata={"tier": target_model}, - level="DEFAULT", + # --- AGY PROXY --- + if should_try_agy: + agy_span_obj = None + try: + from agy_proxy import try_agy_proxy + + last_prompt = "" + for msg in reversed(messages): + if not isinstance(msg, dict): + continue + if msg.get("role") == "user": + content = msg.get("content") or "" + if isinstance(content, list): + content = "".join( + block.get("text") or "" + for block in content + if isinstance(block, dict) and block.get("type") == "text" ) - except Exception: - pass - - is_stream_requested = body.get("stream", False) - agy_response = await try_agy_proxy( - prompt=last_prompt, - messages=messages, - session_id=session_id, - total_timeout=300.0, - stream=is_stream_requested, - target_tier=target_model, - client=get_http_client(), - cooldown_persistence=ValkeyCooldownPersistence(), + last_prompt = str(content) + break + + session_id = ( + body.get("session_id") + or body.get("session") + or request.headers.get("x-session-id") ) - if agy_response: - model_name = agy_response.get("model", "gemini-3.5-flash (via agy)") - - if "stream" in agy_response: - # Real native stream generator - async def native_agy_stream_generator(stream_gen, model_name): - """Asynchronous generator yielding native OpenAI-compatible streaming chunks from the real agy daemon.""" - import uuid - - created_time = int(time.time()) - chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" - token_count = 0 + if session_id: + session_id = str(session_id) + + if last_prompt: + # --- Langfuse child span: agy proxy --- + if langfuse_trace_id: + lf_agy = get_langfuse() + if lf_agy: try: - async for token in stream_gen: - if not token: - continue - token_count += 1 - chunk_data = { + agy_span_obj = lf_agy.start_observation( + trace_context={"trace_id": langfuse_trace_id}, + name="agy-proxy", + input=last_prompt[:200], + metadata={"tier": target_model}, + level="DEFAULT", + ) + except Exception: + pass + + is_stream_requested = body.get("stream", False) + agy_response = await try_agy_proxy( + prompt=last_prompt, + messages=messages, + session_id=session_id, + total_timeout=300.0, + stream=is_stream_requested, + target_tier=target_model, + client=get_http_client(), + cooldown_persistence=ValkeyCooldownPersistence(), + ) + if agy_response: + model_name = agy_response.get("model", "gemini-3.5-flash (via agy)") + + if "stream" in agy_response: + # Real native stream generator + async def native_agy_stream_generator(stream_gen, model_name): + """Asynchronous generator yielding native OpenAI-compatible streaming chunks from the real agy daemon.""" + import uuid + + created_time = int(time.time()) + chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + token_count = 0 + finalized = False + _native_agy_prop = ( + _make_prop_ctx(_trace_session_id, _trace_user_id) + or nullcontext() + ) + _native_agy_prop.__enter__() + try: + async for token in stream_gen: + if not token: + continue + token_count += 1 + chunk_data = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model_name, + "choices": [ + { + "index": 0, + "delta": {"content": token}, + "finish_reason": None, + } + ], + } + yield f"data: {json.dumps(chunk_data)}\n\n".encode( + "utf-8" + ) + + # End of stream chunk + finish_data = { "id": chunk_id, "object": "chat.completion.chunk", "created": created_time, @@ -2078,565 +2273,706 @@ async def native_agy_stream_generator(stream_gen, model_name): "choices": [ { "index": 0, - "delta": {"content": token}, - "finish_reason": None, + "delta": {}, + "finish_reason": "stop", } ], } - yield f"data: {json.dumps(chunk_data)}\n\n".encode( + yield f"data: {json.dumps(finish_data)}\n\n".encode( "utf-8" ) + yield b"data: [DONE]\n\n" + + # Success telemetry + latency_ms = (time.time() - start_time) * 1000.0 + stats["total_proxy_time_ms"] += latency_ms + stats["avg_proxy_latency_ms"] = ( + stats["total_proxy_time_ms"] / stats["total_requests"] + ) + approx_prompt_tokens = estimate_prompt_tokens(body) + + record_tool_usage(ToolUsageRecord( + active_tool, + approx_prompt_tokens, + token_count, + model_name, + latency_ms, + route="google_oauth_direct", + )) + logger.info( + f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms" + ) + _end_child_span(agy_span_obj, + output={ + "model": model_name, + "tokens": token_count, + }, + metadata={ + "latency_ms": latency_ms, + "tier": target_model, + }, + ) + # Finalize parent trace for native agy stream + _end_parent_obs(parent_obs, + output={"model": model_name, "stream": True, + "tier": target_model, "route": "google_oauth_direct"}, + metadata={"latency_ms": latency_ms, + "completion_tokens": token_count}) + _close_prop_ctx(_native_agy_prop) + finalized = True + except Exception as stream_err: + logger.error( + f"Error during native agy stream generation: {type(stream_err).__name__}" + ) + _end_child_span(agy_span_obj, + output={"error": type(stream_err).__name__}, + metadata={"status": "failed"}, + ) + # End parent trace on stream error + _end_parent_obs(parent_obs, + output={"error": type(stream_err).__name__, + "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_native_agy_prop) + finalized = True + raise + finally: + if not finalized: + _end_child_span(agy_span_obj, + output={"error": "cancelled"}, + metadata={"status": "cancelled"}, + ) + _end_parent_obs(parent_obs, + output={"error": "cancelled", + "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_native_agy_prop) + + return StreamingResponse( + native_agy_stream_generator( + agy_response["stream"], model_name + ), + media_type="text/event-stream", + ) + else: + latency_ms = (time.time() - start_time) * 1000.0 + stats["total_proxy_time_ms"] += latency_ms + stats["avg_proxy_latency_ms"] = ( + stats["total_proxy_time_ms"] / stats["total_requests"] + ) + usage = agy_response.get("usage") or {} + prompt_tokens = usage.get("prompt_tokens") or 0 + completion_tokens = usage.get("completion_tokens") or 0 + record_tool_usage(ToolUsageRecord( + active_tool, + prompt_tokens, + completion_tokens, + model_name, + latency_ms, + route="google_oauth_direct", + )) + logger.info( + f"✅ agy proxy succeeded: {model_name}, {latency_ms:.0f}ms" + ) - # End of stream chunk - finish_data = { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created_time, + # Finalize agy span + _end_child_span(agy_span_obj, + output={ "model": model_name, - "choices": [ - { - "index": 0, - "delta": {}, - "finish_reason": "stop", + "tokens": completion_tokens, + }, + metadata={ + "latency_ms": latency_ms, + "tier": target_model, + }, + ) + + if is_stream_requested: + # Robust fallback: simulate stream if we requested stream but got buffered response + content = (agy_response.get("choices") or [{}])[0].get( + "message", {} + ).get("content") or "" + + async def agy_stream_generator(): + """Asynchronous generator yielding simulated OpenAI-compatible streaming chunks from a static agy response.""" + import uuid + + created_time = int(time.time()) + chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + chunk_size = 40 + finalized = False + _agy_gen_prop = ( + _make_prop_ctx(_trace_session_id, _trace_user_id) + or nullcontext() + ) + _agy_gen_prop.__enter__() + try: + for i in range(0, len(content), chunk_size): + chunk_text = content[i : i + chunk_size] + chunk_data = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model_name, + "choices": [ + { + "index": 0, + "delta": {"content": chunk_text}, + "finish_reason": None, + } + ], + } + yield f"data: {json.dumps(chunk_data)}\n\n".encode( + "utf-8" + ) + await asyncio.sleep(0.005) + + finish_data = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model_name, + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], } - ], - } - yield f"data: {json.dumps(finish_data)}\n\n".encode( - "utf-8" + yield f"data: {json.dumps(finish_data)}\n\n".encode( + "utf-8" + ) + yield b"data: [DONE]\n\n" + # Finalize parent trace for simulated agy stream + _end_parent_obs(parent_obs, + output={"model": model_name, "stream": True, + "tier": target_model, "route": "google_oauth_direct"}, + metadata={"latency_ms": latency_ms, + "completion_tokens": len(content) // 4}) + _close_prop_ctx(_agy_gen_prop) + finalized = True + except Exception as e: + logger.error( + f"Error during agy stream generation: {type(e).__name__}" + ) + _end_parent_obs(parent_obs, + output={"error": type(e).__name__, + "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_agy_gen_prop) + finalized = True + raise + finally: + if not finalized: + _end_parent_obs(parent_obs, + output={"error": "cancelled", + "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_agy_gen_prop) + + return StreamingResponse( + agy_stream_generator(), media_type="text/event-stream" ) - yield b"data: [DONE]\n\n" + else: + # Finalize parent trace for non-streaming agy + _end_parent_obs(parent_obs, + output={"model": model_name, "tier": target_model, + "route": "google_oauth_direct"}, + metadata={"latency_ms": latency_ms, + "completion_tokens": completion_tokens}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + return agy_response + # agy_response was falsy (None) — finalize agy span before falling back + _end_child_span(agy_span_obj, + output={"error": "no_response"}, + metadata={"status": "failed"}, + ) + logger.warning("agy proxy returned no response, falling back to LiteLLM") + except ImportError: + _end_child_span(agy_span_obj, + output={"error": "module_not_available"}, + metadata={"status": "skipped"}, + ) + logger.warning("agy_proxy module not available, falling back to LiteLLM") + except Exception as e: + _end_child_span(agy_span_obj, + output={"error": type(e).__name__}, + metadata={"status": "failed"}, + ) + logger.error(f"agy proxy failed: {type(e).__name__}, falling back to LiteLLM") + + if target_model == "llm-routing-agy": + target_model = "agent-advanced-core" + original_target_model = target_model + + # --- OLLAMA (via LiteLLM) --- + # LiteLLM's ollama_chat provider handles the native Ollama API call. + # We just proxy to LiteLLM with the appropriate model name. + # LiteLLM's fallback chain handles failures. + if should_try_ollama: + if client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama"): + if target_model in ("agent-advanced-core", "agent-reasoning-core"): + target_model = "ollama-deepseek-v4-pro" + elif target_model == "agent-complex-core": + target_model = "ollama-deepseek-v4-flash" + elif client_model == "llm-routing-ollama": + if target_model in ("agent-advanced-core", "agent-reasoning-core"): + target_model = "ollama-deepseek-v4-pro" + else: + target_model = "ollama-deepseek-v4-flash" + else: + # Fallback (e.g. if LiteLLM fallback loops back with model: llm-routing-ollama) + if target_model in ("agent-advanced-core", "agent-reasoning-core"): + target_model = "ollama-deepseek-v4-pro" + else: + target_model = "ollama-deepseek-v4-flash" + 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.""" + nonlocal _non_streaming_finalized + # Resolve backend connection parameters + backend_conf = backends.get(model_name) + if not backend_conf: + logger.error(f"Backend '{model_name}' not found in configuration backends.") + raise HTTPException( + status_code=500, detail=f"Backend {model_name} misconfigured" + ) + + backend_api_base = backend_conf["api_base"] + backend_api_key = backend_conf["api_key"] + if backend_api_key == "DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER": + backend_api_key = os.getenv("LITELLM_MASTER_KEY", backend_api_key) - # Success telemetry - latency_ms = (time.time() - start_time) * 1000.0 - stats["total_proxy_time_ms"] += latency_ms + logger.info(f"Proxying to LiteLLM as model={model_name}") + + # --- Langfuse child span: LiteLLM proxy --- + litellm_span_obj = None + if langfuse_trace_id: + lf_litellm = get_langfuse() + if lf_litellm: + try: + litellm_span_obj = lf_litellm.start_observation( + trace_context={"trace_id": langfuse_trace_id}, + name="litellm-proxy", + input=model_name, + metadata={"model": model_name}, + level="DEFAULT", + ) + except Exception: + pass + + client = get_http_client() + try: + headers = {"Authorization": f"Bearer {backend_api_key}"} + if langfuse_trace_id: + headers["X-Langfuse-Trace-Id"] = langfuse_trace_id + + # Handle streaming vs non-streaming proxying (LiteLLM handles fallback internally) + proxy_start = time.time() + + # --- Pre-screening: clamp max_tokens to fit within downstream model context limits --- + try: + body_to_send = body.copy() + body_to_send["model"] = model_name + requested_max_tokens = body_to_send.get("max_tokens", 4096) + + # Tier-aware minimum context length (from actual roster data): + # - agent-simple-core: 32K (includes tiny liquid/dolphin models) + # - agent-medium-core+: 256K (smallest non-tiny model is nemotron-nano-omni at 256K) + # - ollama-deepseek-v4-*: 1M (DeepSeek V4 native context) + _tier_min_ctx = { + "agent-simple-core": 32768, + "ollama-deepseek-v4-pro": 524288, + "ollama-deepseek-v4-flash": 524288, + } + _min_ctx = _tier_min_ctx.get(model_name, 262144) + _est_input = estimate_prompt_tokens(body_to_send) + _safe_max = _min_ctx - _est_input - 2048 # 2K safety margin + if _safe_max < 1024: + raise HTTPException( + status_code=400, + detail=f"Context window exceeded. Estimated input tokens ({_est_input}) plus safety margin (2048) exceeds model context limit ({_min_ctx}).", + ) + if requested_max_tokens > _safe_max: + logger.warning( + f"⛔ Clamping max_tokens: {requested_max_tokens} → {_safe_max} " + f"(est_input={_est_input}, min_ctx={_min_ctx}, tier={model_name})" + ) + body_to_send["max_tokens"] = _safe_max + except HTTPException: + _end_child_span(litellm_span_obj, + output={"error": "Context window exceeded"}, + metadata={"status": "failed"}, + ) + raise + except Exception as e: + logger.warning(f"Pre-screening failed (non-fatal): {e}") + body_to_send = body.copy() + body_to_send["model"] = model_name + if "metadata" not in body_to_send or not isinstance( + body_to_send["metadata"], dict + ): + body_to_send["metadata"] = {} + body_to_send["metadata"]["trace_name"] = "agent-completion" + + if body.get("stream", False): + logger.info(f"Proxying streaming to LiteLLM as model={model_name}") + req = client.build_request( + "POST", + f"{backend_api_base}/chat/completions", + json=body_to_send, + headers=headers, + ) + r = await client.send(req, stream=True) + if r.status_code == 200: + + async def stream_generator(): + """Asynchronous generator that yields streaming chunks from LiteLLM completions response and logs usage stats on completion.""" + import codecs + + completion_chars = 0 + request_tokens = estimate_prompt_tokens(body_to_send) + sse_buffer = "" + decoder = codecs.getincrementaldecoder("utf-8")() + finalized = False + _litellm_gen_prop = ( + _make_prop_ctx(_trace_session_id, _trace_user_id) + or nullcontext() + ) + _litellm_gen_prop.__enter__() + try: + async for chunk in r.aiter_bytes(): + yield chunk + try: + sse_buffer += decoder.decode(chunk) + while "\n" in sse_buffer: + line, sse_buffer = sse_buffer.split("\n", 1) + line = line.strip() + if line.startswith("data:"): + data_str = line[5:].strip() + if data_str == "[DONE]": + continue + try: + data_json = json.loads(data_str) + choices = data_json.get("choices", []) + if choices and isinstance( + choices[0], dict + ): + delta = choices[0].get("delta") + if isinstance(delta, dict): + content = ( + delta.get("content") or "" + ) + completion_chars += len(content) + except Exception: + pass + except Exception: + pass + proxy_latency = (time.time() - proxy_start) * 1000.0 + stats["total_proxy_time_ms"] += proxy_latency stats["avg_proxy_latency_ms"] = ( stats["total_proxy_time_ms"] / stats["total_requests"] ) - approx_prompt_tokens = estimate_prompt_tokens(body) - record_tool_usage(ToolUsageRecord( active_tool, - approx_prompt_tokens, - token_count, + request_tokens, + completion_chars // 4, model_name, - latency_ms, - route="google_oauth_direct", + proxy_latency, + route="litellm_fallback", )) - logger.info( - f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms" + # Finalize LiteLLM span (streaming path) + _end_child_span(litellm_span_obj, + output={"model": model_name, "stream": True}, + metadata={ + "latency_ms": proxy_latency, + "tokens": completion_chars // 4, + }, ) - if agy_span_obj: - try: - agy_span_obj.end( - output={ - "model": model_name, - "tokens": token_count, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass - except Exception as stream_err: - logger.error( - f"Error during native agy stream generation: {type(stream_err).__name__}" + # Finalize parent trace (streaming path) + _end_parent_obs(parent_obs, + output={"model": model_name, "stream": True, + "tier": target_model, "route": "litellm_fallback"}, + metadata={"latency_ms": proxy_latency, + "completion_tokens": completion_chars // 4}) + _close_prop_ctx(_litellm_gen_prop) + finalized = True + except Exception as ex: + logger.error(f"Stream error: {ex}") + # End child span before parent on stream error (CodeRabbit: missing finalization) + _end_child_span(litellm_span_obj, + output={"error": type(ex).__name__}, + metadata={"status": "failed"}, ) - if agy_span_obj: + # End parent trace on stream error (before any cooldown logic) + _end_parent_obs(parent_obs, + output={"error": type(ex).__name__, "route": "litellm_fallback", + "stream": True}) + _close_prop_ctx(_litellm_gen_prop) + finalized = True + if model_name.startswith("ollama-"): + global _ollama_cooldown_until + _ollama_cooldown_until = ( + time.monotonic() + OLLAMA_COOLDOWN_SECONDS + ) try: - agy_span_obj.end( - output={"error": type(stream_err).__name__}, - metadata={"status": "failed"}, + await save_cooldowns_to_valkey() + logger.error( + f"🧊 Ollama failed midway through stream, activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" ) - except Exception: - pass - raise + except Exception as save_err: + logger.warning( + f"Failed to save cooldowns to Valkey: {save_err}" + ) + finally: + if not finalized: + _end_child_span(litellm_span_obj, + output={"error": "cancelled"}, + metadata={"status": "cancelled"}, + ) + _end_parent_obs(parent_obs, + output={"error": "cancelled", "route": "litellm_fallback", + "stream": True}) + _close_prop_ctx(_litellm_gen_prop) + await r.aclose() return StreamingResponse( - native_agy_stream_generator( - agy_response["stream"], model_name - ), - media_type="text/event-stream", + stream_generator(), media_type="text/event-stream" ) else: - latency_ms = (time.time() - start_time) * 1000.0 - stats["total_proxy_time_ms"] += latency_ms + error_body = await r.aread() if r else b"" + logger.warning( + f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}" + ) + await r.aclose() + # Finalize child span before raising on stream connection failure + # parent_obs finalized by outer handler (HTTPException → except block) + _end_child_span(litellm_span_obj, + output={"status": r.status_code, "error": "litellm_stream_failed"}, + metadata={"status": "failed"}, + ) + raise HTTPException( + status_code=r.status_code, + detail="LiteLLM upstream request failed", + ) + else: + logger.info(f"Proxying to LiteLLM as model={model_name}") + response = await client.post( + f"{backend_api_base}/chat/completions", + json=body_to_send, + headers=headers, + ) + if response.status_code == 200: + proxy_latency = (time.time() - proxy_start) * 1000.0 + stats["total_proxy_time_ms"] += proxy_latency stats["avg_proxy_latency_ms"] = ( stats["total_proxy_time_ms"] / stats["total_requests"] ) - usage = agy_response.get("usage") or {} - prompt_tokens = usage.get("prompt_tokens") or 0 - completion_tokens = usage.get("completion_tokens") or 0 + resp_json = response.json() + usage = resp_json.get("usage") or {} + prompt_tokens = usage.get( + "prompt_tokens" + ) or estimate_prompt_tokens(body_to_send) + choices = resp_json.get("choices") or [] + fallback_completion = 0 + if choices and isinstance(choices[0], dict): + msg = choices[0].get("message") + if isinstance(msg, dict): + fallback_completion = len(msg.get("content") or "") // 4 + completion_tokens = ( + usage.get("completion_tokens") or fallback_completion + ) record_tool_usage(ToolUsageRecord( active_tool, prompt_tokens, completion_tokens, model_name, - latency_ms, - route="google_oauth_direct", + proxy_latency, + route="litellm_fallback", )) - logger.info( - f"✅ agy proxy succeeded: {model_name}, {latency_ms:.0f}ms" + # Finalize LiteLLM span (non-streaming path) + _end_child_span(litellm_span_obj, + output={ + "model": model_name, + "tokens": completion_tokens, + }, + metadata={"latency_ms": proxy_latency}, + ) + # Finalize parent trace (non-streaming path) + _end_parent_obs(parent_obs, + output={"model": model_name, "tier": target_model, + "route": "litellm_fallback"}, + metadata={"latency_ms": proxy_latency, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + return resp_json + else: + logger.warning( + f"LiteLLM failed ({response.status_code}): {response.text[:300]}" + ) + # Finalize child span before raising on non-200 response + _end_child_span(litellm_span_obj, + output={"status": response.status_code, "error": "litellm_upstream_failed"}, + metadata={"status": "failed"}, + ) + raise HTTPException( + status_code=response.status_code, + detail="LiteLLM upstream request failed", ) - - # Finalize agy span - if agy_span_obj: - try: - agy_span_obj.end( - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass - - if is_stream_requested: - # Robust fallback: simulate stream if we requested stream but got buffered response - content = (agy_response.get("choices") or [{}])[0].get( - "message", {} - ).get("content") or "" - - async def agy_stream_generator(): - """Asynchronous generator yielding simulated OpenAI-compatible streaming chunks from a static agy response.""" - import uuid - - created_time = int(time.time()) - chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" - chunk_size = 40 - for i in range(0, len(content), chunk_size): - chunk_text = content[i : i + chunk_size] - chunk_data = { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created_time, - "model": model_name, - "choices": [ - { - "index": 0, - "delta": {"content": chunk_text}, - "finish_reason": None, - } - ], - } - yield f"data: {json.dumps(chunk_data)}\n\n".encode( - "utf-8" - ) - await asyncio.sleep(0.005) - - finish_data = { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created_time, - "model": model_name, - "choices": [ - { - "index": 0, - "delta": {}, - "finish_reason": "stop", - } - ], - } - yield f"data: {json.dumps(finish_data)}\n\n".encode( - "utf-8" - ) - yield b"data: [DONE]\n\n" - - return StreamingResponse( - agy_stream_generator(), media_type="text/event-stream" - ) - else: - return agy_response - except ImportError: - if agy_span_obj: - try: - agy_span_obj.end( - output={"error": "module_not_available"}, - metadata={"status": "skipped"}, - ) - except Exception: - pass - logger.warning("agy_proxy module not available, falling back to LiteLLM") - except Exception as e: - if agy_span_obj: - try: - agy_span_obj.end( - output={"error": type(e).__name__}, - metadata={"status": "failed"}, - ) - except Exception: - pass - logger.error(f"agy proxy failed: {type(e).__name__}, falling back to LiteLLM") - - if target_model == "llm-routing-agy": - target_model = "agent-advanced-core" - original_target_model = target_model - - # --- OLLAMA (via LiteLLM) --- - # LiteLLM's ollama_chat provider handles the native Ollama API call. - # We just proxy to LiteLLM with the appropriate model name. - # LiteLLM's fallback chain handles failures. - if should_try_ollama: - if client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama"): - if target_model in ("agent-advanced-core", "agent-reasoning-core"): - target_model = "ollama-deepseek-v4-pro" - elif target_model == "agent-complex-core": - target_model = "ollama-deepseek-v4-flash" - elif client_model == "llm-routing-ollama": - if target_model in ("agent-advanced-core", "agent-reasoning-core"): - target_model = "ollama-deepseek-v4-pro" - else: - target_model = "ollama-deepseek-v4-flash" - else: - # Fallback (e.g. if LiteLLM fallback loops back with model: llm-routing-ollama) - if target_model in ("agent-advanced-core", "agent-reasoning-core"): - target_model = "ollama-deepseek-v4-pro" - else: - target_model = "ollama-deepseek-v4-flash" - 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: - logger.error(f"Backend '{model_name}' not found in configuration backends.") - raise HTTPException( - status_code=500, detail=f"Backend {model_name} misconfigured" - ) - - backend_api_base = backend_conf["api_base"] - backend_api_key = backend_conf["api_key"] - if backend_api_key == "DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER": - backend_api_key = os.getenv("LITELLM_MASTER_KEY", backend_api_key) - - logger.info(f"Proxying to LiteLLM as model={model_name}") - - # --- Langfuse child span: LiteLLM proxy --- - litellm_span_obj = None - if langfuse_trace_id: - lf_litellm = get_langfuse() - if lf_litellm: - try: - litellm_span_obj = lf_litellm.start_observation( - trace_context={"trace_id": langfuse_trace_id}, - name="litellm-proxy", - input=model_name, - metadata={"model": model_name}, - level="DEFAULT", - ) - except Exception: - pass - - client = get_http_client() - try: - headers = {"Authorization": f"Bearer {backend_api_key}"} - if langfuse_trace_id: - headers["X-Langfuse-Trace-Id"] = langfuse_trace_id - - # Handle streaming vs non-streaming proxying (LiteLLM handles fallback internally) - proxy_start = time.time() - - # --- Pre-screening: clamp max_tokens to fit within downstream model context limits --- - try: - body_to_send = body.copy() - body_to_send["model"] = model_name - requested_max_tokens = body_to_send.get("max_tokens", 4096) - - # Tier-aware minimum context length (from actual roster data): - # - agent-simple-core: 32K (includes tiny liquid/dolphin models) - # - agent-medium-core+: 256K (smallest non-tiny model is nemotron-nano-omni at 256K) - # - ollama-deepseek-v4-*: 1M (DeepSeek V4 native context) - _tier_min_ctx = { - "agent-simple-core": 32768, - "ollama-deepseek-v4-pro": 524288, - "ollama-deepseek-v4-flash": 524288, - } - _min_ctx = _tier_min_ctx.get(model_name, 262144) - _est_input = estimate_prompt_tokens(body_to_send) - _safe_max = _min_ctx - _est_input - 2048 # 2K safety margin - if _safe_max < 1024: - raise HTTPException( - status_code=400, - detail=f"Context window exceeded. Estimated input tokens ({_est_input}) plus safety margin (2048) exceeds model context limit ({_min_ctx}).", - ) - if requested_max_tokens > _safe_max: - logger.warning( - f"⛔ Clamping max_tokens: {requested_max_tokens} → {_safe_max} " - f"(est_input={_est_input}, min_ctx={_min_ctx}, tier={model_name})" - ) - body_to_send["max_tokens"] = _safe_max except HTTPException: raise - except Exception as e: - logger.warning(f"Pre-screening failed (non-fatal): {e}") - body_to_send = body.copy() - body_to_send["model"] = model_name - if "metadata" not in body_to_send or not isinstance( - body_to_send["metadata"], dict - ): - body_to_send["metadata"] = {} - body_to_send["metadata"]["trace_name"] = "agent-completion" - - if body.get("stream", False): - logger.info(f"Proxying streaming to LiteLLM as model={model_name}") - req = client.build_request( - "POST", - f"{backend_api_base}/chat/completions", - json=body_to_send, - headers=headers, + except Exception as exc: + logger.error(f"httpx call failed: {exc}") + # Finalize child span before raising on proxy exception + _end_child_span(litellm_span_obj, + output={"error": type(exc).__name__}, + metadata={"status": "failed"}, ) - r = await client.send(req, stream=True) - if r.status_code == 200: - - async def stream_generator(): - """Asynchronous generator that yields streaming chunks from LiteLLM completions response and logs usage stats on completion.""" - import codecs - - completion_chars = 0 - request_tokens = estimate_prompt_tokens(body_to_send) - sse_buffer = "" - decoder = codecs.getincrementaldecoder("utf-8")() - try: - async for chunk in r.aiter_bytes(): - yield chunk - try: - sse_buffer += decoder.decode(chunk) - while "\n" in sse_buffer: - line, sse_buffer = sse_buffer.split("\n", 1) - line = line.strip() - if line.startswith("data:"): - data_str = line[5:].strip() - if data_str == "[DONE]": - continue - try: - data_json = json.loads(data_str) - choices = data_json.get("choices", []) - if choices and isinstance( - choices[0], dict - ): - delta = choices[0].get("delta") - if isinstance(delta, dict): - content = ( - delta.get("content") or "" - ) - completion_chars += len(content) - except Exception: - pass - except Exception: - pass - proxy_latency = (time.time() - proxy_start) * 1000.0 - stats["total_proxy_time_ms"] += proxy_latency - stats["avg_proxy_latency_ms"] = ( - stats["total_proxy_time_ms"] / stats["total_requests"] - ) - record_tool_usage(ToolUsageRecord( - active_tool, - request_tokens, - completion_chars // 4, - model_name, - proxy_latency, - route="litellm_fallback", - )) - # Finalize LiteLLM span (streaming path) - if litellm_span_obj: - try: - litellm_span_obj.end( - output={"model": model_name, "stream": True}, - metadata={ - "latency_ms": proxy_latency, - "tokens": completion_chars // 4, - }, - ) - except Exception: - pass - except Exception as ex: - logger.error(f"Stream error: {ex}") - if model_name.startswith("ollama-"): - global _ollama_cooldown_until - _ollama_cooldown_until = ( - time.monotonic() + OLLAMA_COOLDOWN_SECONDS - ) - try: - await save_cooldowns_to_valkey() - logger.error( - f"🧊 Ollama failed midway through stream, activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" - ) - except Exception as save_err: - logger.warning( - f"Failed to save cooldowns to Valkey: {save_err}" - ) - finally: - await r.aclose() - - return StreamingResponse( - stream_generator(), media_type="text/event-stream" + raise HTTPException( + status_code=502, detail="Proxy call failed" + ) from exc + + if should_try_ollama: + # Sync state from Valkey first + await sync_cooldowns_from_valkey() + + # --- Router-side Ollama cooldown check --- + global _ollama_cooldown_until + now_mono = time.monotonic() + if now_mono < _ollama_cooldown_until: + remaining = int(_ollama_cooldown_until - now_mono) + logger.warning( + f"⏳ Ollama cooldown active ({remaining}s remaining), " + f"skipping {target_model}" + ) + if client_model in ( + "llm-routing-auto-ollama", + "llm-routing-auto-agy-ollama", + ): + # Auto mode: silently fall through to the free tier + logger.info( + f"Auto-mode fallback: {target_model} → {original_target_model} (Ollama cooled down)" ) + try: + return await execute_proxy(original_target_model) + except HTTPException: + _end_parent_obs(parent_obs, + output={"error": "all_backends_failed", "route": "ollama_cooldown_fallback"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise else: - error_body = await r.aread() if r else b"" - logger.warning( - f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}" - ) - await r.aclose() + # Direct/fallback llm-routing-ollama: return 429 so LiteLLM + # skips this model group and moves to openrouter-auto + _end_parent_obs(parent_obs, + output={"error": "ollama_cooldown", "route": "ollama"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True raise HTTPException( - status_code=r.status_code, - detail="LiteLLM upstream request failed", - ) - else: - logger.info(f"Proxying to LiteLLM as model={model_name}") - response = await client.post( - f"{backend_api_base}/chat/completions", - json=body_to_send, - headers=headers, - ) - if response.status_code == 200: - proxy_latency = (time.time() - proxy_start) * 1000.0 - stats["total_proxy_time_ms"] += proxy_latency - stats["avg_proxy_latency_ms"] = ( - stats["total_proxy_time_ms"] / stats["total_requests"] + status_code=429, + detail=f"Ollama backend cooled down ({remaining}s remaining)", ) - resp_json = response.json() - usage = resp_json.get("usage") or {} - prompt_tokens = usage.get( - "prompt_tokens" - ) or estimate_prompt_tokens(body_to_send) - choices = resp_json.get("choices") or [] - fallback_completion = 0 - if choices and isinstance(choices[0], dict): - msg = choices[0].get("message") - if isinstance(msg, dict): - fallback_completion = len(msg.get("content") or "") // 4 - completion_tokens = ( - usage.get("completion_tokens") or fallback_completion + + try: + result = await execute_proxy(target_model) + return result + except HTTPException as e: + is_transient = e.status_code in (429, 500, 502, 503, 504) + if is_transient: + # Ollama failure — activate router-side cooldown + _ollama_cooldown_until = time.monotonic() + OLLAMA_COOLDOWN_SECONDS + await save_cooldowns_to_valkey() + logger.error( + f"🧊 Ollama failed ({e.status_code}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" ) - record_tool_usage(ToolUsageRecord( - active_tool, - prompt_tokens, - completion_tokens, - model_name, - proxy_latency, - route="litellm_fallback", - )) - # Finalize LiteLLM span (non-streaming path) - if litellm_span_obj: + if client_model in ( + "llm-routing-auto-ollama", + "llm-routing-auto-agy-ollama", + ): + if is_transient: + logger.warning( + f"Ollama proxy failed ({e.detail}), falling back to free tier {original_target_model}" + ) try: - litellm_span_obj.end( - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={"latency_ms": proxy_latency}, - ) - except Exception: - pass - return resp_json + return await execute_proxy(original_target_model) + except HTTPException: + _end_parent_obs(parent_obs, + output={"error": "all_backends_failed", "route": "ollama_fallback"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise + else: + _end_parent_obs(parent_obs, + output={"error": f"ollama_non_transient_{e.status_code}", "route": "ollama"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise e else: - logger.warning( - f"LiteLLM failed ({response.status_code}): {response.text[:300]}" - ) - raise HTTPException( - status_code=response.status_code, - detail="LiteLLM upstream request failed", - ) - except HTTPException: - raise - except Exception as exc: - logger.error(f"httpx call failed: {exc}") - raise HTTPException( - status_code=502, detail="Proxy call failed" - ) from exc - - if should_try_ollama: - # Sync state from Valkey first - await sync_cooldowns_from_valkey() - - # --- Router-side Ollama cooldown check --- - global _ollama_cooldown_until - now_mono = time.monotonic() - if now_mono < _ollama_cooldown_until: - remaining = int(_ollama_cooldown_until - now_mono) - logger.warning( - f"⏳ Ollama cooldown active ({remaining}s remaining), " - f"skipping {target_model}" - ) - if client_model in ( - "llm-routing-auto-ollama", - "llm-routing-auto-agy-ollama", - ): - # Auto mode: silently fall through to the free tier - logger.info( - f"Auto-mode fallback: {target_model} → {original_target_model} (Ollama cooled down)" - ) - return await execute_proxy(original_target_model) - else: - # Direct/fallback llm-routing-ollama: return 429 so LiteLLM - # skips this model group and moves to openrouter-auto - raise HTTPException( - status_code=429, - detail=f"Ollama backend cooled down ({remaining}s remaining)", - ) - - try: - result = await execute_proxy(target_model) - return result - except HTTPException as e: - is_transient = e.status_code in (429, 500, 502, 503, 504) - if is_transient: - # Ollama failure — activate router-side cooldown + # Direct/fallback llm-routing-ollama request + if is_transient: + logger.error( + f"Ollama proxy failed ({e.detail}) for direct/fallback request, returning 429" + ) + _end_parent_obs(parent_obs, + output={"error": "ollama_rate_limited", "route": "ollama"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise HTTPException( + status_code=429, + detail="Ollama backend rate limited/unavailable", + ) from e + else: + _end_parent_obs(parent_obs, + output={"error": f"ollama_non_transient_{e.status_code}", "route": "ollama"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise e + except Exception as e: + # Unexpected error (timeouts, connection issues) — also cooldown to prevent hammering _ollama_cooldown_until = time.monotonic() + OLLAMA_COOLDOWN_SECONDS await save_cooldowns_to_valkey() logger.error( - f"🧊 Ollama failed ({e.status_code}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" + f"🧊 Ollama unexpected error ({e}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" ) - if client_model in ( - "llm-routing-auto-ollama", - "llm-routing-auto-agy-ollama", - ): - if is_transient: + if client_model in ( + "llm-routing-auto-ollama", + "llm-routing-auto-agy-ollama", + ): logger.warning( - f"Ollama proxy failed ({e.detail}), falling back to free tier {original_target_model}" + f"Ollama proxy error ({e}), falling back to free tier {original_target_model}" ) - return await execute_proxy(original_target_model) + try: + return await execute_proxy(original_target_model) + except HTTPException: + _end_parent_obs(parent_obs, + output={"error": "all_backends_failed", "route": "ollama_unexpected_fallback"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise else: - raise e - else: - # Direct/fallback llm-routing-ollama request - if is_transient: - logger.error( - f"Ollama proxy failed ({e.detail}) for direct/fallback request, returning 429" - ) + _end_parent_obs(parent_obs, + output={"error": type(e).__name__, "route": "ollama"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True raise HTTPException( - status_code=429, - detail="Ollama backend rate limited/unavailable", + status_code=429, detail="Ollama backend rate limited/unavailable" ) from e - else: - raise e - except Exception as e: - # Unexpected error (timeouts, connection issues) — also cooldown to prevent hammering - _ollama_cooldown_until = time.monotonic() + OLLAMA_COOLDOWN_SECONDS - await save_cooldowns_to_valkey() - logger.error( - f"🧊 Ollama unexpected error ({e}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown" - ) - if client_model in ( - "llm-routing-auto-ollama", - "llm-routing-auto-agy-ollama", - ): - logger.warning( - f"Ollama proxy error ({e}), falling back to free tier {original_target_model}" - ) - return await execute_proxy(original_target_model) - else: - raise HTTPException( - status_code=429, detail="Ollama backend rate limited/unavailable" - ) from e - else: - return await execute_proxy(target_model) + else: + try: + return await execute_proxy(target_model) + except HTTPException: + _end_parent_obs(parent_obs, + output={"error": "all_backends_failed", "route": "default_proxy"}) + _close_prop_ctx(_prop_ctx) + _non_streaming_finalized = True + raise + finally: + if not _is_streaming and not _non_streaming_finalized: + _end_parent_obs(parent_obs, + output={"error": "cancelled", "route": "non_streaming"}) + _prop_ctx = _close_prop_ctx(_prop_ctx) + @app.get("/metrics") @@ -2758,7 +3094,7 @@ async def get_dashboard_data(): asyncio.wait_for(sync_cooldowns_from_valkey(), timeout=2.0), check_tcp_port("127.0.0.1", _valkey_port()), check_http_endpoint(f"http://127.0.0.1:{os.getenv('LITELLM_PORT') or '4000'}/"), - check_http_endpoint(f"{LLAMA_SERVER_URL}/health"), + asyncio.wait_for(_check_llama_health(), timeout=3.0), check_http_endpoint(f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}"), get_gemini_oauth_status(), asyncio.wait_for(get_best_free_model(), timeout=5.0), diff --git a/scripts/verification/verify_canonical_endpoints.py b/scripts/verification/verify_canonical_endpoints.py index 23be879f..a0a17045 100644 --- a/scripts/verification/verify_canonical_endpoints.py +++ b/scripts/verification/verify_canonical_endpoints.py @@ -71,6 +71,8 @@ def _parse(path: Path): "public_base_url": (os.environ.get("PUBLIC_BASE_URL") or env.get("PUBLIC_BASE_URL") or "").rstrip("/"), "minio_s3_port": os.environ.get("MINIO_S3_PORT") or env.get("MINIO_S3_PORT") or "9002", "clickhouse_http_port": os.environ.get("CLICKHOUSE_HTTP_PORT") or env.get("CLICKHOUSE_HTTP_PORT") or "8123", + "langfuse_public_key": os.environ.get("LANGFUSE_PUBLIC_KEY") or env.get("LANGFUSE_PUBLIC_KEY") or "", + "langfuse_secret_key": os.environ.get("LANGFUSE_SECRET_KEY") or env.get("LANGFUSE_SECRET_KEY") or "", } @@ -365,6 +367,194 @@ def test_litellm_direct_chat(cfg: dict) -> tuple[int, int]: return passed, total +def test_langfuse_session_propagation(cfg: dict) -> tuple[int, int]: + """Verify Langfuse session/user propagation and no cross-request leak. + + 1. Send a chat request WITH session_id + user, verify trace has them. + 2. Send a second chat request WITHOUT session info, verify NO leak. + Returns (passed, total). + """ + router_base = f"http://127.0.0.1:{cfg['router_port']}" + lf_host = cfg["langfuse_web_port"] + lf_base = f"http://127.0.0.1:{lf_host}" + pk = cfg["langfuse_public_key"] + sk = cfg["langfuse_secret_key"] + key = cfg["router_api_key"] + + if not pk or not sk: + print( + "\n── Langfuse session propagation — SKIPPED " + "(LANGFUSE_PUBLIC_KEY/SECRET_KEY not set) ──" + ) + return 0, 0 + + auth = (pk, sk) + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + session_id = f"e2e-verify-{os.urandom(4).hex()}" + user_id = "verify-script" + passed = total = 0 + + print("\n── Langfuse session propagation ──") + + # --- Step 1: Request WITH session --- + total += 1 + payload = { + "model": "agent-simple-core", + "messages": [{"role": "user", "content": "Say 'test' and nothing else."}], + "max_tokens": 5, + "session_id": session_id, + "user": user_id, + } + try: + r = httpx.post( + f"{router_base}/v1/chat/completions", + json=payload, + headers=headers, + timeout=120, + ) + if r.status_code != 200: + passed += check( + "Session request (200)", False, + f"HTTP {r.status_code}: {r.text[:100]}" + ) + return passed, total + passed += check("Session request (200)", True, session_id[:16]) + except Exception as e: + passed += check("Session request (200)", False, str(e)[:100]) + return passed, total + + # Wait for Langfuse SDK auto-flush (poll up to 10s) + session_trace_id = None + _attempt = 0 + for _attempt in range(10): + time.sleep(1) + try: + resp = httpx.get( + f"{lf_base}/api/public/traces?page=1&limit=20&orderBy=timestamp.desc", + auth=auth, timeout=10, + ) + if resp.status_code != 200: + continue + traces = resp.json().get("data", []) + for t in traces: + if t.get("sessionId") == session_id and t.get("userId") == user_id: + session_trace_id = t["id"] + break + if session_trace_id: + break + except Exception as e: + if _attempt == 9: + print(f" ⚠ trace poll error: {str(e)[:100]}") + continue + + # Report session trace result + total += 1 + if session_trace_id: + passed += check( + "Trace has session+user", True, + f"session={session_id[:12]} user={user_id} (found after {_attempt+1}s)", + ) + else: + passed += check( + "Trace has session+user", False, + f"session={session_id[:12]} not found after 10s polling", + ) + + # --- Step 2: Request WITHOUT session (leak test) --- + total += 1 + payload_no_session = { + "model": "agent-simple-core", + "messages": [{"role": "user", "content": "Say 'leaktest' and nothing else."}], + "max_tokens": 5, + } + try: + r2 = httpx.post( + f"{router_base}/v1/chat/completions", + json=payload_no_session, + headers=headers, + timeout=120, + ) + if r2.status_code != 200: + passed += check( + "No-session request (200)", False, + f"HTTP {r2.status_code}: {r2.text[:100]}" + ) + return passed, total + passed += check("No-session request (200)", True, "clean") + except Exception as e: + passed += check("No-session request (200)", False, str(e)[:100]) + return passed, total + + # Wait for Langfuse SDK auto-flush (poll up to 10s for leak check) + session_trace_visible = False + traces2 = [] + for _attempt2 in range(10): + time.sleep(1) + try: + resp2 = httpx.get( + f"{lf_base}/api/public/traces?page=1&limit=20&orderBy=timestamp.desc", + auth=auth, timeout=10, + ) + if resp2.status_code != 200: + continue + traces2 = resp2.json().get("data", []) + # Ensure the second request's trace has been flushed before checking for leaks. + # The session trace from step 1 is already in Langfuse, so we wait until + # a NEWER trace appears (session_trace_id is NOT at index 0). + recent_ids = [t["id"] for t in traces2] + if session_trace_id and session_trace_id in recent_ids: + if recent_ids.index(session_trace_id) > 0: + session_trace_visible = True + break + # If session trace hasn't appeared yet, continue waiting + if not session_trace_id: + # Session trace was never found in step 1; leak check is inconclusive. + # Poll for remaining time, then break without a definitive answer. + if _attempt2 >= 5 and len(traces2) > 0: + break + except Exception as e: + if _attempt2 == 9: + print(f" ⚠ trace poll error: {str(e)[:100]}") + continue + + # Report leak test result + # If session_trace_id was never found in step 1, the leak check is inconclusive. + # We cannot distinguish between "no leak" and "session trace still unflushed." + # Don't count this as a failure — just warn and move on. + if session_trace_id is None: + print(" ⚠ No session leak — INCONCLUSIVE (session trace not found in step 1 polling)") + return passed, total + leaked = False + if session_trace_visible and traces2: + leaked = any( + t.get("sessionId") == session_id and t.get("id") != session_trace_id + for t in traces2 + ) + elif _attempt2 == 9: + print( + " ⚠ No session leak — INCONCLUSIVE " + "(second request trace not flushed within 10s)" + ) + return passed, total + total += 1 + + if leaked: + passed += check( + "No session leak", False, + "Previous session leaked into no-session request!" + ) + else: + passed += check( + "No session leak", True, + f"{len(traces2)} recent traces clean (excluded known session trace)" + ) + + return passed, total + + def test_canonical_urls(cfg: dict) -> tuple[int, int, int]: """Verify canonical HTTPS URLs are reachable (if PUBLIC_BASE_URL is set). Returns (passed, total, skipped).""" @@ -460,6 +650,7 @@ def main(): ("Infrastructure", test_infra_health), ("E2E chat (router)", test_e2e_chat), ("LiteLLM direct chat", test_litellm_direct_chat), + ("Langfuse session propagation", test_langfuse_session_propagation), ("Canonical URLs", test_canonical_urls), ]: result = fn(cfg)