fix: esure Bedrock Converse observation nests properly in chat trace

Author: jake-marsden

src/api/models/bedrock.py

@@ -248,7 +248,6 @@ def validate(self, chat_request: ChatRequest):
                 detail=error,
             )
 
-    @observe(as_type="generation", name="Bedrock Converse")
     async def _invoke_bedrock(self, chat_request: ChatRequest, stream=False):
         """Common logic for invoke bedrock models"""
         if DEBUG:
@@ -259,29 +258,6 @@ async def _invoke_bedrock(self, chat_request: ChatRequest, stream=False):
         if DEBUG:
             logger.info("Bedrock request: " + json.dumps(str(args)))
 
-        # Extract model metadata for Langfuse
-        args_clone = args.copy()
-        messages = args_clone.get('messages', [])
-        model_id = args_clone.get('modelId', 'unknown')
-        model_parameters = {
-            **args_clone.get('inferenceConfig', {}),
-            **args_clone.get('additionalModelRequestFields', {})
-        }
-        
-        # Update Langfuse generation with input metadata
-        langfuse_context.update_current_observation(
-            input=messages,
-            model=model_id,
-            model_parameters=model_parameters,
-            metadata={
-                'system': args_clone.get('system', []),
-                'toolConfig': args_clone.get('toolConfig', {}),
-                'stream': stream
-            }
-        )
-        if DEBUG:
-            logger.info(f"Langfuse: Updated observation with input - model={model_id}, stream={stream}, messages_count={len(messages)}")
-
         try:
             if stream:
                 # Run the blocking boto3 call in a thread pool
@@ -291,93 +267,118 @@ async def _invoke_bedrock(self, chat_request: ChatRequest, stream=False):
             else:
                 # Run the blocking boto3 call in a thread pool
                 response = await run_in_threadpool(bedrock_runtime.converse, **args)
-                
-                # For non-streaming, extract response metadata immediately
-                if response and not stream:
-                    output_message = response.get("output", {}).get("message", {})
-                    usage = response.get("usage", {})
-                    
-                    # Build metadata
-                    metadata = {
-                        "stopReason": response.get("stopReason"),
-                        "ResponseMetadata": response.get("ResponseMetadata", {})
-                    }
-                    
-                    # Check for reasoning content in response
-                    has_reasoning = False
-                    reasoning_text = ""
-                    if output_message and "content" in output_message:
-                        for content_block in output_message.get("content", []):
-                            if "reasoningContent" in content_block:
-                                has_reasoning = True
-                                reasoning_text = content_block.get("reasoningContent", {}).get("reasoningText", {}).get("text", "")
-                                break
-                    
-                    if has_reasoning and reasoning_text:
-                        metadata["has_extended_thinking"] = True
-                        metadata["reasoning_content"] = reasoning_text
-                        metadata["reasoning_tokens_estimate"] = len(reasoning_text) // 4
-                    
-                    langfuse_context.update_current_observation(
-                        output=output_message,
-                        usage={
-                            "input": usage.get("inputTokens", 0),
-                            "output": usage.get("outputTokens", 0),
-                            "total": usage.get("totalTokens", 0)
-                        },
-                        metadata=metadata
-                    )
-                    if DEBUG:
-                        logger.info(f"Langfuse: Updated observation with output - "
-                                  f"input_tokens={usage.get('inputTokens', 0)}, "
-                                  f"output_tokens={usage.get('outputTokens', 0)}, "
-                                  f"has_reasoning={has_reasoning}, "
-                                  f"stop_reason={response.get('stopReason')}")
         except bedrock_runtime.exceptions.ValidationException as e:
             error_message = f"Bedrock validation error for model {chat_request.model}: {str(e)}"
             logger.error(error_message)
-            langfuse_context.update_current_observation(level="ERROR", status_message=error_message)
-            if DEBUG:
-                logger.info("Langfuse: Updated observation with ValidationException error")
             raise HTTPException(status_code=400, detail=str(e))
         except bedrock_runtime.exceptions.ThrottlingException as e:
             error_message = f"Bedrock throttling for model {chat_request.model}: {str(e)}"
             logger.warning(error_message)
-            langfuse_context.update_current_observation(level="WARNING", status_message=error_message)
-            if DEBUG:
-                logger.info("Langfuse: Updated observation with ThrottlingException warning")
             raise HTTPException(status_code=429, detail=str(e))
         except Exception as e:
             error_message = f"Bedrock invocation failed for model {chat_request.model}: {str(e)}"
             logger.error(error_message)
-            langfuse_context.update_current_observation(level="ERROR", status_message=error_message)
-            if DEBUG:
-                logger.info("Langfuse: Updated observation with generic Exception error")
             raise HTTPException(status_code=500, detail=str(e))
         return response
 
     async def chat(self, chat_request: ChatRequest) -> ChatResponse:
         """Default implementation for Chat API."""
-
+        from langfuse import get_client
+        
+        langfuse = get_client()
         message_id = self.generate_message_id()
-        response = await self._invoke_bedrock(chat_request)
-
-        output_message = response["output"]["message"]
-        input_tokens = response["usage"]["inputTokens"]
-        output_tokens = response["usage"]["outputTokens"]
-        finish_reason = response["stopReason"]
-
-        chat_response = self._create_response(
-            model=chat_request.model,
-            message_id=message_id,
-            content=output_message["content"],
-            finish_reason=finish_reason,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
-        )
-        if DEBUG:
-            logger.info("Proxy response :" + chat_response.model_dump_json())
-        return chat_response
+        
+        # Create Langfuse observation for Bedrock call
+        with langfuse.start_as_current_observation(as_type="generation", name="Bedrock Converse") as obs:
+            # Parse request for metadata
+            args = self._parse_request(chat_request)
+            messages = args.get('messages', [])
+            model_id = args.get('modelId', 'unknown')
+            model_parameters = {
+                **args.get('inferenceConfig', {}),
+                **args.get('additionalModelRequestFields', {})
+            }
+            
+            # Update observation with input
+            langfuse_context.update_current_observation(
+                input=messages,
+                model=model_id,
+                model_parameters=model_parameters,
+                metadata={
+                    'system': args.get('system', []),
+                    'toolConfig': args.get('toolConfig', {}),
+                    'stream': False
+                }
+            )
+            if DEBUG:
+                logger.info(f"Langfuse: Updated observation with input - model={model_id}, stream=False, messages_count={len(messages)}")
+            
+            try:
+                response = await self._invoke_bedrock(chat_request)
+                
+                output_message = response["output"]["message"]
+                input_tokens = response["usage"]["inputTokens"]
+                output_tokens = response["usage"]["outputTokens"]
+                finish_reason = response["stopReason"]
+                
+                # Update observation with output
+                metadata = {
+                    "stopReason": finish_reason,
+                    "ResponseMetadata": response.get("ResponseMetadata", {})
+                }
+                
+                # Check for reasoning content in response
+                has_reasoning = False
+                reasoning_text = ""
+                if output_message and "content" in output_message:
+                    for content_block in output_message.get("content", []):
+                        if "reasoningContent" in content_block:
+                            has_reasoning = True
+                            reasoning_text = content_block.get("reasoningContent", {}).get("reasoningText", {}).get("text", "")
+                            break
+                
+                if has_reasoning and reasoning_text:
+                    metadata["has_extended_thinking"] = True
+                    metadata["reasoning_content"] = reasoning_text
+                    metadata["reasoning_tokens_estimate"] = len(reasoning_text) // 4
+                
+                langfuse_context.update_current_observation(
+                    output=output_message,
+                    usage={
+                        "input": input_tokens,
+                        "output": output_tokens,
+                        "total": input_tokens + output_tokens
+                    },
+                    metadata=metadata
+                )
+                if DEBUG:
+                    logger.info(f"Langfuse: Updated observation with output - "
+                              f"input_tokens={input_tokens}, "
+                              f"output_tokens={output_tokens}, "
+                              f"has_reasoning={has_reasoning}, "
+                              f"stop_reason={finish_reason}")
+                
+                chat_response = self._create_response(
+                    model=chat_request.model,
+                    message_id=message_id,
+                    content=output_message["content"],
+                    finish_reason=finish_reason,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                )
+                if DEBUG:
+                    logger.info("Proxy response :" + chat_response.model_dump_json())
+                return chat_response
+            except HTTPException as e:
+                langfuse_context.update_current_observation(level="ERROR", status_message=str(e.detail))
+                if DEBUG:
+                    logger.info(f"Langfuse: Updated observation with error - {str(e.detail)[:100]}")
+                raise
+            except Exception as e:
+                langfuse_context.update_current_observation(level="ERROR", status_message=str(e))
+                if DEBUG:
+                    logger.info(f"Langfuse: Updated observation with error - {str(e)[:100]}")
+                raise
 
     async def _async_iterate(self, stream):
         """Helper method to convert sync iterator to async iterator"""
@@ -386,10 +387,21 @@ async def _async_iterate(self, stream):
             yield chunk
 
     async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
-        """Default implementation for Chat Stream API"""
+        """Default implementation for Chat Stream API
+        
+        Note: For streaming, we work within the parent trace context created by @observe
+        decorator on chat_completions endpoint. We update that trace context with 
+        streaming data as it arrives.
+        """
         try:
             if DEBUG:
                 logger.info(f"Langfuse: Starting streaming request for model={chat_request.model}")
+            
+            # Parse request for metadata to log in parent trace
+            args = self._parse_request(chat_request)
+            messages = args.get('messages', [])
+            model_id = args.get('modelId', 'unknown')
+            
             response = await self._invoke_bedrock(chat_request, stream=True)
             message_id = self.generate_message_id()
             stream = response.get("stream")
@@ -403,8 +415,8 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
             has_reasoning = False
 
             async for chunk in self._async_iterate(stream):
-                args = {"model_id": chat_request.model, "message_id": message_id, "chunk": chunk}
-                stream_response = self._create_response_stream(**args)
+                args_chunk = {"model_id": chat_request.model, "message_id": message_id, "chunk": chunk}
+                stream_response = self._create_response_stream(**args_chunk)
                 if not stream_response:
                     continue
 
@@ -438,49 +450,46 @@ 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 (both observation and trace)
+            # Update Langfuse trace with final streaming output
+            # This updates the parent trace from chat_completions
             if final_usage or accumulated_output:
-                update_params = {}
-                if accumulated_output:
-                    final_output = "".join(accumulated_output)
-                    update_params["output"] = final_output
-                if final_usage:
-                    update_params["usage"] = {
-                        "input": final_usage.prompt_tokens,
-                        "output": final_usage.completion_tokens,
-                        "total": final_usage.total_tokens
-                    }
-                # Build metadata
-                metadata = {}
-                if finish_reason:
-                    metadata["finish_reason"] = finish_reason
-                if has_reasoning and accumulated_reasoning:
-                    reasoning_text = "".join(accumulated_reasoning)
-                    metadata["has_extended_thinking"] = True
-                    metadata["reasoning_content"] = reasoning_text
-                    # Estimate reasoning tokens (rough approximation: ~4 chars per token)
-                    metadata["reasoning_tokens_estimate"] = len(reasoning_text) // 4
-                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
+                final_output = "".join(accumulated_output) if accumulated_output else None
                 trace_output = {
                     "message": {
                         "role": "assistant",
-                        "content": final_output if accumulated_output else None,
+                        "content": final_output,
                     },
                     "finish_reason": finish_reason,
                 }
-                langfuse_context.update_current_trace(output=trace_output)
+                
+                # Build metadata including usage info
+                trace_metadata = {
+                    "model": model_id,
+                    "stream": True,
+                }
+                if finish_reason:
+                    trace_metadata["finish_reason"] = finish_reason
+                if final_usage:
+                    trace_metadata["usage"] = {
+                        "prompt_tokens": final_usage.prompt_tokens,
+                        "completion_tokens": final_usage.completion_tokens,
+                        "total_tokens": final_usage.total_tokens
+                    }
+                if has_reasoning and accumulated_reasoning:
+                    reasoning_text = "".join(accumulated_reasoning)
+                    trace_metadata["has_extended_thinking"] = True
+                    trace_metadata["reasoning_tokens_estimate"] = len(reasoning_text) // 4
+                
+                langfuse_context.update_current_trace(
+                    output=trace_output,
+                    metadata=trace_metadata
+                )
 
                 if DEBUG:
                     output_length = len(accumulated_output)
-                    logger.info(f"Langfuse: Updated observation and trace with streaming output - "
+                    logger.info(f"Langfuse: Updated trace with streaming output - "
                               f"chunks_count={output_length}, "
-                              f"output_chars={len(final_output) if accumulated_output else 0}, "
+                              f"output_chars={len(final_output) if final_output else 0}, "
                               f"input_tokens={final_usage.prompt_tokens if final_usage else 'N/A'}, "
                               f"output_tokens={final_usage.completion_tokens if final_usage else 'N/A'}, "
                               f"has_reasoning={has_reasoning}, "
@@ -490,21 +499,17 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
             yield self.stream_response_to_bytes()
             self.think_emitted = False  # Cleanup
         except HTTPException:
-            # HTTPException already has Langfuse updated in _invoke_bedrock, re-raise it
+            # Re-raise HTTPException as-is
             raise
         except Exception as e:
             logger.error("Stream error for model %s: %s", chat_request.model, str(e))
-            # Update Langfuse with error (both observation and trace)
-            langfuse_context.update_current_observation(
-                level="ERROR",
-                status_message=f"Stream error: {str(e)}"
-            )
+            # Update Langfuse with error
             langfuse_context.update_current_trace(
                 output={"error": str(e)},
-                metadata={"error": True}
+                metadata={"error": True, "error_type": type(e).__name__}
             )
             if DEBUG:
-                logger.info(f"Langfuse: Updated observation with streaming error - error={str(e)[:100]}")
+                logger.info(f"Langfuse: Updated trace with streaming error - error={str(e)[:100]}")
             error_event = Error(error=ErrorMessage(message=str(e)))
             yield self.stream_response_to_bytes(error_event)