Format callbacks.py

Author: Twixes

posthog/ai/langchain/callbacks.py

@@ -1,7 +1,9 @@
 try:
     import langchain  # noqa: F401
 except ImportError:
-    raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
+    raise ModuleNotFoundError(
+        "Please install LangChain to use this feature: 'pip install langchain'"
+    )
 
 import logging
 import time
@@ -19,7 +21,14 @@
 from uuid import UUID
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain_core.messages import AIMessage, BaseMessage, FunctionMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
@@ -111,7 +120,9 @@ def on_chat_model_start(
         **kwargs,
     ):
         self._set_parent_of_run(run_id, parent_run_id)
-        input = [_convert_message_to_dict(message) for row in messages for message in row]
+        input = [
+            _convert_message_to_dict(message) for row in messages for message in row
+        ]
         self._set_run_metadata(serialized, run_id, input, **kwargs)
 
     def on_llm_start(
@@ -161,17 +172,24 @@ def on_llm_end(
         generation_result = response.generations[-1]
         if isinstance(generation_result[-1], ChatGeneration):
             output = [
-                _convert_message_to_dict(cast(ChatGeneration, generation).message) for generation in generation_result
+                _convert_message_to_dict(cast(ChatGeneration, generation).message)
+                for generation in generation_result
             ]
         else:
-            output = [_extract_raw_esponse(generation) for generation in generation_result]
+            output = [
+                _extract_raw_esponse(generation) for generation in generation_result
+            ]
 
         event_properties = {
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
-            "$ai_output_choices": with_privacy_mode(self._client, self._privacy_mode, output),
+            "$ai_input": with_privacy_mode(
+                self._client, self._privacy_mode, run.get("messages")
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._client, self._privacy_mode, output
+            ),
             "$ai_http_status": 200,
             "$ai_input_tokens": input_tokens,
             "$ai_output_tokens": output_tokens,
@@ -219,7 +237,9 @@ def on_llm_error(
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
+            "$ai_input": with_privacy_mode(
+                self._client, self._privacy_mode, run.get("messages")
+            ),
             "$ai_http_status": _get_http_status(error),
             "$ai_latency": latency,
             "$ai_trace_id": trace_id,
@@ -339,7 +359,9 @@ def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
     return message_dict
 
 
-def _parse_usage_model(usage: Union[BaseModel, Dict]) -> Tuple[Union[int, None], Union[int, None]]:
+def _parse_usage_model(
+    usage: Union[BaseModel, Dict],
+) -> Tuple[Union[int, None], Union[int, None]]:
     if isinstance(usage, BaseModel):
         usage = usage.__dict__
 
@@ -363,7 +385,9 @@ def _parse_usage_model(usage: Union[BaseModel, Dict]) -> Tuple[Union[int, None],
         if model_key in usage:
             captured_count = usage[model_key]
             final_count = (
-                sum(captured_count) if isinstance(captured_count, list) else captured_count
+                sum(captured_count)
+                if isinstance(captured_count, list)
+                else captured_count
             )  # For Bedrock, the token count is a list when streamed
 
             parsed_usage[type_key] = final_count
@@ -384,8 +408,12 @@ def _parse_usage(response: LLMResult):
     if hasattr(response, "generations"):
         for generation in response.generations:
             for generation_chunk in generation:
-                if generation_chunk.generation_info and ("usage_metadata" in generation_chunk.generation_info):
-                    llm_usage = _parse_usage_model(generation_chunk.generation_info["usage_metadata"])
+                if generation_chunk.generation_info and (
+                    "usage_metadata" in generation_chunk.generation_info
+                ):
+                    llm_usage = _parse_usage_model(
+                        generation_chunk.generation_info["usage_metadata"]
+                    )
                     break
 
                 message_chunk = getattr(generation_chunk, "message", {})
@@ -397,13 +425,19 @@ def _parse_usage(response: LLMResult):
                     else None
                 )
                 bedrock_titan_usage = (
-                    response_metadata.get("amazon-bedrock-invocationMetrics", None)  # for Bedrock-Titan
+                    response_metadata.get(
+                        "amazon-bedrock-invocationMetrics", None
+                    )  # for Bedrock-Titan
                     if isinstance(response_metadata, dict)
                     else None
                 )
-                ollama_usage = getattr(message_chunk, "usage_metadata", None)  # for Ollama
+                ollama_usage = getattr(
+                    message_chunk, "usage_metadata", None
+                )  # for Ollama
 
-                chunk_usage = bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                chunk_usage = (
+                    bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                )
                 if chunk_usage:
                     llm_usage = _parse_usage_model(chunk_usage)
                     break