From 5d06463b00f6a3f4a375939e1d59b2de600423b9 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:08:02 +0200 Subject: [PATCH 1/4] fix: langfuse v4 SDK compatibility + parent observation lifecycle - Add _end_parent_obs() and _end_child_span() helpers (v4: update() then end()) - Replace all .end(output=, metadata=) calls with helpers - Add _flush_langfuse_async() for async flush after each request - Add session_id, user_id, environment metadata to traces - Fix parent_obs.end() was never called (traces were incomplete) - Use await for non-streaming flush paths, asyncio.create_task for generators --- router/main.py | 128 +++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 119 insertions(+), 9 deletions(-) diff --git a/router/main.py b/router/main.py index 35844949..7f6e205c 100644 --- a/router/main.py +++ b/router/main.py @@ -293,6 +293,47 @@ 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 + try: + if output is not None or metadata is not None: + parent_obs.update(output=output, metadata=metadata) + parent_obs.end() + except Exception: + pass + + +def _end_child_span(span, output=None, metadata=None) -> None: + """Safely finalize a Langfuse child span (SDK v4: update + end). + + Non-fatal — errors are logged but never propagate. + """ + if span is None: + return + try: + if output is not None or metadata is not None: + span.update(output=output, metadata=metadata) + span.end() + except Exception: + pass + + +async def _flush_langfuse_async() -> None: + """Flush pending Langfuse events via thread pool (non-blocking, non-fatal).""" + try: + lf = get_langfuse() + if lf: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lf.flush) + except Exception: + pass + + async def push_aggregate_scores(): """Push aggregate KPIs as Langfuse scores every 5 minutes.""" while True: @@ -1106,7 +1147,7 @@ async def classify_request( if response.status_code != 200: if class_span_obj: try: - class_span_obj.end( + _end_child_span(class_span_obj, output={ "status": response.status_code, "error": "classification_failed", @@ -1143,7 +1184,7 @@ async def classify_request( # Finalize classifier child span if class_span_obj: try: - class_span_obj.end( + _end_child_span(class_span_obj, output={"tier": decision, "raw": raw_result}, metadata={"latency_ms": latency}, ) @@ -1878,6 +1919,21 @@ 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 @@ -1892,6 +1948,13 @@ async def chat_completions(request: Request): name=f"triage-{client_model}", input=last_user_message[:200], level="DEFAULT", + metadata={ + "client_model": client_model, + "environment": os.getenv("ENVIRONMENT", "production"), + "session_id": _trace_session_id or "", + "user_id": _trace_user_id or "", + "tags": [os.getenv("ENVIRONMENT", "production"), "llm-routing"], + }, ) except Exception as e: logger.warning(f"Langfuse trace init failed (non-fatal): {e}") @@ -1922,6 +1985,9 @@ async def chat_completions(request: Request): 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}"}) raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2127,7 +2193,7 @@ async def native_agy_stream_generator(stream_gen, model_name): ) if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={ "model": model_name, "tokens": token_count, @@ -2139,18 +2205,29 @@ async def native_agy_stream_generator(stream_gen, model_name): ) except Exception: pass + # 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}) + asyncio.create_task(_flush_langfuse_async()) except Exception as stream_err: logger.error( f"Error during native agy stream generation: {type(stream_err).__name__}" ) if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": type(stream_err).__name__}, metadata={"status": "failed"}, ) except Exception: pass + # End parent trace on stream error + _end_parent_obs(parent_obs, + output={"error": type(stream_err).__name__, + "route": "google_oauth_direct", "stream": True}) raise return StreamingResponse( @@ -2183,7 +2260,7 @@ async def native_agy_stream_generator(stream_gen, model_name): # Finalize agy span if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={ "model": model_name, "tokens": completion_tokens, @@ -2246,16 +2323,30 @@ async def agy_stream_generator(): "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}) + asyncio.create_task(_flush_langfuse_async()) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" ) 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}) + await _flush_langfuse_async() return agy_response except ImportError: if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": "module_not_available"}, metadata={"status": "skipped"}, ) @@ -2265,7 +2356,7 @@ async def agy_stream_generator(): except Exception as e: if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": type(e).__name__}, metadata={"status": "failed"}, ) @@ -2446,7 +2537,7 @@ async def stream_generator(): # Finalize LiteLLM span (streaming path) if litellm_span_obj: try: - litellm_span_obj.end( + _end_child_span(litellm_span_obj, output={"model": model_name, "stream": True}, metadata={ "latency_ms": proxy_latency, @@ -2455,8 +2546,19 @@ async def stream_generator(): ) except Exception: pass + # 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}) + asyncio.create_task(_flush_langfuse_async()) except Exception as ex: logger.error(f"Stream error: {ex}") + # 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}) if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2525,7 +2627,7 @@ async def stream_generator(): # Finalize LiteLLM span (non-streaming path) if litellm_span_obj: try: - litellm_span_obj.end( + _end_child_span(litellm_span_obj, output={ "model": model_name, "tokens": completion_tokens, @@ -2534,6 +2636,14 @@ async def stream_generator(): ) except Exception: pass + # 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}) + await _flush_langfuse_async() return resp_json else: logger.warning( From 655844db33bac2fae1b6c6959121c34d2e67aa04 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:22:29 +0200 Subject: [PATCH 2/4] =?UTF-8?q?fix:=20address=20all=20PR=20#296=20review?= =?UTF-8?q?=20comments=20=E2=80=94=20Langfuse=20trace=20lifecycle=20harden?= =?UTF-8?q?ing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Dynamic kwargs for _end_parent_obs / _end_child_span update() (Gemini: avoid overwriting None) - Fix _end_child_span docstring (CodeRabbit: 'logged' → 'never propagated') - Register all fire-and-forget flush tasks in _background_tasks set (Gemini+CodeRabbit: RUF006) - Flush Langfuse after error-path _end_parent_obs calls (Gemini: missing flush on stream errors) - End litellm_span_obj before parent on LiteLLM stream error (CodeRabbit: missing finalization) - Remove redundant if : + try/except wrappers around helper calls (CodeRabbit: dead code) - Add trace finalization to non-200 and exception error paths (CodeRabbit: pre-existing gaps) - Add trace finalization to streaming non-200 initial response (CodeRabbit: raised without ending traces) --- router/main.py | 220 ++++++++++++++++++++++++++----------------------- 1 file changed, 119 insertions(+), 101 deletions(-) diff --git a/router/main.py b/router/main.py index 7f6e205c..0f672be8 100644 --- a/router/main.py +++ b/router/main.py @@ -301,8 +301,13 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> None: if parent_obs is None: return try: - if output is not None or metadata is not None: - parent_obs.update(output=output, metadata=metadata) + 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: pass @@ -311,13 +316,18 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> 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 logged but never propagate. + Non-fatal — errors are never propagated. """ if span is None: return try: - if output is not None or metadata is not None: - span.update(output=output, metadata=metadata) + 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: pass @@ -1145,17 +1155,13 @@ async def classify_request( latency = (time.time() - start_time) * 1000.0 if response.status_code != 200: - if class_span_obj: - try: - _end_child_span(class_span_obj, - 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}" ) @@ -1182,14 +1188,10 @@ async def classify_request( decision = "agent-advanced-core" # Finalize classifier child span - if class_span_obj: - try: - _end_child_span(class_span_obj, - 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()) @@ -1988,6 +1990,7 @@ async def chat_completions(request: Request): # guard: end parent obs before raising _end_parent_obs(parent_obs, output={"error": f"Unknown model: {client_model}"}) + await _flush_langfuse_async() raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2191,43 +2194,40 @@ async def native_agy_stream_generator(stream_gen, model_name): logger.info( f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms" ) - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={ - "model": model_name, - "tokens": token_count, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass + _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}) - asyncio.create_task(_flush_langfuse_async()) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) except Exception as stream_err: logger.error( f"Error during native agy stream generation: {type(stream_err).__name__}" ) - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": type(stream_err).__name__}, - metadata={"status": "failed"}, - ) - except Exception: - pass + _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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) raise return StreamingResponse( @@ -2258,20 +2258,16 @@ async def native_agy_stream_generator(stream_gen, model_name): ) # Finalize agy span - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass + _end_child_span(agy_span_obj, + output={ + "model": model_name, + "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 @@ -2329,7 +2325,9 @@ async def agy_stream_generator(): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": len(content) // 4}) - asyncio.create_task(_flush_langfuse_async()) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" @@ -2344,24 +2342,16 @@ async def agy_stream_generator(): await _flush_langfuse_async() return agy_response except ImportError: - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": "module_not_available"}, - metadata={"status": "skipped"}, - ) - except Exception: - pass + _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: - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": type(e).__name__}, - metadata={"status": "failed"}, - ) - except Exception: - pass + _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": @@ -2535,30 +2525,36 @@ async def stream_generator(): route="litellm_fallback", )) # Finalize LiteLLM span (streaming path) - if litellm_span_obj: - try: - _end_child_span(litellm_span_obj, - output={"model": model_name, "stream": True}, - metadata={ - "latency_ms": proxy_latency, - "tokens": completion_chars // 4, - }, - ) - except Exception: - pass + _end_child_span(litellm_span_obj, + output={"model": model_name, "stream": True}, + metadata={ + "latency_ms": proxy_latency, + "tokens": completion_chars // 4, + }, + ) # 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}) - asyncio.create_task(_flush_langfuse_async()) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) 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"}, + ) # 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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2585,6 +2581,15 @@ async def stream_generator(): f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}" ) await r.aclose() + # Finalize traces before raising on stream connection failure + _end_child_span(litellm_span_obj, + output={"status": r.status_code, "error": "litellm_stream_failed"}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": f"HTTP {r.status_code}", "route": "litellm_fallback", + "stream": True}) + await _flush_langfuse_async() raise HTTPException( status_code=r.status_code, detail="LiteLLM upstream request failed", @@ -2625,17 +2630,13 @@ async def stream_generator(): route="litellm_fallback", )) # Finalize LiteLLM span (non-streaming path) - if litellm_span_obj: - try: - _end_child_span(litellm_span_obj, - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={"latency_ms": proxy_latency}, - ) - except Exception: - pass + _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, @@ -2649,6 +2650,15 @@ async def stream_generator(): logger.warning( f"LiteLLM failed ({response.status_code}): {response.text[:300]}" ) + # Finalize traces before raising on non-200 response + _end_child_span(litellm_span_obj, + output={"status": response.status_code, "error": "litellm_upstream_failed"}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": f"HTTP {response.status_code}", + "route": "litellm_fallback"}) + await _flush_langfuse_async() raise HTTPException( status_code=response.status_code, detail="LiteLLM upstream request failed", @@ -2657,6 +2667,14 @@ async def stream_generator(): raise except Exception as exc: logger.error(f"httpx call failed: {exc}") + # Finalize traces before raising on proxy exception + _end_child_span(litellm_span_obj, + output={"error": type(exc).__name__}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": type(exc).__name__, "route": "litellm_fallback"}) + await _flush_langfuse_async() raise HTTPException( status_code=502, detail="Proxy call failed" ) from exc From 505d2ed6508663f52e1ec79d78dcbefaca4885a0 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:33:41 +0200 Subject: [PATCH 3/4] fix: LLAMA_SERVER_URL via config placeholder + canonical URL (was hardcoded wrong port) - Moved LLAMA_SERVER_URL from hardcoded os.getenv to config.yaml os.environ/LLAMA_SERVER_URL - Added LLAMA_SERVER_URL=https://x570.vendeuvre.lan/llm-routing/llama to .env + .env.dev - Added get_llama_client() with verify=False for self-signed TLS (mirrors get_classifier_client) - Added _check_llama_health() using dedicated llama client - Fixes 'Failed to fetch llama.cpp metrics' warning (default was port 8080, actual is 8083/HAProxy) Closes #298 --- .env.dev | 1 + router/config.yaml | 2 ++ router/main.py | 59 +++++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 59 insertions(+), 3 deletions(-) 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 0f672be8..775f1973 100644 --- a/router/main.py +++ b/router/main.py @@ -28,7 +28,6 @@ from typing import Dict, Optional, Union 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" @@ -135,6 +134,35 @@ def get_classifier_client(): 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: + ca_bundle = os.getenv("LLAMA_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 + _llama_client = httpx.AsyncClient( + limits=_http_limits(), timeout=3600.0, verify=verify + ) + 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]') @@ -466,6 +494,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( + f"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 = _raw_llama_url.rstrip("/") + # Default colors for tool visualization badges and charts TOOL_COLORS = { "tree": "#34d399", # Green @@ -1063,6 +1106,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]: @@ -1491,7 +1544,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: @@ -2886,7 +2939,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), From d763ddb331cbd7305ff730a697d44e96351bf5be Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:53:57 +0200 Subject: [PATCH 4/4] fix: handle GeneratorExit in streaming generators + use propagate_attributes for Langfuse sessions GeneratorExit (BaseException, client disconnect) bypasses except Exception in async generators, leaving Langfuse traces open indefinitely. Fix: - native_agy_stream_generator: add finalized flag + finally block - agy_stream_generator (simulated): wrap in try/finally with finalized flag - stream_generator (LiteLLM): extend existing finally to handle cancellation Also fix Langfuse sessions: replace metadata-only session_id/user_id with propagate_attributes() context manager (native session mechanism in v4 SDK). Verified: 193/193 tests pass, 24/24 canonical endpoints pass, Langfuse API confirms sessionId propagation to child observations. --- router/main.py | 131 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 94 insertions(+), 37 deletions(-) diff --git a/router/main.py b/router/main.py index 775f1973..3bbf6765 100644 --- a/router/main.py +++ b/router/main.py @@ -27,6 +27,11 @@ 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("/") LANGFUSE_HOST = (os.getenv("LANGFUSE_HOST") or f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}").rstrip("/") @@ -1992,12 +1997,21 @@ async def chat_completions(request: Request): # --- Langfuse parent trace: create early so child spans can reference it --- langfuse_trace_id = None parent_obs = None + _prop_ctx = None lf = get_langfuse() if lf: try: langfuse_trace_id = lf.create_trace_id( seed=f"triage_{stats['total_requests']}" ) + # Propagate session_id/user_id via Langfuse's native session mechanism + if propagate_attributes and (_trace_session_id or _trace_user_id): + _prop_ctx = propagate_attributes( + session_id=_trace_session_id or None, + user_id=_trace_user_id or None, + tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], + ) + _prop_ctx.__enter__() parent_obs = lf.start_observation( trace_context={"trace_id": langfuse_trace_id}, name=f"triage-{client_model}", @@ -2006,15 +2020,18 @@ async def chat_completions(request: Request): metadata={ "client_model": client_model, "environment": os.getenv("ENVIRONMENT", "production"), - "session_id": _trace_session_id or "", - "user_id": _trace_user_id or "", - "tags": [os.getenv("ENVIRONMENT", "production"), "llm-routing"], }, ) except Exception as e: logger.warning(f"Langfuse trace init failed (non-fatal): {e}") langfuse_trace_id = None parent_obs = None + if _prop_ctx: + try: + _prop_ctx.__exit__(None, None, None) + except Exception: + pass + _prop_ctx = None if client_model in AUTO_MODELS or client_model == "llm-routing-ollama": # Full pipeline: classify → route to best tier @@ -2187,6 +2204,7 @@ async def native_agy_stream_generator(stream_gen, model_name): created_time = int(time.time()) chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" token_count = 0 + finalized = False try: async for token in stream_gen: if not token: @@ -2266,6 +2284,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) + finalized = True except Exception as stream_err: logger.error( f"Error during native agy stream generation: {type(stream_err).__name__}" @@ -2281,7 +2300,20 @@ async def native_agy_stream_generator(stream_gen, model_name): _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) + 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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) return StreamingResponse( native_agy_stream_generator( @@ -2335,9 +2367,29 @@ async def agy_stream_generator(): 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 = { + finalized = False + 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, @@ -2345,42 +2397,33 @@ async def agy_stream_generator(): "choices": [ { "index": 0, - "delta": {"content": chunk_text}, - "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" ) - 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" - # 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}) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + 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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) + finalized = True + finally: + if not finalized: + _end_parent_obs(parent_obs, + output={"error": "cancelled", + "route": "google_oauth_direct", "stream": True}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" @@ -2536,6 +2579,7 @@ async def stream_generator(): request_tokens = estimate_prompt_tokens(body_to_send) sse_buffer = "" decoder = codecs.getincrementaldecoder("utf-8")() + finalized = False try: async for chunk in r.aiter_bytes(): yield chunk @@ -2594,6 +2638,7 @@ async def stream_generator(): _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) + finalized = True except Exception as ex: logger.error(f"Stream error: {ex}") # End child span before parent on stream error (CodeRabbit: missing finalization) @@ -2608,6 +2653,7 @@ async def stream_generator(): _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) + finalized = True if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2623,6 +2669,17 @@ async def stream_generator(): 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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) await r.aclose() return StreamingResponse(