fix: add trace-level metadata for all chat messages

Author: jake-marsden

src/api/models/bedrock.py

@@ -438,7 +438,7 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
                     # All other chunks will also include a usage field, but with a null value.
                     yield self.stream_response_to_bytes(stream_response)
 
-            # Update Langfuse with final streaming metadata
+            # Update Langfuse with final streaming metadata (both observation and trace)
             if final_usage or accumulated_output:
                 update_params = {}
                 if accumulated_output:
@@ -463,10 +463,22 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
                 if metadata:
                     update_params["metadata"] = metadata
 
+                # Update the child observation (Bedrock Converse)
                 langfuse_context.update_current_observation(**update_params)
+                
+                # Also update the parent trace (chat_completion) with final output
+                trace_output = {
+                    "message": {
+                        "role": "assistant",
+                        "content": final_output if accumulated_output else None,
+                    },
+                    "finish_reason": finish_reason,
+                }
+                langfuse_context.update_current_trace(output=trace_output)
+                
                 if DEBUG:
                     output_length = len(accumulated_output)
-                    logger.info(f"Langfuse: Updated observation with streaming output - "
+                    logger.info(f"Langfuse: Updated observation and trace with streaming output - "
                               f"chunks_count={output_length}, "
                               f"output_chars={len(final_output) if accumulated_output else 0}, "
                               f"input_tokens={final_usage.prompt_tokens if final_usage else 'N/A'}, "
@@ -482,11 +494,15 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
             raise
         except Exception as e:
             logger.error("Stream error for model %s: %s", chat_request.model, str(e))
-            # Update Langfuse with error
+            # Update Langfuse with error (both observation and trace)
             langfuse_context.update_current_observation(
                 level="ERROR",
                 status_message=f"Stream error: {str(e)}"
             )
+            langfuse_context.update_current_trace(
+                output={"error": str(e)},
+                metadata={"error": True}
+            )
             if DEBUG:
                 logger.info(f"Langfuse: Updated observation with streaming error - error={str(e)[:100]}")
             error_event = Error(error=ErrorMessage(message=str(e)))

src/api/routers/chat.py

@@ -1,7 +1,8 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Body, Depends
+from fastapi import APIRouter, Body, Depends, Header, Request
 from fastapi.responses import StreamingResponse
+from langfuse.decorators import langfuse_context, observe
 
 from api.auth import api_key_auth
 from api.models.bedrock import BedrockModel
@@ -15,10 +16,52 @@
 )
 
 
+def extract_langfuse_metadata(chat_request: ChatRequest, headers: dict) -> dict:
+    """Extract Langfuse tracing metadata from request body and headers.
+    
+    Metadata can be provided via:
+    1. extra_body.langfuse_metadata dict in the request
+    2. HTTP headers: X-Chat-Id, X-User-Id, X-Session-Id, X-Message-Id
+    3. user field in the request (for user_id)
+    
+    Returns a dict with: user_id, session_id, chat_id, message_id, and any custom metadata
+    """
+    metadata = {}
+    
+    # Extract from extra_body if present
+    if chat_request.extra_body and isinstance(chat_request.extra_body, dict):
+        langfuse_meta = chat_request.extra_body.get("langfuse_metadata", {})
+        if isinstance(langfuse_meta, dict):
+            metadata.update(langfuse_meta)
+    
+    # Extract from headers
+    headers_lower = {k.lower(): v for k, v in headers.items()}
+    
+    # Map headers to metadata fields
+    header_mapping = {
+        "x-chat-id": "chat_id",
+        "x-user-id": "user_id", 
+        "x-session-id": "session_id",
+        "x-message-id": "message_id",
+    }
+    
+    for header_key, meta_key in header_mapping.items():
+        if header_key in headers_lower and headers_lower[header_key]:
+            metadata[meta_key] = headers_lower[header_key]
+    
+    # Use the 'user' field from request as user_id if not already set
+    if "user_id" not in metadata and chat_request.user:
+        metadata["user_id"] = chat_request.user
+    
+    return metadata
+
+
 @router.post(
     "/completions", response_model=ChatResponse | ChatStreamResponse | Error, response_model_exclude_unset=True
 )
+@observe(as_type="generation", name="chat_completion")
 async def chat_completions(
+    request: Request,
     chat_request: Annotated[
         ChatRequest,
         Body(
@@ -34,12 +77,45 @@ async def chat_completions(
         ),
     ],
 ):
+    # Extract metadata for Langfuse tracing
+    metadata = extract_langfuse_metadata(chat_request, dict(request.headers))
+    
+    # Create trace name using chat_id if available
+    trace_name = f"chat:{metadata.get('chat_id', 'unknown')}"
+    
+    # Update trace with metadata, user_id, and session_id
+    langfuse_context.update_current_trace(
+        name=trace_name,
+        user_id=metadata.get("user_id"),
+        session_id=metadata.get("session_id"),
+        metadata=metadata,
+        input={
+            "model": chat_request.model,
+            "messages": [msg.model_dump() for msg in chat_request.messages],
+            "temperature": chat_request.temperature,
+            "max_tokens": chat_request.max_tokens,
+            "tools": [tool.model_dump() for tool in chat_request.tools] if chat_request.tools else None,
+        }
+    )
+    
     if chat_request.model.lower().startswith("gpt-"):
         chat_request.model = DEFAULT_MODEL
 
     # Exception will be raised if model not supported.
     model = BedrockModel()
     model.validate(chat_request)
+    
     if chat_request.stream:
         return StreamingResponse(content=model.chat_stream(chat_request), media_type="text/event-stream")
-    return await model.chat(chat_request)
+    
+    response = await model.chat(chat_request)
+    
+    # Update trace with output for non-streaming
+    langfuse_context.update_current_trace(
+        output={
+            "message": response.choices[0].message.model_dump() if response.choices else None,
+            "finish_reason": response.choices[0].finish_reason if response.choices else None,
+        }
+    )
+    
+    return response

src/api/schema.py

@@ -99,7 +99,7 @@ class ChatRequest(BaseModel):
     stream_options: StreamOptions | None = None
     temperature: float | None = Field(default=1.0, le=2.0, ge=0.0)
     top_p: float | None = Field(default=1.0, le=1.0, ge=0.0)
-    user: str | None = None  # Not used
+    user: str | None = None
     max_tokens: int | None = 2048
     max_completion_tokens: int | None = None
     reasoning_effort: Literal["low", "medium", "high"] | None = None