__init__.py

"""Openlayer lib."""

__all__ = [
    "init",
    "configure",
    "get_tracer_config",
    "auto_instrument",
    "unpatch_all",
    "trace",
    "trace_anthropic",
    "trace_openai",
    "trace_openai_assistant_thread_run",
    "trace_mistral",
    "trace_groq",
    "trace_async_openai",
    "trace_async",
    "trace_bedrock",
    "trace_azure_content_understanding",
    "trace_azure_speech",
    "trace_oci_genai",
    "trace_oci",  # Alias for backward compatibility
    "trace_litellm",
    "trace_portkey",
    "trace_google_adk",
    "unpatch_google_adk",
    "trace_claude_agent_sdk",
    "traced_claude_agent_sdk_query",
    "trace_gemini",
    "trace_google_genai",
    "update_current_trace",
    "update_current_step",
    # Offline buffer management functions
    "replay_buffered_traces",
    "get_buffer_status",
    "clear_offline_buffer",
    # User and session context functions
    "set_user_session_context",
    "update_trace_user_session",
    "get_current_user_id",
    "get_current_session_id",
    "clear_user_session_context",
]

# ---------------------------------- Tracing --------------------------------- #
from .tracing import tracer
from .tracing.context import (
    set_user_session_context,
    update_trace_user_session,
    get_current_user_id,
    get_current_session_id,
    clear_user_session_context,
)
from .integrations._auto import auto_instrument, unpatch_all

init = tracer.init
configure = tracer.configure
get_tracer_config = tracer.get_tracer_config
trace = tracer.trace
trace_async = tracer.trace_async
update_current_trace = tracer.update_current_trace
update_current_step = tracer.update_current_step

# Offline buffer management functions
replay_buffered_traces = tracer.replay_buffered_traces
get_buffer_status = tracer.get_buffer_status
clear_offline_buffer = tracer.clear_offline_buffer


def trace_anthropic(client):
    """Trace Anthropic chat completions."""
    # pylint: disable=import-outside-toplevel
    import anthropic

    from .integrations import anthropic_tracer

    if not isinstance(client, anthropic.Anthropic):
        raise ValueError("Invalid client. Please provide an Anthropic client.")
    return anthropic_tracer.trace_anthropic(client)


def trace_openai(client):
    """Trace OpenAI chat completions."""
    # pylint: disable=import-outside-toplevel
    import openai

    from .integrations import openai_tracer

    if not isinstance(client, (openai.Client, openai.AzureOpenAI)):
        raise ValueError("Invalid client. Please provide an OpenAI client.")
    return openai_tracer.trace_openai(client)


def trace_async_openai(client):
    """Trace OpenAI chat completions."""
    # pylint: disable=import-outside-toplevel
    import openai

    from .integrations import async_openai_tracer

    if not isinstance(client, (openai.AsyncOpenAI, openai.AsyncAzureOpenAI)):
        raise ValueError("Invalid client. Please provide an OpenAI client.")
    return async_openai_tracer.trace_async_openai(client)


def trace_openai_assistant_thread_run(client, run):
    """Trace OpenAI Assistant thread run."""
    # pylint: disable=import-outside-toplevel
    from .integrations import openai_tracer

    return openai_tracer.trace_openai_assistant_thread_run(client, run)


def trace_mistral(client):
    """Trace Mistral chat completions."""
    # pylint: disable=import-outside-toplevel
    import mistralai

    from .integrations import mistral_tracer

    if not isinstance(client, mistralai.Mistral):
        raise ValueError("Invalid client. Please provide a Mistral client.")
    return mistral_tracer.trace_mistral(client)


def trace_groq(client):
    """Trace Groq queries."""
    # pylint: disable=import-outside-toplevel
    import groq

    from .integrations import groq_tracer

    if not isinstance(client, groq.Groq):
        raise ValueError("Invalid client. Please provide a Groq client.")
    return groq_tracer.trace_groq(client)


def trace_bedrock(client):
    """Trace AWS Bedrock model invocations."""
    # pylint: disable=import-outside-toplevel
    try:
        import boto3
    except ImportError:
        raise ImportError("boto3 is required for Bedrock tracing. Install with: pip install boto3")

    from .integrations import bedrock_tracer

    # Check if it's a boto3 client for bedrock-runtime service
    if not hasattr(client, "_service_model") or client._service_model.service_name != "bedrock-runtime":
        raise ValueError("Invalid client. Please provide a boto3 bedrock-runtime client.")
    return bedrock_tracer.trace_bedrock(client)


def trace_azure_content_understanding(client):
    """Trace Azure Content Understanding analyses."""
    # pylint: disable=import-outside-toplevel
    try:
        from azure.ai.contentunderstanding import ContentUnderstandingClient
    except ImportError:
        raise ImportError(
            "azure-ai-contentunderstanding is required for Azure Content Understanding tracing. "
            "Install with: pip install azure-ai-contentunderstanding"
        )

    from .integrations import azure_content_understanding_tracer

    if not isinstance(client, ContentUnderstandingClient):
        raise ValueError("Invalid client. Please provide a ContentUnderstandingClient.")
    return azure_content_understanding_tracer.trace_azure_content_understanding(client)


def trace_azure_speech(client):
    """Trace Azure AI Speech recognition, translation and synthesis calls."""
    # pylint: disable=import-outside-toplevel
    try:
        import azure.cognitiveservices.speech  # noqa: F401
    except ImportError:
        raise ImportError(
            "azure-cognitiveservices-speech is required for Azure Speech tracing. "
            "Install with: pip install azure-cognitiveservices-speech"
        )

    from .integrations import azure_speech_tracer

    return azure_speech_tracer.trace_azure_speech(client)


def trace_oci_genai(client, estimate_tokens: bool = True):
    """Trace OCI GenAI chat completions.

    Args:
        client: OCI GenAI client.
        estimate_tokens: Whether to estimate tokens when not available. Defaults to True.
    """
    # pylint: disable=import-outside-toplevel
    try:
        import oci
    except ImportError:
        raise ImportError("oci is required for OCI GenAI tracing. Install with: pip install oci")

    from .integrations import oci_tracer

    if not isinstance(client, oci.generative_ai_inference.GenerativeAiInferenceClient):
        raise ValueError("Invalid client. Please provide an OCI GenAI client.")

    return oci_tracer.trace_oci_genai(client, estimate_tokens=estimate_tokens)


# --------------------------------- OCI GenAI -------------------------------- #
# Alias for backward compatibility
trace_oci = trace_oci_genai


# --------------------------------- LiteLLM ---------------------------------- #
def trace_litellm():
    """Enable tracing for LiteLLM completions.

    This function patches litellm.completion to automatically trace all completions
    made through the LiteLLM library, which provides a unified interface to 100+ LLM APIs.

    Example:
        >>> import litellm
        >>> from openlayer.lib import trace_litellm
        >>> # Enable tracing
        >>> trace_litellm()
        >>> # Use LiteLLM normally - tracing happens automatically
        >>> response = litellm.completion(
        ...     model="gpt-3.5-turbo",
        ...     messages=[{"role": "user", "content": "Hello!"}],
        ...     inference_id="custom-id-123",  # Optional Openlayer parameter
        ... )
    """
    # pylint: disable=import-outside-toplevel
    try:
        import litellm
    except ImportError:
        raise ImportError("litellm is required for LiteLLM tracing. Install with: pip install litellm")

    from .integrations import litellm_tracer

    return litellm_tracer.trace_litellm()


# ---------------------------------- Portkey ---------------------------------- #
def trace_portkey():
    """Enable tracing for Portkey completions.

    This function patches Portkey's chat.completions.create to automatically trace
    all OpenAI-compatible completions routed via the Portkey AI Gateway.

    Example:
        >>> from portkey_ai import Portkey
        >>> from openlayer.lib import trace_portkey
        >>> # Enable openlayer tracing for all Portkey completions
        >>> trace_portkey()
        >>> # Basic portkey client initialization
        >>> portkey = Portkey(
        >>>     api_key = os.environ['PORTKEY_API_KEY'],
        >>>     config = "YOUR_PORTKEY_CONFIG_ID", # optional your portkey config id
        >>> )
        >>> # use portkey normally - tracing happens automatically
        >>> response = portkey.chat.completions.create(
        >>> # model = "@YOUR_PORTKEY_SLUG/YOUR_MODEL_NAME", # optional if giving config
        >>>     messages = [
        >>>         {"role": "system", "content": "You are a helpful assistant."},
        >>>         {"role": "user", "content": "Write a poem on Argentina, least 100 words."}
        >>>     ]
        >>> )
    """
    # pylint: disable=import-outside-toplevel
    try:
        from portkey_ai import Portkey  # noqa: F401
    except ImportError:
        raise ImportError("portkey-ai is required for Portkey tracing. Install with: pip install portkey-ai")

    from .integrations import portkey_tracer

    return portkey_tracer.trace_portkey()


# ------------------------------ Google ADK ---------------------------------- #
def trace_google_adk(disable_adk_otel: bool = False):
    """Enable tracing for Google Agent Development Kit (ADK).

    This function patches Google ADK to automatically trace agent execution,
    LLM calls, and tool calls made through the ADK framework.

    By default, ADK's built-in OpenTelemetry tracing remains active, allowing
    you to send telemetry to both Google Cloud and Openlayer simultaneously.

    Args:
        disable_adk_otel: If True, disables ADK's built-in OpenTelemetry tracing.
            When False (default), ADK's OTel tracing works alongside Openlayer,
            allowing data to be sent to both Google Cloud (Cloud Trace, etc.)
            and Openlayer. Set to True only if you want Openlayer exclusively.

    Requirements:
        Google ADK and wrapt must be installed:
        pip install google-adk wrapt

    Example:
        >>> import os
        >>> os.environ["OPENLAYER_API_KEY"] = "your-api-key"
        >>> os.environ["OPENLAYER_INFERENCE_PIPELINE_ID"] = "your-pipeline-id"
        >>> from openlayer.lib import trace_google_adk
        >>>
        >>> # Enable tracing with both Google Cloud OTel and Openlayer (default)
        >>> trace_google_adk()
        >>>
        >>> # OR: Enable tracing with Openlayer only (disable ADK's OTel)
        >>> # trace_google_adk(disable_adk_otel=True)
        >>>
        >>> # Now create and run your ADK agents
        >>> from google.adk.agents import Agent
        >>> agent = Agent(name="Assistant", model="gemini-2.5-flash")
        >>> result = await agent.run_async(...)
    """
    # pylint: disable=import-outside-toplevel
    from .integrations import google_adk_tracer

    return google_adk_tracer.trace_google_adk(disable_adk_otel=disable_adk_otel)


def unpatch_google_adk():
    """Remove Google ADK tracing patches.

    This function restores Google ADK's original behavior by removing all
    Openlayer instrumentation.

    Example:
        >>> from openlayer.lib import unpatch_google_adk
        >>> unpatch_google_adk()
    """
    # pylint: disable=import-outside-toplevel
    from .integrations import google_adk_tracer

    return google_adk_tracer.unpatch_google_adk()


# ------------------------------ Claude Agent SDK ---------------------------- #
def trace_claude_agent_sdk(
    *,
    inference_pipeline_id=None,
    truncate_tool_output_chars: int = 8192,
    capture_thinking: bool = True,
    redact_mcp_env: bool = True,
):
    """Enable Openlayer tracing for the Claude Agent SDK.

    Monkey-patches ``claude_agent_sdk.query`` and ``ClaudeSDKClient`` so every
    call becomes an Openlayer trace with nested steps for assistant turns,
    tool calls (including MCP and subagent calls), session metadata, cost,
    and tokens.

    Requirements:
        ``claude-agent-sdk>=0.1.81`` must be installed:
        ``pip install 'claude-agent-sdk>=0.1.81'``

    Args:
        inference_pipeline_id: Optional Openlayer inference pipeline ID. Falls
            back to the ``OPENLAYER_INFERENCE_PIPELINE_ID`` env var.
        truncate_tool_output_chars: Maximum characters of tool output to
            capture per TOOL step. Defaults to 8192.
        capture_thinking: Whether to capture ``ThinkingBlock`` content into
            chat-completion step metadata. Defaults to True.
        redact_mcp_env: Whether to strip ``env`` and ``headers`` from MCP
            server config dicts in trace metadata. Defaults to True.

    Example:
        >>> import os
        >>> os.environ["OPENLAYER_API_KEY"] = "..."
        >>> os.environ["OPENLAYER_INFERENCE_PIPELINE_ID"] = "..."
        >>> os.environ["ANTHROPIC_API_KEY"] = "..."
        >>> from openlayer.lib import trace_claude_agent_sdk
        >>> trace_claude_agent_sdk()
        >>>
        >>> from claude_agent_sdk import query, ClaudeAgentOptions
        >>> async for m in query(prompt="hello", options=ClaudeAgentOptions(model="claude-haiku-4-5")):
        ...     ...
    """
    # pylint: disable=import-outside-toplevel
    from .integrations import claude_agent_sdk as _integration

    return _integration.trace_claude_agent_sdk(
        inference_pipeline_id=inference_pipeline_id,
        truncate_tool_output_chars=truncate_tool_output_chars,
        capture_thinking=capture_thinking,
        redact_mcp_env=redact_mcp_env,
    )


def traced_claude_agent_sdk_query(*, prompt, options=None, inference_pipeline_id=None, **kwargs):
    """Per-call wrapper around ``claude_agent_sdk.query()`` (alternative to global init).

    Returns an async iterator that yields the same messages as ``query()`` while
    emitting an Openlayer trace as a side effect.

    Example:
        >>> from openlayer.lib import traced_claude_agent_sdk_query
        >>> async for m in traced_claude_agent_sdk_query(prompt="hello"):
        ...     ...
    """
    # pylint: disable=import-outside-toplevel
    from .integrations import claude_agent_sdk as _integration

    return _integration.traced_query(
        prompt=prompt,
        options=options,
        inference_pipeline_id=inference_pipeline_id,
        **kwargs,
    )


# -------------------------------- Google Gemini --------------------------------- #
def _legacy_gemini_model_class():
    """``google.generativeai.GenerativeModel``, or None if not installed."""
    # pylint: disable=import-outside-toplevel
    try:
        import google.generativeai as legacy_genai

        return legacy_genai.GenerativeModel
    except ImportError:
        return None


def _google_genai_client_class():
    """``google.genai.Client``, or None if not installed."""
    # pylint: disable=import-outside-toplevel
    try:
        from google import genai as google_genai

        return google_genai.Client
    except ImportError:
        return None


def trace_gemini(client):
    """Trace Google Gemini chat completions for either Google SDK.

    Accepts both Gemini clients and dispatches to the matching tracer:

    - ``google.genai.Client`` (package ``google-genai``, the current unified SDK,
      including ``vertexai=True`` mode) -> ``trace_google_genai``
    - ``google.generativeai.GenerativeModel`` (package ``google-generativeai``,
      legacy / maintenance mode) -> the legacy Gemini tracer

    Only the SDK matching the client passed in is imported, so having just one of
    the two installed is fine.

    Example:
        >>> from google import genai
        >>> from openlayer.lib import trace_gemini
        >>> client = trace_gemini(genai.Client(vertexai=True, project="p", location="us-central1"))
        >>> client.models.generate_content(model="gemini-2.5-flash", contents="hi")
    """
    # pylint: disable=import-outside-toplevel
    legacy_class = _legacy_gemini_model_class()
    unified_class = _google_genai_client_class()

    if legacy_class is not None and isinstance(client, legacy_class):
        from .integrations import gemini_tracer

        return gemini_tracer.trace_gemini(client)

    if unified_class is not None and isinstance(client, unified_class):
        from .integrations import google_genai_tracer

        return google_genai_tracer.trace_google_genai(client)

    if legacy_class is None and unified_class is None:
        raise ImportError(
            "Gemini tracing requires either google-genai (recommended) or "
            "google-generativeai. Install with: pip install google-genai"
        )

    raise ValueError(
        "Invalid client. Please provide a google.genai.Client (package "
        "google-genai) or a google.generativeai.GenerativeModel (package "
        f"google-generativeai). Got: {type(client).__module__}.{type(client).__qualname__}"
    )


def trace_google_genai(client):
    """Trace Google Gen AI (``google-genai``) content generation.

    Covers ``client.models.generate_content`` and ``generate_content_stream``,
    their async equivalents under ``client.aio.models``, and chat sessions via
    ``client.chats`` — in both Vertex AI and AI Studio mode.

    Example:
        >>> from google import genai
        >>> from openlayer.lib import trace_google_genai
        >>> client = trace_google_genai(genai.Client(vertexai=True, project="p", location="us-central1"))
        >>> client.models.generate_content(model="gemini-2.5-flash", contents="hi")
    """
    # pylint: disable=import-outside-toplevel
    unified_class = _google_genai_client_class()
    if unified_class is None:
        raise ImportError("google-genai is required for Google Gen AI tracing. Install with: pip install google-genai")

    if not isinstance(client, unified_class):
        raise ValueError(
            "Invalid client. Please provide a google.genai.Client instance. Got: "
            f"{type(client).__module__}.{type(client).__qualname__}"
        )

    from .integrations import google_genai_tracer

    return google_genai_tracer.trace_google_genai(client)