chore: reorganize imports

llm_observability_examples.py

@@ -2,7 +2,7 @@
 import uuid
 
 import posthog
-from posthog.ai import AsyncOpenAI, OpenAI
+from posthog.ai.openai import AsyncOpenAI, OpenAI
 
 # Example credentials - replace these with your own or use environment variables
 posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api-key")

posthog/ai/__init__.py

@@ -1,4 +0,0 @@
-from .providers.openai.openai import OpenAI
-from .providers.openai.openai_async import AsyncOpenAI
-
-__all__ = ["OpenAI", "AsyncOpenAI"]

posthog/ai/langchain/__init__.py

@@ -0,0 +1,3 @@
+from .callbacks import PosthogCallbackHandler
+
+__all__ = ["PosthogCallbackHandler"]

posthog/ai/providers/langchain.py → posthog/ai/langchain/callbacks.py

@@ -1,5 +1,5 @@
 try:
-    import langchain
+    import langchain  # noqa: F401
 except ImportError:
     raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
 
@@ -42,6 +42,10 @@ class RunMetadata(TypedDict, total=False):
 
 
 class PosthogCallbackHandler(BaseCallbackHandler):
+    """
+    A callback handler for LangChain that sends events to PostHog LLM Observability.
+    """
+
     _client: Client
     """PostHog client instance."""
     _distinct_id: Optional[str]
@@ -371,20 +375,19 @@ def _parse_usage(response: LLMResult):
                 message_chunk = getattr(generation_chunk, "message", {})
                 response_metadata = getattr(message_chunk, "response_metadata", {})
 
-                chunk_usage = (
-                    (
-                        response_metadata.get("usage", None)  # for Bedrock-Anthropic
-                        if isinstance(response_metadata, dict)
-                        else None
-                    )
-                    or (
-                        response_metadata.get("amazon-bedrock-invocationMetrics", None)  # for Bedrock-Titan
-                        if isinstance(response_metadata, dict)
-                        else None
-                    )
-                    or getattr(message_chunk, "usage_metadata", None)  # for Ollama
+                bedrock_anthropic_usage = (
+                    response_metadata.get("usage", None)  # for Bedrock-Anthropic
+                    if isinstance(response_metadata, dict)
+                    else None
+                )
+                bedrock_titan_usage = (
+                    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
 
+                chunk_usage = bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
                 if chunk_usage:
                     llm_usage = _parse_usage_model(chunk_usage)
                     break

posthog/ai/openai/__init__.py

@@ -0,0 +1,4 @@
+from .openai import OpenAI
+from .openai_async import AsyncOpenAI
+
+__all__ = ["OpenAI", "AsyncOpenAI"]

posthog/ai/providers/openai/openai.py → posthog/ai/openai/openai.py

@@ -2,13 +2,13 @@
 import uuid
 from typing import Any, Dict, Optional
 
-import openai.resources
-
 try:
     import openai
 except ImportError:
     raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
 
+import openai.resources
+
 from posthog.ai.utils import call_llm_and_track_usage, get_model_params
 from posthog.client import Client as PostHogClient
 

posthog/ai/providers/openai/openai_async.py → posthog/ai/openai/openai_async.py

@@ -2,13 +2,13 @@
 import uuid
 from typing import Any, Dict, Optional
 
-import openai.resources
-
 try:
     import openai
 except ImportError:
     raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
 
+import openai.resources
+
 from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params
 from posthog.client import Client as PostHogClient
 

posthog/test/ai/providers/test_langchain.py → posthog/test/ai/langchain/test_callbacks.py

@@ -12,7 +12,7 @@
 from langchain_openai.chat_models import ChatOpenAI
 import pytest
 
-from posthog.ai.providers.langchain import PosthogCallbackHandler
+from posthog.ai.langchain import PosthogCallbackHandler
 
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")