fix(llma): fix test

Author: Radu-Raicea

posthog/ai/anthropic/anthropic_async.py

@@ -136,7 +136,7 @@ async def _create_streaming(
         content_blocks: List[StreamingContentBlock] = []
         tools_in_progress: Dict[str, ToolInProgress] = {}
         current_text_block: Optional[StreamingContentBlock] = None
-        response = super().create(**kwargs)
+        response = await super().create(**kwargs)
 
         async def generator():
             nonlocal usage_stats

posthog/ai/gemini/gemini_converter.py

@@ -303,8 +303,8 @@ def extract_gemini_usage_from_chunk(chunk: Any) -> StreamingUsageStats:
     usage["input_tokens"] = getattr(chunk.usage_metadata, "prompt_token_count", 0)
     usage["output_tokens"] = getattr(chunk.usage_metadata, "candidates_token_count", 0)
 
-    # Calculate total if both are present
-    if usage["input_tokens"] and usage["output_tokens"]:
+    # Calculate total if both values are defined (including 0)
+    if "input_tokens" in usage and "output_tokens" in usage:
         usage["total_tokens"] = usage["input_tokens"] + usage["output_tokens"]
 
     return usage

posthog/ai/utils.py

@@ -503,8 +503,6 @@ def capture_streaming_event(
         ph_client: PostHog client instance
         event_data: Standardized streaming event data containing all necessary information
     """
-    import uuid
-
     trace_id = event_data.get("trace_id") or str(uuid.uuid4())
 
     # Build base event properties

posthog/test/ai/anthropic/test_anthropic.py

@@ -480,11 +480,14 @@ async def mock_async_stream():
         )
         yield final_msg
 
-    # Mock create to return the async generator directly (not wrapped in a coroutine)
-    # This matches the actual behavior when stream=True
+    # Mock create to return a coroutine that yields the async generator
+    # This matches the actual behavior when stream=True with await
+    async def async_create_wrapper(**kwargs):
+        return mock_async_stream()
+    
     with patch(
         "anthropic.resources.messages.AsyncMessages.create",
-        return_value=mock_async_stream(),
+        side_effect=async_create_wrapper,
     ):
         client = AsyncAnthropic(posthog_client=mock_client)
         response = await client.messages.create(
@@ -926,14 +929,18 @@ def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with
     """Test that tool calls are properly captured in async streaming mode."""
     import asyncio
 
-    async def mock_async_create(**kwargs):
+    async def mock_async_generator():
         # Convert regular generator to async generator
         for event in mock_anthropic_stream_with_tools:
             yield event
+    
+    async def mock_async_create(**kwargs):
+        # Return the async generator (to be awaited by the implementation)
+        return mock_async_generator()
 
     with patch(
         "anthropic.resources.AsyncMessages.create",
-        side_effect=lambda **kwargs: mock_async_create(),
+        side_effect=mock_async_create,
     ):
         async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)