insights-python/hanzo_insights/ai/openai/openai.py

import time
import uuid
from typing import Any, Dict, List, Optional

from hanzo_insights.ai.types import TokenUsage

try:
    import openai
except ImportError:
    raise ModuleNotFoundError(
        "Please install the OpenAI SDK to use this feature: 'pip install openai'"
    )

from hanzo_insights.ai.utils import (
    call_llm_and_track_usage,
    extract_available_tool_calls,
    merge_usage_stats,
    with_privacy_mode,
)
from hanzo_insights.ai.openai.openai_converter import (
    extract_openai_usage_from_chunk,
    extract_openai_content_from_chunk,
    extract_openai_tool_calls_from_chunk,
    accumulate_openai_tool_calls,
)
from hanzo_insights.ai.sanitization import sanitize_openai, sanitize_openai_response
from hanzo_insights.client import Client as InsightsClient
from hanzo_insights import setup


class OpenAI(openai.OpenAI):
    """
    A wrapper around the OpenAI SDK that automatically sends LLM usage events to Insights.
    """

    _ph_client: InsightsClient

    def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
        """
        Args:
            api_key: OpenAI API key.
            insights_client: If provided, events will be captured via this client instead of the global client.
            **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
        """

        super().__init__(**kwargs)
        self._ph_client = insights_client or setup()

        # Store original objects after parent initialization (only if they exist)
        self._original_chat = getattr(self, "chat", None)
        self._original_embeddings = getattr(self, "embeddings", None)
        self._original_beta = getattr(self, "beta", None)
        self._original_responses = getattr(self, "responses", None)

        # Replace with wrapped versions (only if originals exist)
        if self._original_chat is not None:
            self.chat = WrappedChat(self, self._original_chat)

        if self._original_embeddings is not None:
            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)

        if self._original_beta is not None:
            self.beta = WrappedBeta(self, self._original_beta)

        if self._original_responses is not None:
            self.responses = WrappedResponses(self, self._original_responses)


class WrappedResponses:
    """Wrapper for OpenAI responses that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_responses):
        self._client = client
        self._original = original_responses

    def __getattr__(self, name):
        """Fallback to original responses object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    def create(
        self,
        insights_distinct_id: Optional[str] = None,
        insights_trace_id: Optional[str] = None,
        insights_properties: Optional[Dict[str, Any]] = None,
        insights_privacy_mode: bool = False,
        insights_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        if insights_trace_id is None:
            insights_trace_id = str(uuid.uuid4())

        if kwargs.get("stream", False):
            return self._create_streaming(
                insights_distinct_id,
                insights_trace_id,
                insights_properties,
                insights_privacy_mode,
                insights_groups,
                **kwargs,
            )

        return call_llm_and_track_usage(
            insights_distinct_id,
            self._client._ph_client,
            "openai",
            insights_trace_id,
            insights_properties,
            insights_privacy_mode,
            insights_groups,
            self._client.base_url,
            self._original.create,
            **kwargs,
        )

    def _create_streaming(
        self,
        insights_distinct_id: Optional[str],
        insights_trace_id: Optional[str],
        insights_properties: Optional[Dict[str, Any]],
        insights_privacy_mode: bool,
        insights_groups: Optional[Dict[str, Any]],
        **kwargs: Any,
    ):
        start_time = time.time()
        usage_stats: TokenUsage = TokenUsage()
        final_content = []
        model_from_response: Optional[str] = None
        response = self._original.create(**kwargs)

        def generator():
            nonlocal usage_stats
            nonlocal final_content  # noqa: F824
            nonlocal model_from_response

            try:
                for chunk in response:
                    # Extract model from response object in chunk (for stored prompts)
                    if hasattr(chunk, "response") and chunk.response:
                        if model_from_response is None and hasattr(
                            chunk.response, "model"
                        ):
                            model_from_response = chunk.response.model

                    # Extract usage stats from chunk
                    chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")

                    if chunk_usage:
                        merge_usage_stats(usage_stats, chunk_usage)

                    # Extract content from chunk
                    content = extract_openai_content_from_chunk(chunk, "responses")

                    if content is not None:
                        final_content.append(content)

                    yield chunk

            finally:
                end_time = time.time()
                latency = end_time - start_time
                output = final_content
                self._capture_streaming_event(
                    insights_distinct_id,
                    insights_trace_id,
                    insights_properties,
                    insights_privacy_mode,
                    insights_groups,
                    kwargs,
                    usage_stats,
                    latency,
                    output,
                    None,  # Responses API doesn't have tools
                    model_from_response,
                )

        return generator()

    def _capture_streaming_event(
        self,
        insights_distinct_id: Optional[str],
        insights_trace_id: Optional[str],
        insights_properties: Optional[Dict[str, Any]],
        insights_privacy_mode: bool,
        insights_groups: Optional[Dict[str, Any]],
        kwargs: Dict[str, Any],
        usage_stats: TokenUsage,
        latency: float,
        output: Any,
        available_tool_calls: Optional[List[Dict[str, Any]]] = None,
        model_from_response: Optional[str] = None,
    ):
        from hanzo_insights.ai.types import StreamingEventData
        from hanzo_insights.ai.openai.openai_converter import (
            format_openai_streaming_input,
            format_openai_streaming_output,
        )
        from hanzo_insights.ai.utils import capture_streaming_event

        # Prepare standardized event data
        formatted_input = format_openai_streaming_input(kwargs, "responses")
        sanitized_input = sanitize_openai_response(formatted_input)

        # Use model from kwargs, fallback to model from response
        model = kwargs.get("model") or model_from_response or "unknown"

        event_data = StreamingEventData(
            provider="openai",
            model=model,
            base_url=str(self._client.base_url),
            kwargs=kwargs,
            formatted_input=sanitized_input,
            formatted_output=format_openai_streaming_output(output, "responses"),
            usage_stats=usage_stats,
            latency=latency,
            distinct_id=insights_distinct_id,
            trace_id=insights_trace_id,
            properties=insights_properties,
            privacy_mode=insights_privacy_mode,
            groups=insights_groups,
        )

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

    def parse(
        self,
        insights_distinct_id: Optional[str] = None,
        insights_trace_id: Optional[str] = None,
        insights_properties: Optional[Dict[str, Any]] = None,
        insights_privacy_mode: bool = False,
        insights_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        """
        Parse structured output using OpenAI's 'responses.parse' method, but also track usage in Insights.

        Args:
            insights_distinct_id: Optional ID to associate with the usage event.
            insights_trace_id: Optional trace UUID for linking events.
            insights_properties: Optional dictionary of extra properties to include in the event.
            insights_privacy_mode: Whether to anonymize the input and output.
            insights_groups: Optional dictionary of groups to associate with the event.
            **kwargs: Any additional parameters for the OpenAI Responses Parse API.

        Returns:
            The response from OpenAI's responses.parse call.
        """
        return call_llm_and_track_usage(
            insights_distinct_id,
            self._client._ph_client,
            "openai",
            insights_trace_id,
            insights_properties,
            insights_privacy_mode,
            insights_groups,
            self._client.base_url,
            self._original.parse,
            **kwargs,
        )


class WrappedChat:
    """Wrapper for OpenAI chat that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_chat):
        self._client = client
        self._original = original_chat

    def __getattr__(self, name):
        """Fallback to original chat object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    @property
    def completions(self):
        return WrappedCompletions(self._client, self._original.completions)


class WrappedCompletions:
    """Wrapper for OpenAI chat completions that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_completions):
        self._client = client
        self._original = original_completions

    def __getattr__(self, name):
        """Fallback to original completions object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    def create(
        self,
        insights_distinct_id: Optional[str] = None,
        insights_trace_id: Optional[str] = None,
        insights_properties: Optional[Dict[str, Any]] = None,
        insights_privacy_mode: bool = False,
        insights_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        if insights_trace_id is None:
            insights_trace_id = str(uuid.uuid4())

        if kwargs.get("stream", False):
            return self._create_streaming(
                insights_distinct_id,
                insights_trace_id,
                insights_properties,
                insights_privacy_mode,
                insights_groups,
                **kwargs,
            )

        return call_llm_and_track_usage(
            insights_distinct_id,
            self._client._ph_client,
            "openai",
            insights_trace_id,
            insights_properties,
            insights_privacy_mode,
            insights_groups,
            self._client.base_url,
            self._original.create,
            **kwargs,
        )

    def _create_streaming(
        self,
        insights_distinct_id: Optional[str],
        insights_trace_id: Optional[str],
        insights_properties: Optional[Dict[str, Any]],
        insights_privacy_mode: bool,
        insights_groups: Optional[Dict[str, Any]],
        **kwargs: Any,
    ):
        start_time = time.time()
        usage_stats: TokenUsage = TokenUsage()
        accumulated_content = []
        accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
        model_from_response: Optional[str] = None
        if "stream_options" not in kwargs:
            kwargs["stream_options"] = {}
        kwargs["stream_options"]["include_usage"] = True
        response = self._original.create(**kwargs)

        def generator():
            nonlocal usage_stats
            nonlocal accumulated_content  # noqa: F824
            nonlocal accumulated_tool_calls
            nonlocal model_from_response

            try:
                for chunk in response:
                    # Extract model from chunk (Chat Completions chunks have model field)
                    if model_from_response is None and hasattr(chunk, "model"):
                        model_from_response = chunk.model

                    # Extract usage stats from chunk
                    chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")

                    if chunk_usage:
                        merge_usage_stats(usage_stats, chunk_usage)

                    # Extract content from chunk
                    content = extract_openai_content_from_chunk(chunk, "chat")

                    if content is not None:
                        accumulated_content.append(content)

                    # Extract and accumulate tool calls from chunk
                    chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk)
                    if chunk_tool_calls:
                        accumulate_openai_tool_calls(
                            accumulated_tool_calls, chunk_tool_calls
                        )

                    yield chunk

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

                # Convert accumulated tool calls dict to list
                tool_calls_list = (
                    list(accumulated_tool_calls.values())
                    if accumulated_tool_calls
                    else None
                )

                self._capture_streaming_event(
                    insights_distinct_id,
                    insights_trace_id,
                    insights_properties,
                    insights_privacy_mode,
                    insights_groups,
                    kwargs,
                    usage_stats,
                    latency,
                    accumulated_content,
                    tool_calls_list,
                    extract_available_tool_calls("openai", kwargs),
                    model_from_response,
                )

        return generator()

    def _capture_streaming_event(
        self,
        insights_distinct_id: Optional[str],
        insights_trace_id: Optional[str],
        insights_properties: Optional[Dict[str, Any]],
        insights_privacy_mode: bool,
        insights_groups: Optional[Dict[str, Any]],
        kwargs: Dict[str, Any],
        usage_stats: TokenUsage,
        latency: float,
        output: Any,
        tool_calls: Optional[List[Dict[str, Any]]] = None,
        available_tool_calls: Optional[List[Dict[str, Any]]] = None,
        model_from_response: Optional[str] = None,
    ):
        from hanzo_insights.ai.types import StreamingEventData
        from hanzo_insights.ai.openai.openai_converter import (
            format_openai_streaming_input,
            format_openai_streaming_output,
        )
        from hanzo_insights.ai.utils import capture_streaming_event

        # Prepare standardized event data
        formatted_input = format_openai_streaming_input(kwargs, "chat")
        sanitized_input = sanitize_openai(formatted_input)

        # Use model from kwargs, fallback to model from response
        model = kwargs.get("model") or model_from_response or "unknown"

        event_data = StreamingEventData(
            provider="openai",
            model=model,
            base_url=str(self._client.base_url),
            kwargs=kwargs,
            formatted_input=sanitized_input,
            formatted_output=format_openai_streaming_output(output, "chat", tool_calls),
            usage_stats=usage_stats,
            latency=latency,
            distinct_id=insights_distinct_id,
            trace_id=insights_trace_id,
            properties=insights_properties,
            privacy_mode=insights_privacy_mode,
            groups=insights_groups,
        )

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


class WrappedEmbeddings:
    """Wrapper for OpenAI embeddings that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_embeddings):
        self._client = client
        self._original = original_embeddings

    def __getattr__(self, name):
        """Fallback to original embeddings object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    def create(
        self,
        insights_distinct_id: Optional[str] = None,
        insights_trace_id: Optional[str] = None,
        insights_properties: Optional[Dict[str, Any]] = None,
        insights_privacy_mode: bool = False,
        insights_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        """
        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in Insights.

        Args:
            insights_distinct_id: Optional ID to associate with the usage event.
            insights_trace_id: Optional trace UUID for linking events.
            insights_properties: Optional dictionary of extra properties to include in the event.
            insights_privacy_mode: Whether to anonymize the input and output.
            insights_groups: Optional dictionary of groups to associate with the event.
            **kwargs: Any additional parameters for the OpenAI Embeddings API.

        Returns:
            The response from OpenAI's embeddings.create call.
        """

        if insights_trace_id is None:
            insights_trace_id = str(uuid.uuid4())

        start_time = time.time()
        response = self._original.create(**kwargs)
        end_time = time.time()

        # Extract usage statistics if available
        usage_stats = {}
        if hasattr(response, "usage") and response.usage:
            usage_stats = {
                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
                "total_tokens": getattr(response.usage, "total_tokens", 0),
            }

        latency = end_time - start_time

        # Build the event properties
        event_properties = {
            "$ai_provider": "openai",
            "$ai_model": kwargs.get("model"),
            "$ai_input": with_privacy_mode(
                self._client._ph_client,
                insights_privacy_mode,
                sanitize_openai_response(kwargs.get("input")),
            ),
            "$ai_http_status": 200,
            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
            "$ai_latency": latency,
            "$ai_trace_id": insights_trace_id,
            "$ai_base_url": str(self._client.base_url),
            **(insights_properties or {}),
        }

        if insights_distinct_id is None:
            event_properties["$process_person_profile"] = False

        # Send capture event for embeddings
        if hasattr(self._client._ph_client, "capture"):
            self._client._ph_client.capture(
                distinct_id=insights_distinct_id or insights_trace_id,
                event="$ai_embedding",
                properties=event_properties,
                groups=insights_groups,
            )

        return response


class WrappedBeta:
    """Wrapper for OpenAI beta features that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_beta):
        self._client = client
        self._original = original_beta

    def __getattr__(self, name):
        """Fallback to original beta object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    @property
    def chat(self):
        return WrappedBetaChat(self._client, self._original.chat)


class WrappedBetaChat:
    """Wrapper for OpenAI beta chat that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_beta_chat):
        self._client = client
        self._original = original_beta_chat

    def __getattr__(self, name):
        """Fallback to original beta chat object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    @property
    def completions(self):
        return WrappedBetaCompletions(self._client, self._original.completions)


class WrappedBetaCompletions:
    """Wrapper for OpenAI beta chat completions that tracks usage in Insights."""

    def __init__(self, client: OpenAI, original_beta_completions):
        self._client = client
        self._original = original_beta_completions

    def __getattr__(self, name):
        """Fallback to original beta completions object for any methods we don't explicitly handle."""
        return getattr(self._original, name)

    def parse(
        self,
        insights_distinct_id: Optional[str] = None,
        insights_trace_id: Optional[str] = None,
        insights_properties: Optional[Dict[str, Any]] = None,
        insights_privacy_mode: bool = False,
        insights_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        return call_llm_and_track_usage(
            insights_distinct_id,
            self._client._ph_client,
            "openai",
            insights_trace_id,
            insights_properties,
            insights_privacy_mode,
            insights_groups,
            self._client.base_url,
            self._original.parse,
            **kwargs,
        )