From 5a9b19e964fd0ab056a6c0fc96fe5341c27a43ef Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 17:20:45 -0800 Subject: [PATCH 01/17] feat: anthropic --- posthog/ai/anthropic/__init__.py | 3 + posthog/ai/anthropic/anthropic.py | 191 ++++++++++++++++++ posthog/ai/langchain/callbacks.py | 2 +- posthog/ai/openai/openai.py | 12 +- posthog/ai/openai/openai_async.py | 12 +- posthog/ai/utils.py | 38 +++- posthog/test/ai/anthropic/test_anthropic.py | 209 ++++++++++++++++++++ 7 files changed, 441 insertions(+), 26 deletions(-) create mode 100644 posthog/ai/anthropic/__init__.py create mode 100644 posthog/ai/anthropic/anthropic.py create mode 100644 posthog/test/ai/anthropic/test_anthropic.py diff --git a/posthog/ai/anthropic/__init__.py b/posthog/ai/anthropic/__init__.py new file mode 100644 index 00000000..dddd9285 --- /dev/null +++ b/posthog/ai/anthropic/__init__.py @@ -0,0 +1,3 @@ +from .anthropic import Anthropic + +__all__ = ["Anthropic"] diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py new file mode 100644 index 00000000..7d5cc4b8 --- /dev/null +++ b/posthog/ai/anthropic/anthropic.py @@ -0,0 +1,191 @@ +try: + import anthropic + from anthropic.resources import Messages +except ImportError: + raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'") + +from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode +from posthog.client import Client as PostHogClient +from typing import Any, Dict, Optional +import uuid +import time + + +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], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + **kwargs: Any, + ): + 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] = {} + 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 = { + "input_tokens": getattr(event.usage, "input_tokens", 0), + "output_tokens": getattr(event.usage, "output_tokens", 0), + } + + 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/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 9ce17b9f..cbcdab4f 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..97b5a4e0 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, @@ -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..eedcd94e 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -28,16 +28,38 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: return model_params -def format_response(response): +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 +71,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, @@ -85,11 +108,11 @@ def call_llm_and_track_usage( 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, @@ -120,6 +143,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, @@ -152,11 +176,11 @@ async def call_llm_and_track_usage_async( 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, diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py new file mode 100644 index 00000000..8cd4582a --- /dev/null +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -0,0 +1,209 @@ +import time +from unittest.mock import patch + +import pytest +from anthropic.types import Message, Usage, MessageStreamEvent + +from posthog.ai.anthropic import Anthropic + + +@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(): + yield from [ + MessageStreamEvent( + id="msg_123", + type="message", + role="assistant", + content="A", + ), + MessageStreamEvent( + id="msg_123", + type="message", + role="assistant", + content="B", + ), + MessageStreamEvent( + id="msg_123", + type="message", + role="assistant", + content="C", + ), + ] + + +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" + 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_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" + 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_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 From 79b4fe878223e3cc67666c9e657aea0fe397e912 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 17:59:25 -0800 Subject: [PATCH 02/17] feat: anthroipic --- CHANGELOG.md | 5 ++ posthog/ai/anthropic/anthropic.py | 23 ++++++--- posthog/ai/utils.py | 33 +++++++++---- posthog/test/ai/anthropic/test_anthropic.py | 55 +++++++++++---------- posthog/version.py | 2 +- 5 files changed, 72 insertions(+), 46 deletions(-) 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/anthropic.py b/posthog/ai/anthropic/anthropic.py index 7d5cc4b8..ddd85a64 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -80,13 +80,16 @@ def create( def stream( 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]], + 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, @@ -106,7 +109,7 @@ def _create_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: Dict[str, int] = {} + usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0} accumulated_content = [] response = super().create(**kwargs) @@ -117,8 +120,11 @@ def generator(): for event in response: if hasattr(event, "usage") and event.usage: usage_stats = { - "input_tokens": getattr(event.usage, "input_tokens", 0), - "output_tokens": getattr(event.usage, "output_tokens", 0), + k: getattr(event.usage, k, 0) + for k in [ + "input_tokens", + "output_tokens", + ] } if hasattr(event, "content") and event.content: @@ -130,6 +136,7 @@ def generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) + self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index eedcd94e..d2e996de 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -28,6 +28,23 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: return model_params +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. @@ -103,10 +120,8 @@ 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": provider, "$ai_model": kwargs.get("model"), @@ -114,8 +129,8 @@ def call_llm_and_track_usage( "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")), "$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), @@ -171,10 +186,8 @@ 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": provider, "$ai_model": kwargs.get("model"), @@ -182,8 +195,8 @@ async def call_llm_and_track_usage_async( "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")), "$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 index 8cd4582a..0c60f01f 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from anthropic.types import Message, Usage, MessageStreamEvent +from anthropic.types import Message, Usage from posthog.ai.anthropic import Anthropic @@ -33,26 +33,23 @@ def mock_anthropic_response(): @pytest.fixture def mock_anthropic_stream(): - yield from [ - MessageStreamEvent( - id="msg_123", - type="message", - role="assistant", - content="A", - ), - MessageStreamEvent( - id="msg_123", - type="message", - role="assistant", - content="B", - ), - MessageStreamEvent( - id="msg_123", - type="message", - role="assistant", - content="C", - ), - ] + 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): @@ -101,6 +98,9 @@ def test_streaming(mock_client, mock_anthropic_stream): 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] @@ -111,12 +111,11 @@ def test_streaming(mock_client, mock_anthropic_stream): 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_output_choices"] == [{"role": "assistant", "content": "ABC"}] 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) + assert props["foo"] == "bar" def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream): @@ -135,6 +134,9 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream): 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] @@ -145,12 +147,11 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream): 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_output_choices"] == [{"role": "assistant", "content": "ABC"}] 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) + assert props["foo"] == "bar" def test_groups(mock_client, mock_anthropic_response): 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 From 59dc44de93fb853e15a31274a133455427652b14 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 20:33:57 -0800 Subject: [PATCH 03/17] feat: add vertex + bedrock --- posthog/ai/anthropic/anthropic_async.py | 198 ++++++++++++++++++++ posthog/ai/anthropic/anthropic_providers.py | 59 ++++++ posthog/ai/openai/openai_async.py | 4 +- 3 files changed, 259 insertions(+), 2 deletions(-) create mode 100644 posthog/ai/anthropic/anthropic_async.py create mode 100644 posthog/ai/anthropic/anthropic_providers.py diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py new file mode 100644 index 00000000..8a74fced --- /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'") + +from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, with_privacy_mode +from posthog.client import Client as PostHogClient +from typing import Any, Dict, Optional +import uuid +import time + + +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..a869b709 --- /dev/null +++ b/posthog/ai/anthropic/anthropic_providers.py @@ -0,0 +1,59 @@ +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/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 97b5a4e0..1915002e 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -119,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, @@ -133,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], From e0cd9c145640b0b109d1ce204a49b9fd1fcc3fe8 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 20:34:49 -0800 Subject: [PATCH 04/17] feat: export anthropic --- posthog/ai/anthropic/__init__.py | 4 +++- setup.py | 1 + setup_analytics.py | 1 + 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/posthog/ai/anthropic/__init__.py b/posthog/ai/anthropic/__init__.py index dddd9285..a28bf45b 100644 --- a/posthog/ai/anthropic/__init__.py +++ b/posthog/ai/anthropic/__init__.py @@ -1,3 +1,5 @@ from .anthropic import Anthropic +from .anthropic_async import AsyncAnthropic +from .anthropic_providers import AnthropicBedrock, AsyncAnthropicBedrock, AnthropicVertex, AsyncAnthropicVertex -__all__ = ["Anthropic"] +__all__ = ["Anthropic", "AsyncAnthropic", "AnthropicBedrock", "AsyncAnthropicBedrock", "AnthropicVertex", "AsyncAnthropicVertex"] diff --git a/setup.py b/setup.py index e1f125ce..330e83da 100644 --- a/setup.py +++ b/setup.py @@ -61,6 +61,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", From 5b287e0f35c62b7f3c478d6b494de77e62481939 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 20:36:55 -0800 Subject: [PATCH 05/17] fi9x: isort + black --- posthog/ai/anthropic/__init__.py | 11 +++++++++-- posthog/ai/anthropic/anthropic.py | 16 ++++++++-------- posthog/ai/anthropic/anthropic_async.py | 16 ++++++++-------- posthog/ai/anthropic/anthropic_providers.py | 1 + posthog/ai/utils.py | 8 ++++++-- posthog/test/ai/anthropic/test_anthropic.py | 4 ++-- 6 files changed, 34 insertions(+), 22 deletions(-) diff --git a/posthog/ai/anthropic/__init__.py b/posthog/ai/anthropic/__init__.py index a28bf45b..db4db080 100644 --- a/posthog/ai/anthropic/__init__.py +++ b/posthog/ai/anthropic/__init__.py @@ -1,5 +1,12 @@ from .anthropic import Anthropic from .anthropic_async import AsyncAnthropic -from .anthropic_providers import AnthropicBedrock, AsyncAnthropicBedrock, AnthropicVertex, AsyncAnthropicVertex +from .anthropic_providers import AnthropicBedrock, AnthropicVertex, AsyncAnthropicBedrock, AsyncAnthropicVertex -__all__ = ["Anthropic", "AsyncAnthropic", "AnthropicBedrock", "AsyncAnthropicBedrock", "AnthropicVertex", "AsyncAnthropicVertex"] +__all__ = [ + "Anthropic", + "AsyncAnthropic", + "AnthropicBedrock", + "AsyncAnthropicBedrock", + "AnthropicVertex", + "AsyncAnthropicVertex", +] diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py index ddd85a64..a6b12b8b 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -4,11 +4,12 @@ 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 -from typing import Any, Dict, Optional -import uuid -import time class Anthropic(anthropic.Anthropic): @@ -43,7 +44,7 @@ def create( ): """ 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 @@ -77,7 +78,7 @@ def create( super().create, **kwargs, ) - + def stream( self, posthog_distinct_id: Optional[str] = None, @@ -89,7 +90,7 @@ def stream( ): if posthog_trace_id is None: posthog_trace_id = uuid.uuid4() - + return self._create_streaming( posthog_distinct_id, posthog_trace_id, @@ -136,7 +137,7 @@ def generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) - + self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -195,4 +196,3 @@ def _capture_streaming_event( properties=event_properties, groups=posthog_groups, ) - diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py index 8a74fced..31a3b355 100644 --- a/posthog/ai/anthropic/anthropic_async.py +++ b/posthog/ai/anthropic/anthropic_async.py @@ -4,11 +4,12 @@ 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 -from typing import Any, Dict, Optional -import uuid -import time class AsyncAnthropic(anthropic.AsyncAnthropic): @@ -43,7 +44,7 @@ async def create( ): """ 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 @@ -77,7 +78,7 @@ async def create( super().create, **kwargs, ) - + async def stream( self, posthog_distinct_id: Optional[str] = None, @@ -89,7 +90,7 @@ async def stream( ): if posthog_trace_id is None: posthog_trace_id = uuid.uuid4() - + return await self._create_streaming( posthog_distinct_id, posthog_trace_id, @@ -136,7 +137,7 @@ async def generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) - + await self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -195,4 +196,3 @@ async def _capture_streaming_event( properties=event_properties, groups=posthog_groups, ) - diff --git a/posthog/ai/anthropic/anthropic_providers.py b/posthog/ai/anthropic/anthropic_providers.py index a869b709..57288046 100644 --- a/posthog/ai/anthropic/anthropic_providers.py +++ b/posthog/ai/anthropic/anthropic_providers.py @@ -7,6 +7,7 @@ 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. diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index d2e996de..2c44039e 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -127,7 +127,9 @@ def call_llm_and_track_usage( "$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_choices": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response, provider)), + "$ai_output_choices": with_privacy_mode( + ph_client, posthog_privacy_mode, format_response(response, provider) + ), "$ai_http_status": http_status, "$ai_input_tokens": usage.get("input_tokens", 0), "$ai_output_tokens": usage.get("output_tokens", 0), @@ -193,7 +195,9 @@ async def call_llm_and_track_usage_async( "$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_choices": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response, provider)), + "$ai_output_choices": with_privacy_mode( + ph_client, posthog_privacy_mode, format_response(response, provider) + ), "$ai_http_status": http_status, "$ai_input_tokens": usage.get("input_tokens", 0), "$ai_output_tokens": usage.get("output_tokens", 0), diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 0c60f01f..2cb1e411 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -98,7 +98,7 @@ def test_streaming(mock_client, mock_anthropic_stream): 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 @@ -134,7 +134,7 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream): 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 From a193a77395138936e5c7eea5701c59e5dd583723 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Thu, 16 Jan 2025 20:40:51 -0800 Subject: [PATCH 06/17] fix: tests ai_output --- posthog/test/ai/langchain/test_callbacks.py | 28 ++++++++++----------- posthog/test/ai/openai/test_openai.py | 6 ++--- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 321f5fa0..570a0b6b 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -120,7 +120,7 @@ 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"] == { + assert props["$ai_output_choices"] == { "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}] } assert props["$ai_input_tokens"] == 10 @@ -165,7 +165,7 @@ 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"] == { + assert props["$ai_output_choices"] == { "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}] } assert props["$ai_input_tokens"] == 10 @@ -200,7 +200,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"] == {"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 +231,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"] == {"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 +258,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"] == {"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 +270,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"] == {"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 +338,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"] == {"choices": [{"role": "assistant", "content": "Bar"}]} assert first_call_props["$ai_http_status"] == 200 assert isinstance(first_call_props["$ai_latency"], float) @@ -392,7 +392,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,7 +443,7 @@ 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"] == { + assert first_call_props["$ai_output_choices"] == { "choices": [ { "role": "assistant", @@ -486,7 +486,7 @@ 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"] == { + assert first_call_props["$ai_output_choices"] == { "choices": [ { "role": "assistant", @@ -544,7 +544,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"] == {"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 +576,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"] == {"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 +630,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,4 +649,4 @@ 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 diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 69acee3c..90d66203 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"] == {"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 From 212f99a7d59489c7326c6c3446e5d27a521edc7b Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:28:40 +0100 Subject: [PATCH 07/17] fix: test dependency --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 330e83da..46506daf 100644 --- a/setup.py +++ b/setup.py @@ -42,6 +42,7 @@ "django", "langchain-community>=0.2.0", "langchain-openai>=0.2.0", + "anthropic", ], "sentry": ["sentry-sdk", "django"], "langchain": ["langchain>=0.2.0"], From fcd08f0f01195bc1a7ddd6c15b2ccbea2919b223 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:41:15 +0100 Subject: [PATCH 08/17] test: integration test --- posthog/test/ai/anthropic/test_anthropic.py | 31 +++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 2cb1e411..42ca63e3 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -1,3 +1,4 @@ +import os import time from unittest.mock import patch @@ -208,3 +209,33 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response): props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None + + +@pytest.mark.skipif(not os.getenv("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) From 589c2e3647118781867c4f4c74265907fc29419d Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:43:24 +0100 Subject: [PATCH 09/17] fix: langchain tests --- posthog/test/ai/langchain/test_callbacks.py | 66 ++++++++++----------- 1 file changed, 31 insertions(+), 35 deletions(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 570a0b6b..54d937e3 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -120,9 +120,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"] == { - "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 +165,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"] == { - "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 +200,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"] == {"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 +231,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"] == {"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 +258,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"] == {"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 +270,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"] == {"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 +338,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"] == {"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) @@ -443,15 +443,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"] == { - "choices": [ - { - "role": "assistant", - "content": "Bar", - "additional_kwargs": {"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 +484,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"] == { - "choices": [ - { - "role": "assistant", - "content": "Bar", - "additional_kwargs": {"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 +540,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"] == {"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 +572,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"] == {"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 From b64d5878a5a92e344fd1e3068fa44844db62000b Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:44:41 +0100 Subject: [PATCH 10/17] test: add async test case --- posthog/test/ai/anthropic/test_anthropic.py | 32 ++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 42ca63e3..5ce5a67a 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -5,7 +5,7 @@ import pytest from anthropic.types import Message, Usage -from posthog.ai.anthropic import Anthropic +from posthog.ai.anthropic import Anthropic, AsyncAnthropic @pytest.fixture @@ -239,3 +239,33 @@ def test_basic_integration(mock_client): assert props["$ai_http_status"] == 200 assert props["foo"] == "bar" assert isinstance(props["$ai_latency"], float) + + +@pytest.mark.skipif(not os.getenv("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) From 69eba21d93bfb711488cc04213ff0f33227d46af Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:45:40 +0100 Subject: [PATCH 11/17] fix: linter --- posthog/test/ai/anthropic/test_anthropic.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 5ce5a67a..ce2169d4 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -7,6 +7,8 @@ from posthog.ai.anthropic import Anthropic, AsyncAnthropic +ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY") + @pytest.fixture def mock_client(): @@ -211,7 +213,7 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response): assert props["$ai_output_choices"] is None -@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY is not set") +@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( @@ -241,7 +243,7 @@ def test_basic_integration(mock_client): assert isinstance(props["$ai_latency"], float) -@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY is not set") +@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( From 616931add353d495d8dc5df6d9a527b7f9310ec0 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:57:58 +0100 Subject: [PATCH 12/17] fix: include packages for testing --- setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 46506daf..06fb4580 100644 --- a/setup.py +++ b/setup.py @@ -40,9 +40,11 @@ "pytest-timeout", "pytest-asyncio", "django", + "openai", + "anthropic", "langchain-community>=0.2.0", "langchain-openai>=0.2.0", - "anthropic", + "langchain-anthropic>=0.2.0", ], "sentry": ["sentry-sdk", "django"], "langchain": ["langchain>=0.2.0"], From bd497b35b767e264869277ebd3450bc7d37328c1 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 14:58:23 +0100 Subject: [PATCH 13/17] feat: langchain callbacks for anthropic --- posthog/ai/utils.py | 3 +- posthog/test/ai/langchain/test_callbacks.py | 87 ++++++++++++++++++++- posthog/test/ai/openai/test_openai.py | 2 +- 3 files changed, 89 insertions(+), 3 deletions(-) diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 2c44039e..728ddd4d 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -21,7 +21,8 @@ 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] diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 54d937e3..269d88e9 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -11,10 +11,12 @@ from langchain_core.prompts import ChatPromptTemplate from langchain_core.runnables import RunnableLambda from langchain_openai.chat_models import ChatOpenAI - +from langchain_anthropic.chat_models import ChatAnthropic 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") @@ -646,3 +648,86 @@ def test_privacy_mode_global(mock_client): call = mock_client.capture.call_args[1] assert call["properties"]["$ai_input"] 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 diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 90d66203..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"] == {"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 From cea91e8ea50d5321af31b66c7d972e0cca130676 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 17 Jan 2025 16:23:15 +0100 Subject: [PATCH 14/17] fix: sort imports --- posthog/test/ai/langchain/test_callbacks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 269d88e9..92d422ff 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -5,15 +5,15 @@ 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 from langchain_core.prompts import ChatPromptTemplate from langchain_core.runnables import RunnableLambda from langchain_openai.chat_models import ChatOpenAI -from langchain_anthropic.chat_models import ChatAnthropic -from posthog.ai.langchain import CallbackHandler +from posthog.ai.langchain import CallbackHandler OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY") From 46c519595e0d91b6d837a45144893d2481cbc03c Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Fri, 17 Jan 2025 12:00:52 -0800 Subject: [PATCH 15/17] fix: broken merge --- posthog/test/ai/langchain/test_callbacks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index e06b4a54..08fbcfcd 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -731,7 +731,6 @@ async def test_async_anthropic_streaming(mock_client): 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 - assert call["properties"]["$ai_output"] is None def test_tool_calls(mock_client): From 0601945106f52b959e5f27bbdd009b46a101474e Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Fri, 17 Jan 2025 12:04:26 -0800 Subject: [PATCH 16/17] fix: always miss one --- posthog/test/ai/langchain/test_callbacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 08fbcfcd..e79e35c7 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -760,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", From 5405e442cb2b53337a57e38f4e4d2b2975332ef3 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Fri, 17 Jan 2025 12:30:40 -0800 Subject: [PATCH 17/17] fix: always miss two* --- posthog/test/ai/langchain/test_callbacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index e79e35c7..697f5f54 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -770,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]