diff --git a/router/main.py b/router/main.py index 35844949..0f672be8 100644 --- a/router/main.py +++ b/router/main.py @@ -293,6 +293,57 @@ 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: + 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 + + +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: + 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: @@ -1104,17 +1155,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 +1188,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()) @@ -1878,6 +1921,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 +1950,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 +1987,10 @@ 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}"}) + await _flush_langfuse_async() raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2125,32 +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: - agy_span_obj.end( - 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}) + _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: - agy_span_obj.end( - 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( @@ -2181,20 +2258,16 @@ async def native_agy_stream_generator(stream_gen, model_name): ) # 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 + _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 @@ -2246,31 +2319,39 @@ 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}) + _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" ) 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( - 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: - agy_span_obj.end( - 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": @@ -2444,19 +2525,36 @@ async def stream_generator(): 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 + _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}) + _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 = ( @@ -2483,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", @@ -2523,22 +2630,35 @@ async def stream_generator(): route="litellm_fallback", )) # Finalize LiteLLM span (non-streaming path) - if litellm_span_obj: - try: - litellm_span_obj.end( - 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, + "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( 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", @@ -2547,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