feat: add langfuse observability to bedrock gateway

Author: jake-marsden

src/api/models/bedrock.py

@@ -1,6 +1,7 @@
 import base64
 import json
 import logging
+import os
 import re
 import time
 from abc import ABC
@@ -13,6 +14,7 @@
 import tiktoken
 from botocore.config import Config
 from fastapi import HTTPException
+from langfuse import observe
 from starlette.concurrency import run_in_threadpool
 
 from api.models.base import BaseChatModel, BaseEmbeddingsModel
@@ -230,6 +232,7 @@ 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:
@@ -240,6 +243,30 @@ async def _invoke_bedrock(self, chat_request: ChatRequest, stream=False):
         if DEBUG:
             logger.info("Bedrock request: " + json.dumps(str(args)))
 
+        # Import langfuse to update current generation (must be done inside the decorated function)
+        from langfuse import langfuse_context
+        
+        # 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
+            }
+        )
+
         try:
             if stream:
                 # Run the blocking boto3 call in a thread pool
@@ -249,14 +276,38 @@ 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", {})
+                    
+                    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={
+                            "stopReason": response.get("stopReason"),
+                            "ResponseMetadata": response.get("ResponseMetadata", {})
+                        }
+                    )
         except bedrock_runtime.exceptions.ValidationException as e:
-            logger.error("Bedrock validation error for model %s: %s", chat_request.model, str(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)
             raise HTTPException(status_code=400, detail=str(e))
         except bedrock_runtime.exceptions.ThrottlingException as e:
-            logger.warning("Bedrock throttling for model %s: %s", chat_request.model, str(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)
             raise HTTPException(status_code=429, detail=str(e))
         except Exception as e:
-            logger.error("Bedrock invocation failed for model %s: %s", chat_request.model, str(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)
             raise HTTPException(status_code=500, detail=str(e))
         return response
 

src/requirements.txt

@@ -7,3 +7,4 @@ requests==2.32.4
 numpy==2.2.5
 boto3==1.40.4
 botocore==1.40.4
+langfuse