redact api_key and fix cost calculation for custom providers

Author: alrocar

examples/litellm_async.py

@@ -21,7 +21,12 @@ async def main():
         messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
         stream=True,
         user="test_user",
-        metadata={"organization": "tinybird", "environment": "dev", "project": "litellm_test", "chat_id": "1234567890"}
+        metadata={
+            "organization": "tinybird",
+            "environment": "dev",
+            "project": "litellm_test",
+            "chat_id": "1234567890",
+        },
     )
 
     async for chunk in response:

examples/litellm_sync.py

@@ -18,7 +18,12 @@
     messages=[{"role": "user", "content": "Hi 👋 - i'm claude"}],
     stream=True,
     user="test_user",
-    metadata={"organization": "tinybird", "environment": "dev", "project": "litellm_test", "chat_id": "1234567890"}
+    metadata={
+        "organization": "tinybird",
+        "environment": "dev",
+        "project": "litellm_test",
+        "chat_id": "1234567890",
+    },
 )
 
 for chunk in response:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tinybird-python-sdk"
-version = "0.3.1"
+version = "0.3.2"
 description = "Python SDK for Tinybird"
 readme = "README.md"
 authors = [

tb/litellm/handler.py

@@ -48,6 +48,12 @@ def __init__(
         )
 
     def _extract_data(self, kwargs, response_obj, start_time, end_time):
+        api_key = kwargs.get("api_key")
+        if api_key and len(api_key) > 8:
+            api_key = api_key[:4] + "****" + api_key[-4:]
+        else:
+            api_key = "****"
+
         data = {
             "model": kwargs.get("model"),
             "messages": kwargs.get("messages"),
@@ -58,77 +64,80 @@ def _extract_data(self, kwargs, response_obj, start_time, end_time):
             "stream": kwargs.get("stream", False),
             "call_type": kwargs.get("call_type", "completion"),
             "provider": kwargs.get("custom_llm_provider"),
-            "api_key": kwargs.get("api_key"),
             "log_event_type": kwargs.get("log_event_type"),
             "llm_api_duration_ms": kwargs.get("llm_api_duration_ms"),
             "response_headers": kwargs.get("response_headers", {}),
             "cache_hit": kwargs.get("cache_hit", False),
-            "standard_logging_object_id": kwargs.get("standard_logging_object", {}).get("id"),
-            "standard_logging_object_status": kwargs.get("standard_logging_object", {}).get("status"),
-            "standard_logging_object_response_time": kwargs.get("standard_logging_object", {}).get
-            ("response_time"),
-            "standard_logging_object_saved_cache_cost": kwargs.get("standard_logging_object", {}).get("saved_cache_cost"),
-            "standard_logging_object_hidden_params": kwargs.get("standard_logging_object", {}).get("status"),
+            "standard_logging_object_id": kwargs.get("standard_logging_object", {}).get(
+                "id"
+            ),
+            "standard_logging_object_status": kwargs.get(
+                "standard_logging_object", {}
+            ).get("status"),
+            "standard_logging_object_response_time": kwargs.get(
+                "standard_logging_object", {}
+            ).get("response_time"),
+            "standard_logging_object_saved_cache_cost": kwargs.get(
+                "standard_logging_object", {}
+            ).get("saved_cache_cost"),
+            "standard_logging_object_hidden_params": kwargs.get(
+                "standard_logging_object", {}
+            ).get("status"),
+            "api_key": api_key,
         }
-        
+
         # response = litellm.completion(model="gpt-3.5-turbo", messages=messages, metadata={"hello": "world"})
         litellm_params = kwargs.get("litellm_params", {})
         data["proxy_metadata"] = litellm_params.get("metadata", {})
         data["response"] = response_obj.json()
-        data["cost"] = litellm.completion_cost(completion_response=response_obj)
+        data["cost"] = litellm.completion_cost(
+            completion_response=response_obj,
+            custom_llm_provider=kwargs.get("custom_llm_provider"),
+        )
 
         data["exception"] = kwargs.get("exception", None)
         data["traceback"] = kwargs.get("traceback_exception", None)
         if isinstance(start_time, datetime) and isinstance(end_time, datetime):
             duration = (end_time - start_time).total_seconds()
         else:
-            duration = (end_time - start_time)
+            duration = end_time - start_time
         data["duration"] = duration
         return safe_json_dumps(data)
-    
+
+
 class TinybirdLitellmSyncHandler(TinybirdLitellmHandler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
         try:
             data = self._extract_data(kwargs, response_obj, start_time, end_time)
-            self.api.send(
-                f"events?name={self.datasource_name}",
-                data=data
-            )
+            self.api.send(f"events?name={self.datasource_name}", data=data)
         except Exception as e:
             logging.error(f"Error logging success event: {e}")
+
     def log_failure_event(self, kwargs, response_obj, start_time, end_time):
         try:
             data = self._extract_data(kwargs, response_obj, start_time, end_time)
-            self.api.send(
-                f"events?name={self.datasource_name}",
-                data=data
-            )
+            self.api.send(f"events?name={self.datasource_name}", data=data)
         except Exception as e:
             logging.error(f"Error logging failure event: {e}")
 
+
 class TinybirdLitellmAsyncHandler(TinybirdLitellmHandler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         try:
             data = self._extract_data(kwargs, response_obj, start_time, end_time)
-            await self.async_api.send(
-                f"events?name={self.datasource_name}",
-                data=data
-            )
+            await self.async_api.send(f"events?name={self.datasource_name}", data=data)
         except Exception as e:
             logging.error(f"Error logging success event: {e}")
 
     async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
         try:
             data = self._extract_data(kwargs, response_obj, start_time, end_time)
-            await self.async_api.send(
-                f"events?name={self.datasource_name}",
-                data=data
-            )
+            await self.async_api.send(f"events?name={self.datasource_name}", data=data)
         except Exception as e:
             logging.error(f"Error logging failure event: {e}")