ai-station/.venv/lib/python3.12/site-packages/posthog/ai/anthropic/anthropic_providers.py

66 lines
2.1 KiB
Python

try:
import anthropic
except ImportError:
raise ModuleNotFoundError(
"Please install the Anthropic SDK to use this feature: 'pip install anthropic'"
)
from typing import Optional
from posthog.ai.anthropic.anthropic import WrappedMessages
from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages
from posthog.client import Client as PostHogClient
from posthog import setup
class AnthropicBedrock(anthropic.AnthropicBedrock):
"""
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
"""
_ph_client: PostHogClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self.messages = WrappedMessages(self)
class AsyncAnthropicBedrock(anthropic.AsyncAnthropicBedrock):
"""
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
"""
_ph_client: PostHogClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self.messages = AsyncWrappedMessages(self)
class AnthropicVertex(anthropic.AnthropicVertex):
"""
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
"""
_ph_client: PostHogClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self.messages = WrappedMessages(self)
class AsyncAnthropicVertex(anthropic.AsyncAnthropicVertex):
"""
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
"""
_ph_client: PostHogClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self.messages = AsyncWrappedMessages(self)