anthropic.py

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, List, Optional

from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
from posthog.ai.utils import (
    call_llm_and_track_usage,
    merge_usage_stats,
)
from posthog.ai.anthropic.anthropic_converter import (
    extract_anthropic_usage_from_event,
    handle_anthropic_content_block_start,
    handle_anthropic_text_delta,
    handle_anthropic_tool_delta,
    finalize_anthropic_tool_input,
)
from posthog.ai.sanitization import sanitize_anthropic
from posthog.client import Client as PostHogClient
from posthog import setup


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: Optional[PostHogClient] = None, **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 or setup()
        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 = str(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 = str(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: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
        accumulated_content = ""
        content_blocks: List[StreamingContentBlock] = []
        tools_in_progress: Dict[str, ToolInProgress] = {}
        current_text_block: Optional[StreamingContentBlock] = None
        stop_reason: Optional[str] = None
        response = super().create(**kwargs)

        def generator():
            nonlocal usage_stats
            nonlocal accumulated_content
            nonlocal content_blocks
            nonlocal tools_in_progress
            nonlocal current_text_block
            nonlocal stop_reason

            try:
                for event in response:
                    # Extract usage stats from event
                    event_usage = extract_anthropic_usage_from_event(event)
                    merge_usage_stats(usage_stats, event_usage)

                    # Handle content block start events
                    if hasattr(event, "type") and event.type == "content_block_start":
                        block, tool = handle_anthropic_content_block_start(event)

                        if block:
                            content_blocks.append(block)

                            if block.get("type") == "text":
                                current_text_block = block
                            else:
                                current_text_block = None

                        if tool:
                            tool_id = tool["block"].get("id")
                            if tool_id:
                                tools_in_progress[tool_id] = tool

                    # Handle text delta events
                    delta_text = handle_anthropic_text_delta(event, current_text_block)

                    if delta_text:
                        accumulated_content += delta_text

                    # Handle tool input delta events
                    handle_anthropic_tool_delta(
                        event, content_blocks, tools_in_progress
                    )

                    # Handle content block stop events
                    if hasattr(event, "type") and event.type == "content_block_stop":
                        current_text_block = None
                        finalize_anthropic_tool_input(
                            event, content_blocks, tools_in_progress
                        )

                    # Capture stop reason from message_delta events
                    if hasattr(event, "type") and event.type == "message_delta":
                        delta = getattr(event, "delta", None)
                        if delta is not None:
                            delta_stop_reason = getattr(delta, "stop_reason", None)
                            if delta_stop_reason is not None:
                                stop_reason = delta_stop_reason

                    yield event

            finally:
                end_time = time.time()
                latency = end_time - start_time

                self._capture_streaming_event(
                    posthog_distinct_id,
                    posthog_trace_id,
                    posthog_properties,
                    posthog_privacy_mode,
                    posthog_groups,
                    kwargs,
                    usage_stats,
                    latency,
                    content_blocks,
                    accumulated_content,
                    stop_reason=stop_reason,
                )

        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: TokenUsage,
        latency: float,
        content_blocks: List[StreamingContentBlock],
        accumulated_content: str,
        stop_reason: Optional[str] = None,
    ):
        from posthog.ai.types import StreamingEventData
        from posthog.ai.anthropic.anthropic_converter import (
            format_anthropic_streaming_input,
            format_anthropic_streaming_output_complete,
        )
        from posthog.ai.utils import capture_streaming_event

        # Prepare standardized event data
        formatted_input = format_anthropic_streaming_input(kwargs)
        sanitized_input = sanitize_anthropic(formatted_input)

        event_data = StreamingEventData(
            provider="anthropic",
            model=kwargs.get("model", "unknown"),
            base_url=str(self._client.base_url),
            kwargs=kwargs,
            formatted_input=sanitized_input,
            formatted_output=format_anthropic_streaming_output_complete(
                content_blocks, accumulated_content
            ),
            usage_stats=usage_stats,
            latency=latency,
            distinct_id=posthog_distinct_id,
            trace_id=posthog_trace_id,
            properties=posthog_properties,
            privacy_mode=posthog_privacy_mode,
            groups=posthog_groups,
            stop_reason=stop_reason,
        )

        # Use the common capture function
        capture_streaming_event(self._client._ph_client, event_data)