diff --git a/CHANGELOG.md b/CHANGELOG.md index dc7c91d8..fa5681a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +## 3.8.4 - 2025-01-17 + +1. Add Anthropic support for LLM Observability. +2. Update LLM Observability to use output_choices. + ## 3.8.3 - 2025-01-14 1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package. diff --git a/posthog/ai/anthropic/__init__.py b/posthog/ai/anthropic/__init__.py new file mode 100644 index 00000000..db4db080 --- /dev/null +++ b/posthog/ai/anthropic/__init__.py @@ -0,0 +1,12 @@ +from .anthropic import Anthropic +from .anthropic_async import AsyncAnthropic +from .anthropic_providers import AnthropicBedrock, AnthropicVertex, AsyncAnthropicBedrock, AsyncAnthropicVertex + +__all__ = [ + "Anthropic", + "AsyncAnthropic", + "AnthropicBedrock", + "AsyncAnthropicBedrock", + "AnthropicVertex", + "AsyncAnthropicVertex", +] diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py new file mode 100644 index 00000000..a6b12b8b --- /dev/null +++ b/posthog/ai/anthropic/anthropic.py @@ -0,0 +1,198 @@ +try: + import anthropic + from anthropic.resources import Messages +except ImportError: + raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'") + +import time +import uuid +from typing import Any, Dict, Optional + +from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode +from posthog.client import Client as PostHogClient + + +class Anthropic(anthropic.Anthropic): + """ + A wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog. + """ + + _ph_client: PostHogClient + + def __init__(self, posthog_client: PostHogClient, **kwargs): + """ + Args: + posthog_client: PostHog client for tracking usage + **kwargs: Additional arguments passed to the Anthropic client + """ + super().__init__(**kwargs) + self._ph_client = posthog_client + self.messages = WrappedMessages(self) + + +class WrappedMessages(Messages): + _client: Anthropic + + def create( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + """ + Create a message using Anthropic's API while tracking usage in PostHog. + + Args: + posthog_distinct_id: Optional ID to associate with the usage event + posthog_trace_id: Optional trace UUID for linking events + posthog_properties: Optional dictionary of extra properties to include in the event + posthog_privacy_mode: Whether to redact sensitive information in tracking + posthog_groups: Optional group analytics properties + **kwargs: Arguments passed to Anthropic's messages.create + """ + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + if kwargs.get("stream", False): + return self._create_streaming( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + + return call_llm_and_track_usage( + posthog_distinct_id, + self._client._ph_client, + "anthropic", + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + self._client.base_url, + super().create, + **kwargs, + ) + + def stream( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + return self._create_streaming( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + + def _create_streaming( + self, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + **kwargs: Any, + ): + start_time = time.time() + usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0} + accumulated_content = [] + response = super().create(**kwargs) + + def generator(): + nonlocal usage_stats + nonlocal accumulated_content + try: + for event in response: + if hasattr(event, "usage") and event.usage: + usage_stats = { + k: getattr(event.usage, k, 0) + for k in [ + "input_tokens", + "output_tokens", + ] + } + + if hasattr(event, "content") and event.content: + accumulated_content.append(event.content) + + yield event + + finally: + end_time = time.time() + latency = end_time - start_time + output = "".join(accumulated_content) + + self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + usage_stats, + latency, + output, + ) + + return generator() + + def _capture_streaming_event( + self, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + kwargs: Dict[str, Any], + usage_stats: Dict[str, int], + latency: float, + output: str, + ): + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + event_properties = { + "$ai_provider": "anthropic", + "$ai_model": kwargs.get("model"), + "$ai_model_parameters": get_model_params(kwargs), + "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")), + "$ai_output_choices": with_privacy_mode( + self._client._ph_client, + posthog_privacy_mode, + [{"content": output, "role": "assistant"}], + ), + "$ai_http_status": 200, + "$ai_input_tokens": usage_stats.get("input_tokens", 0), + "$ai_output_tokens": usage_stats.get("output_tokens", 0), + "$ai_latency": latency, + "$ai_trace_id": posthog_trace_id, + "$ai_base_url": str(self._client.base_url), + **(posthog_properties or {}), + } + + if posthog_distinct_id is None: + event_properties["$process_person_profile"] = False + + if hasattr(self._client._ph_client, "capture"): + self._client._ph_client.capture( + distinct_id=posthog_distinct_id or posthog_trace_id, + event="$ai_generation", + properties=event_properties, + groups=posthog_groups, + ) diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py new file mode 100644 index 00000000..31a3b355 --- /dev/null +++ b/posthog/ai/anthropic/anthropic_async.py @@ -0,0 +1,198 @@ +try: + import anthropic + from anthropic.resources import AsyncMessages +except ImportError: + raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'") + +import time +import uuid +from typing import Any, Dict, Optional + +from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, with_privacy_mode +from posthog.client import Client as PostHogClient + + +class AsyncAnthropic(anthropic.AsyncAnthropic): + """ + An async wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog. + """ + + _ph_client: PostHogClient + + def __init__(self, posthog_client: PostHogClient, **kwargs): + """ + Args: + posthog_client: PostHog client for tracking usage + **kwargs: Additional arguments passed to the Anthropic client + """ + super().__init__(**kwargs) + self._ph_client = posthog_client + self.messages = AsyncWrappedMessages(self) + + +class AsyncWrappedMessages(AsyncMessages): + _client: AsyncAnthropic + + async def create( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + """ + Create a message using Anthropic's API while tracking usage in PostHog. + + Args: + posthog_distinct_id: Optional ID to associate with the usage event + posthog_trace_id: Optional trace UUID for linking events + posthog_properties: Optional dictionary of extra properties to include in the event + posthog_privacy_mode: Whether to redact sensitive information in tracking + posthog_groups: Optional group analytics properties + **kwargs: Arguments passed to Anthropic's messages.create + """ + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + if kwargs.get("stream", False): + return await self._create_streaming( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + + return await call_llm_and_track_usage_async( + posthog_distinct_id, + self._client._ph_client, + "anthropic", + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + self._client.base_url, + super().create, + **kwargs, + ) + + async def stream( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + return await self._create_streaming( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + + async def _create_streaming( + self, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + **kwargs: Any, + ): + start_time = time.time() + usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0} + accumulated_content = [] + response = await super().create(**kwargs) + + async def generator(): + nonlocal usage_stats + nonlocal accumulated_content + try: + async for event in response: + if hasattr(event, "usage") and event.usage: + usage_stats = { + k: getattr(event.usage, k, 0) + for k in [ + "input_tokens", + "output_tokens", + ] + } + + if hasattr(event, "content") and event.content: + accumulated_content.append(event.content) + + yield event + + finally: + end_time = time.time() + latency = end_time - start_time + output = "".join(accumulated_content) + + await self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + usage_stats, + latency, + output, + ) + + return generator() + + async def _capture_streaming_event( + self, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + kwargs: Dict[str, Any], + usage_stats: Dict[str, int], + latency: float, + output: str, + ): + if posthog_trace_id is None: + posthog_trace_id = uuid.uuid4() + + event_properties = { + "$ai_provider": "anthropic", + "$ai_model": kwargs.get("model"), + "$ai_model_parameters": get_model_params(kwargs), + "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")), + "$ai_output_choices": with_privacy_mode( + self._client._ph_client, + posthog_privacy_mode, + [{"content": output, "role": "assistant"}], + ), + "$ai_http_status": 200, + "$ai_input_tokens": usage_stats.get("input_tokens", 0), + "$ai_output_tokens": usage_stats.get("output_tokens", 0), + "$ai_latency": latency, + "$ai_trace_id": posthog_trace_id, + "$ai_base_url": str(self._client.base_url), + **(posthog_properties or {}), + } + + if posthog_distinct_id is None: + event_properties["$process_person_profile"] = False + + if hasattr(self._client._ph_client, "capture"): + self._client._ph_client.capture( + distinct_id=posthog_distinct_id or posthog_trace_id, + event="$ai_generation", + properties=event_properties, + groups=posthog_groups, + ) diff --git a/posthog/ai/anthropic/anthropic_providers.py b/posthog/ai/anthropic/anthropic_providers.py new file mode 100644 index 00000000..57288046 --- /dev/null +++ b/posthog/ai/anthropic/anthropic_providers.py @@ -0,0 +1,60 @@ +try: + import anthropic +except ImportError: + raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'") + +from posthog.ai.anthropic.anthropic import WrappedMessages +from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages +from posthog.client import Client as PostHogClient + + +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: PostHogClient, **kwargs): + super().__init__(**kwargs) + self._ph_client = posthog_client + 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: PostHogClient, **kwargs): + super().__init__(**kwargs) + self._ph_client = posthog_client + 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: PostHogClient, **kwargs): + super().__init__(**kwargs) + self._ph_client = posthog_client + 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: PostHogClient, **kwargs): + super().__init__(**kwargs) + self._ph_client = posthog_client + self.messages = AsyncWrappedMessages(self) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 4c08a257..7a513b21 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -171,7 +171,7 @@ def on_llm_end( "$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": with_privacy_mode(self._client, self._privacy_mode, {"choices": output}), + "$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, diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 8d34f79a..5987ff4e 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -69,6 +69,7 @@ def create( return call_llm_and_track_usage( posthog_distinct_id, self._client._ph_client, + "openai", posthog_trace_id, posthog_properties, posthog_privacy_mode, @@ -155,17 +156,10 @@ def _capture_streaming_event( "$ai_model": kwargs.get("model"), "$ai_model_parameters": get_model_params(kwargs), "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")), - "$ai_output": with_privacy_mode( + "$ai_output_choices": with_privacy_mode( self._client._ph_client, posthog_privacy_mode, - { - "choices": [ - { - "content": output, - "role": "assistant", - } - ] - }, + [{"content": output, "role": "assistant"}], ), "$ai_http_status": 200, "$ai_input_tokens": usage_stats.get("prompt_tokens", 0), diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 31ab1eb1..1915002e 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -69,6 +69,7 @@ async def create( response = await call_llm_and_track_usage_async( posthog_distinct_id, self._client._ph_client, + "openai", posthog_trace_id, posthog_properties, self._client.base_url, @@ -118,7 +119,7 @@ async def async_generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) - self._capture_streaming_event( + await self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, posthog_properties, @@ -132,7 +133,7 @@ async def async_generator(): return async_generator() - def _capture_streaming_event( + async def _capture_streaming_event( self, posthog_distinct_id: Optional[str], posthog_trace_id: Optional[str], @@ -152,17 +153,10 @@ def _capture_streaming_event( "$ai_model": kwargs.get("model"), "$ai_model_parameters": get_model_params(kwargs), "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")), - "$ai_output": with_privacy_mode( + "$ai_output_choices": with_privacy_mode( self._client._ph_client, posthog_privacy_mode, - { - "choices": [ - { - "content": output, - "role": "assistant", - } - ] - }, + [{"content": output, "role": "assistant"}], ), "$ai_http_status": 200, "$ai_input_tokens": usage_stats.get("prompt_tokens", 0), diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 6f524179..728ddd4d 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -21,23 +21,63 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: "presence_penalty", "n", "stop", - "stream", + "stream", # OpenAI-specific field + "streaming", # Anthropic-specific field ]: if param in kwargs and kwargs[param] is not None: model_params[param] = kwargs[param] return model_params -def format_response(response): +def get_usage(response, provider: str) -> Dict[str, Any]: + if provider == "anthropic": + return { + "input_tokens": response.usage.input_tokens, + "output_tokens": response.usage.output_tokens, + } + elif provider == "openai": + return { + "input_tokens": response.usage.prompt_tokens, + "output_tokens": response.usage.completion_tokens, + } + return { + "input_tokens": 0, + "output_tokens": 0, + } + + +def format_response(response, provider: str): """ Format a regular (non-streaming) response. """ - output = {"choices": []} + output = [] if response is None: return output + if provider == "anthropic": + return format_response_anthropic(response) + elif provider == "openai": + return format_response_openai(response) + return output + + +def format_response_anthropic(response): + output = [] + for choice in response.content: + if choice.text: + output.append( + { + "role": "assistant", + "content": choice.text, + } + ) + return output + + +def format_response_openai(response): + output = [] for choice in response.choices: if choice.message.content: - output["choices"].append( + output.append( { "content": choice.message.content, "role": choice.message.role, @@ -49,6 +89,7 @@ def format_response(response): def call_llm_and_track_usage( posthog_distinct_id: Optional[str], ph_client: PostHogClient, + provider: str, posthog_trace_id: Optional[str], posthog_properties: Optional[Dict[str, Any]], posthog_privacy_mode: bool, @@ -80,19 +121,19 @@ def call_llm_and_track_usage( posthog_trace_id = uuid.uuid4() if response and hasattr(response, "usage"): - usage = response.usage.model_dump() + usage = get_usage(response, provider) - input_tokens = usage.get("prompt_tokens", 0) - output_tokens = usage.get("completion_tokens", 0) event_properties = { - "$ai_provider": "openai", + "$ai_provider": provider, "$ai_model": kwargs.get("model"), "$ai_model_parameters": get_model_params(kwargs), "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")), - "$ai_output": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response)), + "$ai_output_choices": with_privacy_mode( + ph_client, posthog_privacy_mode, format_response(response, provider) + ), "$ai_http_status": http_status, - "$ai_input_tokens": input_tokens, - "$ai_output_tokens": output_tokens, + "$ai_input_tokens": usage.get("input_tokens", 0), + "$ai_output_tokens": usage.get("output_tokens", 0), "$ai_latency": latency, "$ai_trace_id": posthog_trace_id, "$ai_base_url": str(base_url), @@ -120,6 +161,7 @@ def call_llm_and_track_usage( async def call_llm_and_track_usage_async( posthog_distinct_id: Optional[str], ph_client: PostHogClient, + provider: str, posthog_trace_id: Optional[str], posthog_properties: Optional[Dict[str, Any]], posthog_privacy_mode: bool, @@ -147,19 +189,19 @@ async def call_llm_and_track_usage_async( posthog_trace_id = uuid.uuid4() if response and hasattr(response, "usage"): - usage = response.usage.model_dump() + usage = get_usage(response, provider) - input_tokens = usage.get("prompt_tokens", 0) - output_tokens = usage.get("completion_tokens", 0) event_properties = { - "$ai_provider": "openai", + "$ai_provider": provider, "$ai_model": kwargs.get("model"), "$ai_model_parameters": get_model_params(kwargs), "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")), - "$ai_output": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response)), + "$ai_output_choices": with_privacy_mode( + ph_client, posthog_privacy_mode, format_response(response, provider) + ), "$ai_http_status": http_status, - "$ai_input_tokens": input_tokens, - "$ai_output_tokens": output_tokens, + "$ai_input_tokens": usage.get("input_tokens", 0), + "$ai_output_tokens": usage.get("output_tokens", 0), "$ai_latency": latency, "$ai_trace_id": posthog_trace_id, "$ai_base_url": str(base_url), diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py new file mode 100644 index 00000000..ce2169d4 --- /dev/null +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -0,0 +1,273 @@ +import os +import time +from unittest.mock import patch + +import pytest +from anthropic.types import Message, Usage + +from posthog.ai.anthropic import Anthropic, AsyncAnthropic + +ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY") + + +@pytest.fixture +def mock_client(): + with patch("posthog.client.Client") as mock_client: + mock_client.privacy_mode = False + yield mock_client + + +@pytest.fixture +def mock_anthropic_response(): + return Message( + id="msg_123", + type="message", + role="assistant", + content=[{"type": "text", "text": "Test response"}], + model="claude-3-opus-20240229", + usage=Usage( + input_tokens=20, + output_tokens=10, + ), + stop_reason="end_turn", + stop_sequence=None, + ) + + +@pytest.fixture +def mock_anthropic_stream(): + class MockStreamEvent: + def __init__(self, content, usage=None): + self.content = content + self.usage = usage + + def stream_generator(): + yield MockStreamEvent("A") + yield MockStreamEvent("B") + yield MockStreamEvent( + "C", + usage=Usage( + input_tokens=20, + output_tokens=10, + ), + ) + + return stream_generator() + + +def test_basic_completion(mock_client, mock_anthropic_response): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + assert response == mock_anthropic_response + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "anthropic" + assert props["$ai_model"] == "claude-3-opus-20240229" + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}] + assert props["$ai_input_tokens"] == 20 + assert props["$ai_output_tokens"] == 10 + assert props["$ai_http_status"] == 200 + assert props["foo"] == "bar" + assert isinstance(props["$ai_latency"], float) + + +def test_streaming(mock_client, mock_anthropic_stream): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + # Consume the stream + chunks = list(response) + assert len(chunks) == 3 + assert chunks[0].content == "A" + assert chunks[1].content == "B" + assert chunks[2].content == "C" + + # Wait a bit to ensure the capture is called + time.sleep(0.1) + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "anthropic" + assert props["$ai_model"] == "claude-3-opus-20240229" + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}] + assert props["$ai_input_tokens"] == 20 + assert props["$ai_output_tokens"] == 10 + assert isinstance(props["$ai_latency"], float) + assert props["foo"] == "bar" + + +def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.stream( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + # Consume the stream + chunks = list(response) + assert len(chunks) == 3 + assert chunks[0].content == "A" + assert chunks[1].content == "B" + assert chunks[2].content == "C" + + # Wait a bit to ensure the capture is called + time.sleep(0.1) + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "anthropic" + assert props["$ai_model"] == "claude-3-opus-20240229" + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}] + assert props["$ai_input_tokens"] == 20 + assert props["$ai_output_tokens"] == 10 + assert isinstance(props["$ai_latency"], float) + assert props["foo"] == "bar" + + +def test_groups(mock_client, mock_anthropic_response): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + posthog_distinct_id="test-id", + posthog_groups={"company": "test_company"}, + ) + + assert response == mock_anthropic_response + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + assert call_args["groups"] == {"company": "test_company"} + + +def test_privacy_mode_local(mock_client, mock_anthropic_response): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + posthog_distinct_id="test-id", + posthog_privacy_mode=True, + ) + + assert response == mock_anthropic_response + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + assert props["$ai_input"] is None + assert props["$ai_output_choices"] is None + + +def test_privacy_mode_global(mock_client, mock_anthropic_response): + with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response): + mock_client.privacy_mode = True + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Hello"}], + posthog_distinct_id="test-id", + posthog_privacy_mode=False, + ) + + assert response == mock_anthropic_response + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + assert props["$ai_input"] is None + assert props["$ai_output_choices"] is None + + +@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set") +def test_basic_integration(mock_client): + client = Anthropic(posthog_client=mock_client) + client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "You must always answer with 'Bar'."}], + max_tokens=1, + temperature=0, + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "anthropic" + assert props["$ai_model"] == "claude-3-opus-20240229" + assert props["$ai_input"] == [{"role": "user", "content": "You must always answer with 'Bar'."}] + assert props["$ai_output_choices"][0]["role"] == "assistant" + assert props["$ai_input_tokens"] == 16 + assert props["$ai_output_tokens"] == 1 + assert props["$ai_http_status"] == 200 + assert props["foo"] == "bar" + assert isinstance(props["$ai_latency"], float) + + +@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set") +async def test_basic_async_integration(mock_client): + client = AsyncAnthropic(posthog_client=mock_client) + await client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "You must always answer with 'Bar'."}], + max_tokens=1, + temperature=0, + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "anthropic" + assert props["$ai_model"] == "claude-3-opus-20240229" + assert props["$ai_input"] == [{"role": "user", "content": "You must always answer with 'Bar'."}] + assert props["$ai_output_choices"][0]["role"] == "assistant" + assert props["$ai_input_tokens"] == 16 + assert props["$ai_output_tokens"] == 1 + assert props["$ai_http_status"] == 200 + assert props["foo"] == "bar" + assert isinstance(props["$ai_latency"], float) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 2753d6ce..697f5f54 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest +from langchain_anthropic.chat_models import ChatAnthropic from langchain_community.chat_models.fake import FakeMessagesListChatModel from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM from langchain_core.messages import AIMessage @@ -15,6 +16,7 @@ from posthog.ai.langchain import CallbackHandler OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY") @pytest.fixture(scope="function") @@ -120,9 +122,9 @@ def test_basic_chat_chain(mock_client, stream): {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Who won the world series in 2020?"}, ] - assert props["$ai_output"] == { - "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}] - } + assert props["$ai_output_choices"] == [ + {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."} + ] assert props["$ai_input_tokens"] == 10 assert props["$ai_output_tokens"] == 10 assert props["$ai_http_status"] == 200 @@ -165,9 +167,9 @@ async def test_async_basic_chat_chain(mock_client, stream): {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Who won the world series in 2020?"}, ] - assert props["$ai_output"] == { - "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}] - } + assert props["$ai_output_choices"] == [ + {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."} + ] assert props["$ai_input_tokens"] == 10 assert props["$ai_output_tokens"] == 10 assert props["$ai_http_status"] == 200 @@ -200,7 +202,7 @@ def test_basic_llm_chain(mock_client, Model, stream): assert "$ai_model" in props assert "$ai_provider" in props assert props["$ai_input"] == ["Who won the world series in 2020?"] - assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]} + assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."] assert props["$ai_http_status"] == 200 assert props["$ai_trace_id"] is not None assert isinstance(props["$ai_latency"], float) @@ -231,7 +233,7 @@ async def test_async_basic_llm_chain(mock_client, Model, stream): assert "$ai_model" in props assert "$ai_provider" in props assert props["$ai_input"] == ["Who won the world series in 2020?"] - assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]} + assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."] assert props["$ai_http_status"] == 200 assert props["$ai_trace_id"] is not None assert isinstance(props["$ai_latency"], float) @@ -258,7 +260,7 @@ def test_trace_id_for_multiple_chains(mock_client): assert "$ai_model" in first_call_props assert "$ai_provider" in first_call_props assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}] - assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]} + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] assert first_call_props["$ai_http_status"] == 200 assert first_call_props["$ai_trace_id"] is not None assert isinstance(first_call_props["$ai_latency"], float) @@ -270,7 +272,7 @@ def test_trace_id_for_multiple_chains(mock_client): assert "$ai_model" in second_call_props assert "$ai_provider" in second_call_props assert second_call_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}] - assert second_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]} + assert second_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] assert second_call_props["$ai_http_status"] == 200 assert second_call_props["$ai_trace_id"] is not None assert isinstance(second_call_props["$ai_latency"], float) @@ -338,7 +340,7 @@ def test_metadata(mock_client): assert first_call_props["$ai_trace_id"] == "test-trace-id" assert first_call_props["foo"] == "bar" assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}] - assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]} + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] assert first_call_props["$ai_http_status"] == 200 assert isinstance(first_call_props["$ai_latency"], float) @@ -392,7 +394,7 @@ def test_openai_error(mock_client): props = args["properties"] assert props["$ai_http_status"] == 401 assert props["$ai_input"] == [{"role": "user", "content": "Foo"}] - assert "$ai_output" not in props + assert "$ai_output_choices" not in props @pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set") @@ -443,15 +445,13 @@ def test_openai_chain(mock_client): {"role": "system", "content": 'You must always answer with "Bar".'}, {"role": "user", "content": "Foo"}, ] - assert first_call_props["$ai_output"] == { - "choices": [ - { - "role": "assistant", - "content": "Bar", - "refusal": None, - } - ] - } + assert first_call_props["$ai_output_choices"] == [ + { + "role": "assistant", + "content": "Bar", + "additional_kwargs": {"refusal": None}, + } + ] assert first_call_props["$ai_http_status"] == 200 assert isinstance(first_call_props["$ai_latency"], float) assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency @@ -486,19 +486,17 @@ def test_openai_captures_multiple_generations(mock_client): {"role": "system", "content": 'You must always answer with "Bar".'}, {"role": "user", "content": "Foo"}, ] - assert first_call_props["$ai_output"] == { - "choices": [ - { - "role": "assistant", - "content": "Bar", - "refusal": None, - }, - { - "role": "assistant", - "content": "Bar", - }, - ] - } + assert first_call_props["$ai_output_choices"] == [ + { + "role": "assistant", + "content": "Bar", + "additional_kwargs": {"refusal": None}, + }, + { + "role": "assistant", + "content": "Bar", + }, + ] # langchain-openai for langchain v3 if "max_completion_tokens" in first_call_props["$ai_model_parameters"]: @@ -544,7 +542,7 @@ def test_openai_streaming(mock_client): {"role": "system", "content": 'You must always answer with "Bar".'}, {"role": "user", "content": "Foo"}, ] - assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]} + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] assert first_call_props["$ai_http_status"] == 200 assert first_call_props["$ai_input_tokens"] == 20 assert first_call_props["$ai_output_tokens"] == 1 @@ -576,7 +574,7 @@ async def test_async_openai_streaming(mock_client): {"role": "system", "content": 'You must always answer with "Bar".'}, {"role": "user", "content": "Foo"}, ] - assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]} + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] assert first_call_props["$ai_http_status"] == 200 assert first_call_props["$ai_input_tokens"] == 20 assert first_call_props["$ai_output_tokens"] == 1 @@ -630,7 +628,7 @@ def test_privacy_mode_local(mock_client): assert mock_client.capture.call_count == 1 call = mock_client.capture.call_args[1] assert call["properties"]["$ai_input"] is None - assert call["properties"]["$ai_output"] is None + assert call["properties"]["$ai_output_choices"] is None def test_privacy_mode_global(mock_client): @@ -649,7 +647,90 @@ def test_privacy_mode_global(mock_client): assert mock_client.capture.call_count == 1 call = mock_client.capture.call_args[1] assert call["properties"]["$ai_input"] is None - assert call["properties"]["$ai_output"] is None + assert call["properties"]["$ai_output_choices"] is None + + +@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set") +def test_anthropic_chain(mock_client): + prompt = ChatPromptTemplate.from_messages( + [ + ("system", 'You must always answer with "Bar".'), + ("user", "Foo"), + ] + ) + chain = prompt | ChatAnthropic( + api_key=ANTHROPIC_API_KEY, + model="claude-3-opus-20240229", + temperature=0, + max_tokens=1, + ) + callbacks = CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"}) + start_time = time.time() + result = chain.invoke({}, config={"callbacks": [callbacks]}) + approximate_latency = math.floor(time.time() - start_time) + + assert result.content == "Bar" + assert mock_client.capture.call_count == 1 + + first_call_args = mock_client.capture.call_args[1] + first_call_props = first_call_args["properties"] + assert first_call_args["event"] == "$ai_generation" + assert first_call_props["$ai_trace_id"] == "test-trace-id" + assert first_call_props["$ai_provider"] == "anthropic" + assert first_call_props["$ai_model"] == "claude-3-opus-20240229" + assert first_call_props["foo"] == "bar" + + assert first_call_props["$ai_model_parameters"] == { + "temperature": 0.0, + "max_tokens": 1, + "streaming": False, + } + assert first_call_props["$ai_input"] == [ + {"role": "system", "content": 'You must always answer with "Bar".'}, + {"role": "user", "content": "Foo"}, + ] + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] + assert first_call_props["$ai_http_status"] == 200 + assert isinstance(first_call_props["$ai_latency"], float) + assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency + assert first_call_props["$ai_input_tokens"] == 17 + assert first_call_props["$ai_output_tokens"] == 1 + + +@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set") +async def test_async_anthropic_streaming(mock_client): + prompt = ChatPromptTemplate.from_messages( + [ + ("system", 'You must always answer with "Bar".'), + ("user", "Foo"), + ] + ) + chain = prompt | ChatAnthropic( + api_key=ANTHROPIC_API_KEY, + model="claude-3-opus-20240229", + temperature=0, + max_tokens=1, + streaming=True, + stream_usage=True, + ) + callbacks = CallbackHandler(mock_client) + result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})] + result = sum(result[1:], result[0]) + + assert result.content == "Bar" + assert mock_client.capture.call_count == 1 + + first_call_args = mock_client.capture.call_args[1] + first_call_props = first_call_args["properties"] + assert first_call_props["$ai_model_parameters"]["streaming"] + assert first_call_props["$ai_input"] == [ + {"role": "system", "content": 'You must always answer with "Bar".'}, + {"role": "user", "content": "Foo"}, + ] + assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}] + assert first_call_props["$ai_http_status"] == 200 + assert first_call_props["$ai_input_tokens"] == 17 + assert first_call_props["$ai_output_tokens"] is not None def test_tool_calls(mock_client): @@ -679,7 +760,7 @@ def test_tool_calls(mock_client): assert mock_client.capture.call_count == 1 call = mock_client.capture.call_args[1] - assert call["properties"]["$ai_output"]["choices"][0]["tool_calls"] == [ + assert call["properties"]["$ai_output_choices"][0]["tool_calls"] == [ { "type": "function", "id": "123", @@ -689,4 +770,4 @@ def test_tool_calls(mock_client): }, } ] - assert "additional_kwargs" not in call["properties"]["$ai_output"]["choices"][0] + assert "additional_kwargs" not in call["properties"]["$ai_output_choices"][0] diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 69acee3c..d2dbc06f 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -83,7 +83,7 @@ def test_basic_completion(mock_client, mock_openai_response): assert props["$ai_provider"] == "openai" assert props["$ai_model"] == "gpt-4" assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] - assert props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Test response"}]} + assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}] assert props["$ai_input_tokens"] == 20 assert props["$ai_output_tokens"] == 10 assert props["$ai_http_status"] == 200 @@ -152,7 +152,7 @@ def test_privacy_mode_local(mock_client, mock_openai_response): call_args = mock_client.capture.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None - assert props["$ai_output"] is None + assert props["$ai_output_choices"] is None def test_privacy_mode_global(mock_client, mock_openai_response): @@ -172,4 +172,4 @@ def test_privacy_mode_global(mock_client, mock_openai_response): call_args = mock_client.capture.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None - assert props["$ai_output"] is None + assert props["$ai_output_choices"] is None diff --git a/posthog/version.py b/posthog/version.py index 5d5e3c1b..b784d3d0 100644 --- a/posthog/version.py +++ b/posthog/version.py @@ -1,4 +1,4 @@ -VERSION = "3.8.3" +VERSION = "3.8.4" if __name__ == "__main__": print(VERSION, end="") # noqa: T201 diff --git a/setup.py b/setup.py index e1f125ce..06fb4580 100644 --- a/setup.py +++ b/setup.py @@ -40,8 +40,11 @@ "pytest-timeout", "pytest-asyncio", "django", + "openai", + "anthropic", "langchain-community>=0.2.0", "langchain-openai>=0.2.0", + "langchain-anthropic>=0.2.0", ], "sentry": ["sentry-sdk", "django"], "langchain": ["langchain>=0.2.0"], @@ -61,6 +64,7 @@ "posthog.ai", "posthog.ai.langchain", "posthog.ai.openai", + "posthog.ai.anthropic", "posthog.test", "posthog.sentry", "posthog.exception_integrations", diff --git a/setup_analytics.py b/setup_analytics.py index 55db1d76..9c91a5dc 100644 --- a/setup_analytics.py +++ b/setup_analytics.py @@ -32,6 +32,7 @@ "posthoganalytics.ai", "posthoganalytics.ai.langchain", "posthoganalytics.ai.openai", + "posthoganalytics.ai.anthropic", "posthoganalytics.test", "posthoganalytics.sentry", "posthoganalytics.exception_integrations",