diff --git a/src/openlayer/lib/__init__.py b/src/openlayer/lib/__init__.py index 35ee2107..ee61dc3c 100644 --- a/src/openlayer/lib/__init__.py +++ b/src/openlayer/lib/__init__.py @@ -4,6 +4,8 @@ "init", "configure", "get_tracer_config", + "auto_instrument", + "unpatch_all", "trace", "trace_anthropic", "trace_openai", @@ -49,6 +51,9 @@ configure = tracer.configure get_tracer_config = tracer.get_tracer_config trace = tracer.trace + +# Auto-instrumentation entry points +from .integrations._auto import auto_instrument, unpatch_all # noqa: E402 trace_async = tracer.trace_async update_current_trace = tracer.update_current_trace update_current_step = tracer.update_current_step diff --git a/src/openlayer/lib/integrations/_auto.py b/src/openlayer/lib/integrations/_auto.py new file mode 100644 index 00000000..a793e02b --- /dev/null +++ b/src/openlayer/lib/integrations/_auto.py @@ -0,0 +1,258 @@ +"""Auto-instrumentation registry and orchestrator. + +This module powers ``openlayer.lib.init(auto_instrument=True)`` — when called, +it walks a registry of supported LLM SDK integrations, detects which ones are +installed (via ``importlib.util.find_spec``), and applies their patch function +so that every newly-constructed client is auto-traced. + +The two helper functions ``_patch_class_init`` and ``_unpatch_class_init`` are +the shared patch/unpatch primitives used by each ``_tracer.py`` module to +wrap an SDK's client class ``__init__``. They are idempotent and reversible. +""" + +from __future__ import annotations + +import functools +import importlib +import importlib.util +import logging +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union + +logger = logging.getLogger(__name__) + + +# Attribute marker placed on a client *class* once its __init__ has been +# wrapped. Distinct from `_openlayer_patched` which marks individual *instances* +# after a per-client wrap. Both are needed: the class marker prevents re-wrapping +# __init__, the instance marker prevents re-wrapping methods if the user mixes +# auto-instrument with an explicit trace_(client) call on the same instance. +_CLASS_PATCHED_ATTR = "_openlayer_class_patched" + + +def _patch_class_init( + cls: Type[Any], + wrap_fn: Callable[[Any], Any], +) -> None: + """Wrap ``cls.__init__`` so every newly-constructed instance is auto-traced. + + Idempotent: re-applying on an already-patched class is a no-op. The wrap + failure is caught and logged so a broken integration doesn't break client + construction. + """ + if getattr(cls, _CLASS_PATCHED_ATTR, False): + return + + original_init = cls.__init__ + + @functools.wraps(original_init) + def wrapped_init(self, *args, **kwargs): # type: ignore[no-untyped-def] + original_init(self, *args, **kwargs) + try: + wrap_fn(self) + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Openlayer: failed to auto-trace %s instance: %s", + cls.__name__, + e, + ) + + cls.__init__ = wrapped_init # type: ignore[method-assign] + setattr(cls, _CLASS_PATCHED_ATTR, True) + + +def _unpatch_class_init(cls: Type[Any]) -> None: + """Restore the original ``cls.__init__``. No-op if not patched.""" + if not getattr(cls, _CLASS_PATCHED_ATTR, False): + return + wrapped = cls.__init__ + original = getattr(wrapped, "__wrapped__", None) + if original is not None: + cls.__init__ = original # type: ignore[method-assign] + try: + delattr(cls, _CLASS_PATCHED_ATTR) + except AttributeError: + pass + + +@dataclass(frozen=True) +class IntegrationSpec: + """Declarative entry in the auto-instrumentation registry.""" + + name: str + """Public name users pass in ``auto_instrument=[...]``.""" + + probe: str + """Top-level package name used with ``importlib.util.find_spec`` to detect + whether the SDK is installed. Only the first dotted segment is probed.""" + + patch: Callable[[], None] + """Idempotent function that patches the SDK's client class(es).""" + + unpatch: Optional[Callable[[], None]] = None + """Optional restorer. Some integrations (litellm, portkey) don't have one + today; they're still registered so the patch path works.""" + + +# ----------------------------- Lazy import helpers ----------------------------- # +# +# The registry uses lambdas that import the tracer module on first call. This +# keeps the import-time cost of `openlayer.lib` low — we don't load every +# LLM SDK's tracer module just because the user did `from openlayer.lib import init`. + + +def _patch_via(module_name: str, attr: str) -> Callable[[], None]: + """Return a lazy patcher: imports ``.`` and calls its ``attr``.""" + + def _do_patch() -> None: + mod = importlib.import_module(f".{module_name}", package=__package__) + getattr(mod, attr)() + + return _do_patch + + +# ----------------------------- Registry ----------------------------- # + +REGISTRY: Tuple[IntegrationSpec, ...] = ( + IntegrationSpec( + "openai", + "openai", + _patch_via("openai_tracer", "_patch_openai"), + _patch_via("openai_tracer", "_unpatch_openai"), + ), + IntegrationSpec( + "anthropic", + "anthropic", + _patch_via("anthropic_tracer", "_patch_anthropic"), + _patch_via("anthropic_tracer", "_unpatch_anthropic"), + ), + IntegrationSpec( + "mistral", + "mistralai", + _patch_via("mistral_tracer", "_patch_mistral"), + _patch_via("mistral_tracer", "_unpatch_mistral"), + ), + IntegrationSpec( + "groq", + "groq", + _patch_via("groq_tracer", "_patch_groq"), + _patch_via("groq_tracer", "_unpatch_groq"), + ), + IntegrationSpec( + "gemini", + "google.generativeai", + _patch_via("gemini_tracer", "_patch_gemini"), + _patch_via("gemini_tracer", "_unpatch_gemini"), + ), + IntegrationSpec( + "oci", + "oci", + _patch_via("oci_tracer", "_patch_oci"), + _patch_via("oci_tracer", "_unpatch_oci"), + ), + IntegrationSpec( + "azure_content_understanding", + "azure.ai.contentunderstanding", + _patch_via("azure_content_understanding_tracer", "_patch_acu"), + _patch_via("azure_content_understanding_tracer", "_unpatch_acu"), + ), + IntegrationSpec( + "litellm", + "litellm", + _patch_via("litellm_tracer", "trace_litellm"), + None, + ), + IntegrationSpec( + "portkey", + "portkey_ai", + _patch_via("portkey_tracer", "trace_portkey"), + None, + ), + IntegrationSpec( + "google_adk", + "google.adk", + _patch_via("google_adk_tracer", "trace_google_adk"), + _patch_via("google_adk_tracer", "unpatch_google_adk"), + ), +) + + +_REGISTRY_BY_NAME: Dict[str, IntegrationSpec] = {s.name: s for s in REGISTRY} + + +def _is_installed(probe: str) -> bool: + """Return True if the full dotted ``probe`` path is importable. + + Uses ``find_spec`` so the SDK itself is not imported at probe time — only + the patch function imports it, and only when we've committed to patching. + The full path matters for namespace-package collisions: e.g. ``google`` + can be installed via ``google.generativeai`` without ``google.adk`` being + available, so we have to probe ``google.adk`` not just ``google``. + """ + try: + return importlib.util.find_spec(probe) is not None + except (ImportError, ValueError, ModuleNotFoundError): + return False + + +def auto_instrument( + targets: Union[bool, List[str]] = True, +) -> Dict[str, bool]: + """Detect and patch installed LLM SDKs so new client instances are auto-traced. + + Args: + targets: ``True`` (default) patches every installed supported SDK. + ``False`` is a no-op. A list of names patches only that subset + (e.g. ``["openai", "anthropic"]``). + + Returns: + Dict mapping integration name to a boolean: ``True`` if patched + successfully, ``False`` if skipped (not installed, or patch raised). + """ + if targets is False: + return {} + + if targets is True: + enabled = {s.name for s in REGISTRY} + else: + enabled = set(targets) + unknown = enabled - set(_REGISTRY_BY_NAME) + if unknown: + logger.warning( + "Openlayer: unknown auto_instrument targets: %s", + sorted(unknown), + ) + enabled &= set(_REGISTRY_BY_NAME) + + results: Dict[str, bool] = {} + for spec in REGISTRY: + if spec.name not in enabled: + continue + if not _is_installed(spec.probe): + results[spec.name] = False + logger.debug("Openlayer: skipped %s (not installed)", spec.name) + continue + try: + spec.patch() + results[spec.name] = True + logger.info("Openlayer: auto-instrumented %s", spec.name) + except Exception as e: # pylint: disable=broad-except + logger.warning("Openlayer: failed to instrument %s: %s", spec.name, e) + results[spec.name] = False + return results + + +def unpatch_all() -> None: + """Restore the original ``__init__`` for every patched integration. + + Integrations without an unpatch function are skipped. After this call, + previously-constructed (and already-traced) client instances are unaffected + — only future construction goes through the restored ``__init__``. + """ + for spec in REGISTRY: + if spec.unpatch is None: + continue + try: + spec.unpatch() + except Exception as e: # pylint: disable=broad-except + logger.warning("Openlayer: failed to unpatch %s: %s", spec.name, e) diff --git a/src/openlayer/lib/integrations/anthropic_tracer.py b/src/openlayer/lib/integrations/anthropic_tracer.py index 054a89b2..1caaca7c 100644 --- a/src/openlayer/lib/integrations/anthropic_tracer.py +++ b/src/openlayer/lib/integrations/anthropic_tracer.py @@ -53,7 +53,10 @@ def trace_anthropic( raise ImportError( "Anthropic library is not installed. Please install it with: pip install anthropic" ) - + + if getattr(client, "_openlayer_patched", False) is True: + return client + create_func = client.messages.create @wraps(create_func) @@ -76,9 +79,37 @@ def traced_create_func(*args, **kwargs): ) client.messages.create = traced_create_func + client._openlayer_patched = True return client +def _patch_anthropic() -> None: + """Patch ``anthropic.Anthropic`` (and async variant if available) ``__init__`` + so every newly-constructed instance is auto-traced. Idempotent.""" + if not HAVE_ANTHROPIC: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(anthropic.Anthropic, trace_anthropic) + # Async variant: patch only if it exists in this anthropic version. + async_cls = getattr(anthropic, "AsyncAnthropic", None) + if async_cls is not None: + _patch_class_init(async_cls, trace_anthropic) + + +def _unpatch_anthropic() -> None: + if not HAVE_ANTHROPIC: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(anthropic.Anthropic) + async_cls = getattr(anthropic, "AsyncAnthropic", None) + if async_cls is not None: + _unpatch_class_init(async_cls) + + def handle_streaming_create( create_func: callable, *args, diff --git a/src/openlayer/lib/integrations/async_openai_tracer.py b/src/openlayer/lib/integrations/async_openai_tracer.py index 85e5d4ed..3ab5dd5b 100644 --- a/src/openlayer/lib/integrations/async_openai_tracer.py +++ b/src/openlayer/lib/integrations/async_openai_tracer.py @@ -71,6 +71,9 @@ def trace_async_openai( if not HAVE_OPENAI: raise ImportError("OpenAI library is not installed. Please install it with: pip install openai") + if getattr(client, "_openlayer_patched", False) is True: + return client + is_azure_openai = isinstance(client, openai.AsyncAzureOpenAI) # Patch Chat Completions API @@ -174,6 +177,7 @@ async def traced_embeddings_create_func(*args, **kwargs): client.embeddings.create = traced_embeddings_create_func + client._openlayer_patched = True return client diff --git a/src/openlayer/lib/integrations/azure_content_understanding_tracer.py b/src/openlayer/lib/integrations/azure_content_understanding_tracer.py index 795de887..abf8325c 100644 --- a/src/openlayer/lib/integrations/azure_content_understanding_tracer.py +++ b/src/openlayer/lib/integrations/azure_content_understanding_tracer.py @@ -57,6 +57,9 @@ def trace_azure_content_understanding( "Please install it with: pip install azure-ai-contentunderstanding" ) + if getattr(client, "_openlayer_patched", False) is True: + return client + begin_analyze_func = client.begin_analyze @wraps(begin_analyze_func) @@ -109,9 +112,30 @@ def traced_result(*result_args, **result_kwargs): return poller client.begin_analyze = traced_begin_analyze + client._openlayer_patched = True return client +def _patch_acu() -> None: + """Patch ``azure.ai.contentunderstanding.ContentUnderstandingClient.__init__`` + so every newly-constructed client is auto-traced. Idempotent.""" + if not HAVE_AZURE_CONTENT_UNDERSTANDING: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(ContentUnderstandingClient, trace_azure_content_understanding) + + +def _unpatch_acu() -> None: + if not HAVE_AZURE_CONTENT_UNDERSTANDING: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(ContentUnderstandingClient) + + def _extract_usage_from_poller(poller: Any) -> Dict[str, Any]: """Extract UsageDetails from the LRO poller's final pipeline response. diff --git a/src/openlayer/lib/integrations/bedrock_tracer.py b/src/openlayer/lib/integrations/bedrock_tracer.py index a7834d65..5a9ad021 100644 --- a/src/openlayer/lib/integrations/bedrock_tracer.py +++ b/src/openlayer/lib/integrations/bedrock_tracer.py @@ -61,6 +61,9 @@ def trace_bedrock(client: "boto3.client") -> "boto3.client": "boto3 library is not installed. Please install it with: pip install boto3" ) + if getattr(client, "_openlayer_patched", False) is True: + return client + # Patch invoke_model for non-streaming requests invoke_model_func = client.invoke_model invoke_model_stream_func = client.invoke_model_with_response_stream @@ -95,6 +98,7 @@ def traced_invoke_model_stream(*args, **kwargs): client.invoke_model = traced_invoke_model client.invoke_model_with_response_stream = traced_invoke_model_stream + client._openlayer_patched = True return client diff --git a/src/openlayer/lib/integrations/gemini_tracer.py b/src/openlayer/lib/integrations/gemini_tracer.py index f6241cc6..f87c40bc 100644 --- a/src/openlayer/lib/integrations/gemini_tracer.py +++ b/src/openlayer/lib/integrations/gemini_tracer.py @@ -73,6 +73,9 @@ def trace_gemini( "Google Generative AI library is not installed. Please install it with: pip install google-generativeai" ) + if getattr(client, "_openlayer_patched", False) is True: + return client + # Store original methods original_generate_content = client.generate_content original_generate_content_async = client.generate_content_async @@ -125,9 +128,30 @@ async def traced_generate_content_async(*args, **kwargs): client.generate_content = traced_generate_content client.generate_content_async = traced_generate_content_async + client._openlayer_patched = True return client +def _patch_gemini() -> None: + """Patch ``google.generativeai.GenerativeModel.__init__`` so every newly- + constructed model is auto-traced. Idempotent.""" + if not HAVE_GEMINI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(genai.GenerativeModel, trace_gemini) + + +def _unpatch_gemini() -> None: + if not HAVE_GEMINI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(genai.GenerativeModel) + + def handle_streaming_generate( generate_func: callable, *args, diff --git a/src/openlayer/lib/integrations/groq_tracer.py b/src/openlayer/lib/integrations/groq_tracer.py index eb1403da..10771fc3 100644 --- a/src/openlayer/lib/integrations/groq_tracer.py +++ b/src/openlayer/lib/integrations/groq_tracer.py @@ -53,7 +53,10 @@ def trace_groq( raise ImportError( "Groq library is not installed. Please install it with: pip install groq" ) - + + if getattr(client, "_openlayer_patched", False) is True: + return client + create_func = client.chat.completions.create @wraps(create_func) @@ -76,9 +79,36 @@ def traced_create_func(*args, **kwargs): ) client.chat.completions.create = traced_create_func + client._openlayer_patched = True return client +def _patch_groq() -> None: + """Patch ``groq.Groq.__init__`` (and ``AsyncGroq`` if available) so every + newly-constructed instance is auto-traced. Idempotent.""" + if not HAVE_GROQ: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(groq.Groq, trace_groq) + async_cls = getattr(groq, "AsyncGroq", None) + if async_cls is not None: + _patch_class_init(async_cls, trace_groq) + + +def _unpatch_groq() -> None: + if not HAVE_GROQ: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(groq.Groq) + async_cls = getattr(groq, "AsyncGroq", None) + if async_cls is not None: + _unpatch_class_init(async_cls) + + def handle_streaming_create( create_func: callable, *args, diff --git a/src/openlayer/lib/integrations/mistral_tracer.py b/src/openlayer/lib/integrations/mistral_tracer.py index 5939c50e..b00d5d70 100644 --- a/src/openlayer/lib/integrations/mistral_tracer.py +++ b/src/openlayer/lib/integrations/mistral_tracer.py @@ -53,7 +53,10 @@ def trace_mistral( raise ImportError( "Mistral library is not installed. Please install it with: pip install mistralai" ) - + + if getattr(client, "_openlayer_patched", False) is True: + return client + stream_func = client.chat.stream create_func = client.chat.complete @@ -80,9 +83,30 @@ def traced_create_func(*args, **kwargs): client.chat.stream = traced_stream_func client.chat.complete = traced_create_func + client._openlayer_patched = True return client +def _patch_mistral() -> None: + """Patch ``mistralai.Mistral.__init__`` so every newly-constructed instance is + auto-traced. Idempotent.""" + if not HAVE_MISTRAL: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(mistralai.Mistral, trace_mistral) + + +def _unpatch_mistral() -> None: + if not HAVE_MISTRAL: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(mistralai.Mistral) + + def handle_streaming_create( create_func: callable, *args, diff --git a/src/openlayer/lib/integrations/oci_tracer.py b/src/openlayer/lib/integrations/oci_tracer.py index e74fd1a9..cc563d87 100644 --- a/src/openlayer/lib/integrations/oci_tracer.py +++ b/src/openlayer/lib/integrations/oci_tracer.py @@ -59,6 +59,9 @@ def trace_oci_genai( if not HAVE_OCI: raise ImportError("oci library is not installed. Please install it with: pip install oci") + if getattr(client, "_openlayer_patched", False) is True: + return client + chat_func = client.chat @wraps(chat_func) @@ -100,9 +103,30 @@ def traced_chat_func(*args, **kwargs): ) client.chat = traced_chat_func + client._openlayer_patched = True return client +def _patch_oci() -> None: + """Patch ``oci.generative_ai_inference.GenerativeAiInferenceClient.__init__`` + so every newly-constructed client is auto-traced. Idempotent.""" + if not HAVE_OCI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(GenerativeAiInferenceClient, trace_oci_genai) + + +def _unpatch_oci() -> None: + if not HAVE_OCI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(GenerativeAiInferenceClient) + + def handle_streaming_chat( response: Iterator[Any], chat_details: Any, diff --git a/src/openlayer/lib/integrations/openai_tracer.py b/src/openlayer/lib/integrations/openai_tracer.py index 46040052..ae412edc 100644 --- a/src/openlayer/lib/integrations/openai_tracer.py +++ b/src/openlayer/lib/integrations/openai_tracer.py @@ -70,6 +70,9 @@ def trace_openai( "OpenAI library is not installed. Please install it with: pip install openai" ) + if getattr(client, "_openlayer_patched", False) is True: + return client + is_azure_openai = isinstance(client, openai.AzureOpenAI) # Patch Chat Completions API @@ -187,9 +190,36 @@ def traced_embeddings_create_func(*args, **kwargs): client.embeddings.create = traced_embeddings_create_func + client._openlayer_patched = True return client +def _patch_openai() -> None: + """Patch ``openai.OpenAI`` / ``AzureOpenAI`` / ``AsyncOpenAI`` / ``AsyncAzureOpenAI`` + ``__init__`` so every newly-constructed instance is auto-traced. Idempotent.""" + if not HAVE_OPENAI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + from .async_openai_tracer import trace_async_openai + + _patch_class_init(openai.OpenAI, trace_openai) + _patch_class_init(openai.AzureOpenAI, trace_openai) + _patch_class_init(openai.AsyncOpenAI, trace_async_openai) + _patch_class_init(openai.AsyncAzureOpenAI, trace_async_openai) + + +def _unpatch_openai() -> None: + """Restore the original ``__init__`` on every OpenAI client class.""" + if not HAVE_OPENAI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + for cls in (openai.OpenAI, openai.AzureOpenAI, openai.AsyncOpenAI, openai.AsyncAzureOpenAI): + _unpatch_class_init(cls) + + def handle_streaming_create( create_func: callable, *args, diff --git a/src/openlayer/lib/tracing/tracer.py b/src/openlayer/lib/tracing/tracer.py index 4dd34a0c..91e4494d 100644 --- a/src/openlayer/lib/tracing/tracer.py +++ b/src/openlayer/lib/tracing/tracer.py @@ -147,6 +147,7 @@ def init( attachment_upload_enabled: Any = _UNSET, url_upload_enabled: Any = _UNSET, background_publish_enabled: Any = _UNSET, + auto_instrument: Union[bool, List[str]] = True, ) -> None: """Initialize and configure the Openlayer tracer. @@ -188,6 +189,14 @@ def init( platform has a durable copy. Defaults to False. background_publish_enabled: Publish traces from a background thread so the main thread returns immediately. Defaults to True. + auto_instrument: When truthy (default ``True``), detects every installed + supported LLM SDK (openai, anthropic, mistral, groq, gemini, oci, + azure_content_understanding, litellm, portkey, google_adk) and patches + it so newly-constructed clients are auto-traced. Pass ``False`` to skip + patching, or a list of names (e.g. ``["openai", "anthropic"]``) to + patch only that subset. This is a procedural argument — it is NOT + persisted across calls, and ``False`` does not unpatch (use + :func:`openlayer.lib.unpatch_all` for that). Examples: >>> from openlayer.lib import init, trace @@ -231,6 +240,13 @@ def init( reset_uploader() + # Run auto-instrumentation. The patchers themselves are idempotent, so + # repeated init() calls are safe. False short-circuits entirely. + if auto_instrument is not False: + from ..integrations._auto import auto_instrument as _run_auto_instrument + + _run_auto_instrument(auto_instrument) + # Track that the configure() deprecation log has been emitted once per process # so we don't spam logs on repeated calls. @@ -259,6 +275,10 @@ def configure(**kwargs: Any) -> None: "openlayer.lib.configure() is deprecated; use openlayer.lib.init() instead." ) _configure_deprecation_logged = True + # Preserve the historical configure() behavior of NOT auto-instrumenting — + # users who haven't migrated to init() should not be surprised by sudden + # client patching on SDK upgrade. + kwargs.setdefault("auto_instrument", False) init(**kwargs)