fix(llma): fix types

Author: Radu-Raicea

posthog/ai/anthropic/anthropic.py

@@ -10,8 +10,10 @@
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthog.ai.types import StreamingContentBlock, ToolInProgress
 from posthog.ai.utils import (
     call_llm_and_track_usage,
+    merge_usage_stats,
 )
 from posthog.ai.anthropic.anthropic_converter import (
     extract_anthropic_usage_from_event,
@@ -126,9 +128,9 @@ def _create_streaming(
         start_time = time.time()
         usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
         accumulated_content = ""
-        content_blocks: List[Dict[str, Any]] = []
-        tools_in_progress: Dict[str, Dict[str, Any]] = {}
-        current_text_block: Optional[Dict[str, Any]] = None
+        content_blocks: List[StreamingContentBlock] = []
+        tools_in_progress: Dict[str, ToolInProgress] = {}
+        current_text_block: Optional[StreamingContentBlock] = None
         response = super().create(**kwargs)
 
         def generator():
@@ -142,7 +144,7 @@ def generator():
                 for event in response:
                     # Extract usage stats from event
                     event_usage = extract_anthropic_usage_from_event(event)
-                    usage_stats.update(event_usage)
+                    merge_usage_stats(usage_stats, event_usage)
 
                     # Handle content block start events
                     if hasattr(event, "type") and event.type == "content_block_start":
@@ -157,7 +159,9 @@ def generator():
                                 current_text_block = None
 
                         if tool:
-                            tools_in_progress[tool["block"]["id"]] = tool
+                            tool_id = tool["block"].get("id")
+                            if tool_id:
+                                tools_in_progress[tool_id] = tool
 
                     # Handle text delta events
                     delta_text = handle_anthropic_text_delta(event, current_text_block)
@@ -208,7 +212,7 @@ def _capture_streaming_event(
         kwargs: Dict[str, Any],
         usage_stats: Dict[str, int],
         latency: float,
-        content_blocks: List[Dict[str, Any]],
+        content_blocks: List[StreamingContentBlock],
         accumulated_content: str,
     ):
         from posthog.ai.types import StreamingEventData
@@ -225,7 +229,7 @@ def _capture_streaming_event(
 
         event_data = StreamingEventData(
             provider="anthropic",
-            model=kwargs.get("model"),
+            model=kwargs.get("model", "unknown"),
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,

posthog/ai/anthropic/anthropic_async.py

@@ -11,11 +11,13 @@
 from typing import Any, Dict, List, Optional
 
 from posthog import setup
+from posthog.ai.types import StreamingContentBlock, ToolInProgress
 from posthog.ai.utils import (
     call_llm_and_track_usage_async,
     extract_available_tool_calls,
     get_model_params,
     merge_system_prompt,
+    merge_usage_stats,
     with_privacy_mode,
 )
 from posthog.ai.anthropic.anthropic_converter import (
@@ -131,9 +133,9 @@ async def _create_streaming(
         start_time = time.time()
         usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
         accumulated_content = ""
-        content_blocks: List[Dict[str, Any]] = []
-        tools_in_progress: Dict[str, Dict[str, Any]] = {}
-        current_text_block: Optional[Dict[str, Any]] = None
+        content_blocks: List[StreamingContentBlock] = []
+        tools_in_progress: Dict[str, ToolInProgress] = {}
+        current_text_block: Optional[StreamingContentBlock] = None
         response = super().create(**kwargs)
 
         async def generator():
@@ -147,7 +149,7 @@ async def generator():
                 async for event in response:
                     # Extract usage stats from event
                     event_usage = extract_anthropic_usage_from_event(event)
-                    usage_stats.update(event_usage)
+                    merge_usage_stats(usage_stats, event_usage)
 
                     # Handle content block start events
                     if hasattr(event, "type") and event.type == "content_block_start":
@@ -162,7 +164,9 @@ async def generator():
                                 current_text_block = None
 
                         if tool:
-                            tools_in_progress[tool["block"]["id"]] = tool
+                            tool_id = tool["block"].get("id")
+                            if tool_id:
+                                tools_in_progress[tool_id] = tool
 
                     # Handle text delta events
                     delta_text = handle_anthropic_text_delta(event, current_text_block)
@@ -213,7 +217,7 @@ async def _capture_streaming_event(
         kwargs: Dict[str, Any],
         usage_stats: Dict[str, int],
         latency: float,
-        content_blocks: List[Dict[str, Any]],
+        content_blocks: List[StreamingContentBlock],
         accumulated_content: str,
     ):
         if posthog_trace_id is None:

posthog/ai/anthropic/anthropic_converter.py

@@ -31,7 +31,7 @@ def format_anthropic_response(response: Any) -> List[FormattedMessage]:
         List of formatted messages with role and content
     """
 
-    output = []
+    output: List[FormattedMessage] = []
 
     if response is None:
         return output
@@ -127,7 +127,7 @@ def extract_anthropic_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
 
 
 def format_anthropic_streaming_content(
-    content_blocks: List[Dict[str, Any]],
+    content_blocks: List[StreamingContentBlock],
 ) -> List[FormattedContentItem]:
     """
     Format content blocks from Anthropic streaming response.
@@ -145,19 +145,17 @@ def format_anthropic_streaming_content(
 
     for block in content_blocks:
         if block.get("type") == "text":
-            text_content: FormattedTextContent = {
+            formatted.append({
                 "type": "text",
-                "text": block.get("text", ""),
-            }
-            formatted.append(text_content)
+                "text": block.get("text") or "",
+            })
 
         elif block.get("type") == "function":
-            function_call: FormattedFunctionCall = {
+            formatted.append({
                 "type": "function",
                 "id": block.get("id"),
-                "function": block.get("function", {}),
-            }
-            formatted.append(function_call)
+                "function": block.get("function") or {},
+            })
 
     return formatted
 
@@ -222,13 +220,13 @@ def handle_anthropic_content_block_start(
         return content_block, None
 
     elif block.type == "tool_use":
-        content_block: StreamingContentBlock = {
+        tool_block: StreamingContentBlock = {
             "type": "function",
             "id": getattr(block, "id", ""),
             "function": {"name": getattr(block, "name", ""), "arguments": {}},
         }
-        tool_in_progress: ToolInProgress = {"block": content_block, "input_string": ""}
-        return content_block, tool_in_progress
+        tool_in_progress: ToolInProgress = {"block": tool_block, "input_string": ""}
+        return tool_block, tool_in_progress
 
     return None, None
 
@@ -251,7 +249,11 @@ def handle_anthropic_text_delta(
         delta_text = event.delta.text or ""
 
         if current_block is not None and current_block.get("type") == "text":
-            current_block["text"] = current_block.get("text", "") + delta_text
+            text_val = current_block.get("text")
+            if text_val is not None:
+                current_block["text"] = text_val + delta_text
+            else:
+                current_block["text"] = delta_text
 
         return delta_text
 

posthog/ai/gemini/gemini.py

@@ -14,6 +14,7 @@
 from posthog.ai.utils import (
     call_llm_and_track_usage,
     capture_streaming_event,
+    merge_usage_stats,
 )
 from posthog.ai.gemini.gemini_converter import (
     format_gemini_input,
@@ -308,7 +309,8 @@ def generator():
                     chunk_usage = extract_gemini_usage_from_chunk(chunk)
 
                     if chunk_usage:
-                        usage_stats.update(chunk_usage)
+                        # Gemini reports cumulative totals, not incremental values
+                        merge_usage_stats(usage_stats, chunk_usage, mode="cumulative")
 
                     # Extract content from chunk (now returns content blocks)
                     content_block = extract_gemini_content_from_chunk(chunk)

posthog/ai/gemini/gemini_converter.py

@@ -9,9 +9,7 @@
 
 from posthog.ai.types import (
     FormattedContentItem,
-    FormattedFunctionCall,
     FormattedMessage,
-    FormattedTextContent,
     StreamingUsageStats,
     TokenUsage,
 )
@@ -164,7 +162,7 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
         List of formatted messages with role and content
     """
 
-    output = []
+    output: List[FormattedMessage] = []
 
     if response is None:
         return output
@@ -177,43 +175,38 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
                 if hasattr(candidate.content, "parts") and candidate.content.parts:
                     for part in candidate.content.parts:
                         if hasattr(part, "text") and part.text:
-                            text_content: FormattedTextContent = {
+                            content.append({
                                 "type": "text",
                                 "text": part.text,
-                            }
-                            content.append(text_content)
+                            })
 
                         elif hasattr(part, "function_call") and part.function_call:
                             function_call = part.function_call
-                            func_content: FormattedFunctionCall = {
+                            content.append({
                                 "type": "function",
                                 "function": {
                                     "name": function_call.name,
                                     "arguments": function_call.args,
                                 },
-                            }
-                            content.append(func_content)
+                            })
 
                 if content:
-                    message: FormattedMessage = {
+                    output.append({
                         "role": "assistant",
                         "content": content,
-                    }
-                    output.append(message)
+                    })
 
             elif hasattr(candidate, "text") and candidate.text:
-                message: FormattedMessage = {
+                output.append({
                     "role": "assistant",
                     "content": [{"type": "text", "text": candidate.text}],
-                }
-                output.append(message)
+                })
 
     elif hasattr(response, "text") and response.text:
-        message: FormattedMessage = {
+        output.append({
             "role": "assistant",
             "content": [{"type": "text", "text": response.text}],
-        }
-        output.append(message)
+        })
 
     return output
 
@@ -258,7 +251,7 @@ def format_gemini_input(contents: Any) -> List[FormattedMessage]:
 
     # Handle list input
     if isinstance(contents, list):
-        formatted = []
+        formatted: List[FormattedMessage] = []
 
         for item in contents:
             if isinstance(item, str):
@@ -383,27 +376,24 @@ def format_gemini_streaming_output(
                 elif item.get("type") == "function":
                     # If we have accumulated text, add it first
                     if text_parts:
-                        text_content: FormattedTextContent = {
+                        content.append({
                             "type": "text",
                             "text": "".join(text_parts),
-                        }
-                        content.append(text_content)
+                        })
                         text_parts = []
 
                     # Add the function call
-                    func_content: FormattedFunctionCall = {
+                    content.append({
                         "type": "function",
                         "function": item.get("function", {}),
-                    }
-                    content.append(func_content)
+                    })
 
         # Add any remaining text
         if text_parts:
-            text_content: FormattedTextContent = {
+            content.append({
                 "type": "text",
                 "text": "".join(text_parts),
-            }
-            content.append(text_content)
+            })
 
         # If we have content, return it
         if content:

posthog/ai/openai/openai.py

@@ -12,6 +12,7 @@
 from posthog.ai.utils import (
     call_llm_and_track_usage,
     extract_available_tool_calls,
+    merge_usage_stats,
     with_privacy_mode,
 )
 from posthog.ai.openai.openai_converter import (
@@ -133,7 +134,7 @@ def generator():
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
                     if chunk_usage:
-                        usage_stats.update(chunk_usage)
+                        merge_usage_stats(usage_stats, chunk_usage)
 
                     # Extract content from chunk
                     content = extract_openai_content_from_chunk(chunk, "responses")
@@ -189,7 +190,7 @@ def _capture_streaming_event(
 
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model"),
+            model=kwargs.get("model", "unknown"),
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,
@@ -334,7 +335,7 @@ def generator():
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
 
                     if chunk_usage:
-                        usage_stats.update(chunk_usage)
+                        merge_usage_stats(usage_stats, chunk_usage)
 
                     # Extract content from chunk
                     content = extract_openai_content_from_chunk(chunk, "chat")
@@ -406,7 +407,7 @@ def _capture_streaming_event(
 
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model"),
+            model=kwargs.get("model", "unknown"),
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,

posthog/ai/openai/openai_async.py

@@ -14,6 +14,7 @@
     call_llm_and_track_usage_async,
     extract_available_tool_calls,
     get_model_params,
+    merge_usage_stats,
     with_privacy_mode,
 )
 from posthog.ai.openai.openai_converter import (
@@ -137,7 +138,7 @@ async def async_generator():
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
                     if chunk_usage:
-                        usage_stats.update(chunk_usage)
+                        merge_usage_stats(usage_stats, chunk_usage)
 
                     # Extract content from chunk
                     content = extract_openai_content_from_chunk(chunk, "responses")
@@ -354,7 +355,7 @@ async def async_generator():
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
                     if chunk_usage:
-                        usage_stats.update(chunk_usage)
+                        merge_usage_stats(usage_stats, chunk_usage)
 
                     # Extract content from chunk
                     content = extract_openai_content_from_chunk(chunk, "chat")