gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/instrumentation/__init__.py

@@ -0,0 +1,300 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Instrumentation module for gatewaysdk.
+
+Provides automatic metric collection and tracing for various libraries.
+Uses a plugin architecture with an InstrumentorRegistry.
+
+Usage:
+    # Enable specific instrumentors
+    gatewaysdk.instrument("openai")
+    gatewaysdk.instrument("openai", "anthropic", "langchain")
+
+    # Enable all available instrumentors
+    gatewaysdk.instrument_all()
+
+    # Disable instrumentors
+    gatewaysdk.uninstrument("openai")
+    gatewaysdk.uninstrument_all()
+"""
+
+import warnings
+from typing import List, Optional
+
+# Export base classes and registry
+from .base import Instrumentor, MetricSchema
+from .registry import (
+    InstrumentorRegistry,
+)
+from .registry import instrument as registry_instrument
+from .registry import uninstrument as registry_uninstrument
+
+__all__ = [
+    "Instrumentor",
+    "MetricSchema",
+    "InstrumentorRegistry",
+    "instrument",
+    "uninstrument",
+    "instrument_all",
+    "uninstrument_all",
+]
+
+# Legacy instrumentor availability flags
+AGENTOPS_INSTALLED: bool = False
+AGENTOPS_LANGCHAIN_INSTALLED: bool = False
+LITELLM_INSTALLED: bool = False
+VLLM_INSTALLED: bool = False
+WEAVE_INSTALLED: bool = False
+OPENAI_OTEL_INSTALLED: bool = False
+
+try:
+    from . import agentops  # type: ignore
+
+    AGENTOPS_INSTALLED = True  # type: ignore
+except ImportError:
+    pass
+
+try:
+    from . import litellm  # type: ignore
+
+    LITELLM_INSTALLED = True  # type: ignore
+except ImportError:
+    pass
+
+try:
+    from opentelemetry.instrumentation.openai import OpenAIInstrumentor  # type: ignore
+
+    OPENAI_OTEL_INSTALLED = True
+except ImportError:
+    pass
+
+
+def instrument(*names: str, **kwargs) -> List[Optional[Instrumentor]]:
+    """
+    Activate one or more instrumentors by name.
+
+    Uses the InstrumentorRegistry to activate instrumentors.
+    For legacy instrumentors not yet migrated to the registry,
+    falls back to direct activation.
+
+    Args:
+        *names: Instrumentor names to activate
+        **kwargs: Additional arguments passed to instrumentor constructors
+
+    Returns:
+        List of activated instrumentor instances (None for failed activations)
+
+    Examples:
+        # Enable OpenAI instrumentation
+        gatewaysdk.instrument("openai")
+
+        # Enable multiple instrumentors
+        gatewaysdk.instrument("openai", "anthropic", "langchain")
+
+        # Enable all available instrumentors
+        gatewaysdk.instrument_all()
+    """
+    results = []
+    for name in names:
+        # Try registry first
+        try:
+            result = InstrumentorRegistry.instrument(name, **kwargs)
+            results.append(result)
+        except ValueError:
+            # Fall back to legacy instrumentors
+            result = _instrument_legacy(name)
+            results.append(result)
+    return results
+
+
+def uninstrument(*names: str) -> None:
+    """
+    Deactivate one or more instrumentors by name.
+
+    Args:
+        *names: Instrumentor names to deactivate
+    """
+    for name in names:
+        # Try registry first
+        if InstrumentorRegistry.is_active(name):
+            InstrumentorRegistry.uninstrument(name)
+        else:
+            # Fall back to legacy uninstrument
+            _uninstrument_legacy(name)
+
+
+def _instrument_legacy(name: str) -> Optional[Instrumentor]:
+    """Activate a legacy instrumentor by name."""
+    name_lower = name.lower()
+
+    if name_lower == "agentops" and AGENTOPS_INSTALLED:
+        from .agentops import instrument_agentops
+
+        instrument_agentops()
+        return None  # Legacy doesn't return instance
+
+    if name_lower == "litellm" and LITELLM_INSTALLED:
+        from .litellm import instrument_litellm
+
+        instrument_litellm()
+        return None
+
+    if name_lower == "vllm" and VLLM_INSTALLED:
+        from .vllm import instrument_vllm
+
+        instrument_vllm()
+        return None
+
+    if name_lower in ("openai", "openai-otel") and OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().instrument()  # type: ignore
+        except Exception as e:
+            warnings.warn(f"Failed to instrument OpenAI SDK: {e}")
+        return None
+
+    warnings.warn(f"Unknown or unavailable instrumentor: {name}")
+    return None
+
+
+def _uninstrument_legacy(name: str) -> None:
+    """Deactivate a legacy instrumentor by name."""
+    name_lower = name.lower()
+
+    if name_lower == "agentops" and AGENTOPS_INSTALLED:
+        try:
+            from .agentops import uninstrument_agentops
+
+            uninstrument_agentops()
+        except ImportError:
+            pass
+
+    elif name_lower == "litellm" and LITELLM_INSTALLED:
+        try:
+            from .litellm import uninstrument_litellm
+
+            uninstrument_litellm()
+        except ImportError:
+            pass
+
+    elif name_lower == "vllm" and VLLM_INSTALLED:
+        try:
+            from .vllm import uninstrument_vllm
+
+            uninstrument_vllm()
+        except ImportError:
+            pass
+
+    elif name_lower in ("openai", "openai-otel") and OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().uninstrument()  # type: ignore
+        except Exception:
+            pass
+
+
+def instrument_all():
+    """
+    Instrument all available instrumentation libraries.
+
+    This includes both registry-based instrumentors and legacy instrumentors.
+    """
+    # First, activate all registry-based instrumentors
+    for name in InstrumentorRegistry.list_available():
+        try:
+            InstrumentorRegistry.instrument(name)
+        except Exception as e:
+            warnings.warn(f"Failed to instrument {name}: {e}")
+
+    # Then activate legacy instrumentors
+    if AGENTOPS_INSTALLED:
+        from .agentops import instrument_agentops
+
+        instrument_agentops()
+    else:
+        warnings.warn("agentops is not installed. It's therefore not instrumented.")
+
+    if LITELLM_INSTALLED:
+        from .litellm import instrument_litellm
+
+        instrument_litellm()
+    else:
+        warnings.warn("litellm is not installed. It's therefore not instrumented.")
+
+    if VLLM_INSTALLED:
+        from .vllm import instrument_vllm
+
+        instrument_vllm()
+    else:
+        warnings.warn("vllm is not installed. It's therefore not instrumented.")
+
+    if AGENTOPS_LANGCHAIN_INSTALLED:
+        from .agentops_langchain import instrument_agentops_langchain
+
+        instrument_agentops_langchain()
+    else:
+        warnings.warn("Agentops-langchain integration is not installed. It's therefore not instrumented.")
+
+    # Enable standard OpenAI instrumentation if available
+    # This allows client-side spans to be generated for OpenAI SDK calls
+    if OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().instrument()  # type: ignore
+        except Exception as e:
+            warnings.warn(f"Failed to instrument OpenAI SDK: {e}")
+
+
+def uninstrument_all():
+    """
+    Uninstrument all active instrumentation libraries.
+
+    This includes both registry-based instrumentors and legacy instrumentors.
+    """
+    # First, deactivate all registry-based instrumentors
+    InstrumentorRegistry.uninstrument_all()
+
+    # Then deactivate legacy instrumentors
+    if AGENTOPS_INSTALLED:
+        try:
+            from .agentops import uninstrument_agentops
+
+            uninstrument_agentops()
+        except ImportError:
+            warnings.warn("agentops is installed but uninstrument_agentops could not be imported.")
+    else:
+        warnings.warn("agentops is not installed. It's therefore not uninstrumented.")
+
+    if LITELLM_INSTALLED:
+        try:
+            from .litellm import uninstrument_litellm
+
+            uninstrument_litellm()
+        except ImportError:
+            warnings.warn("litellm is installed but uninstrument_litellm could not be imported.")
+    else:
+        warnings.warn("litellm is not installed. It's therefore not uninstrumented.")
+
+    if VLLM_INSTALLED:
+        try:
+            from .vllm import uninstrument_vllm
+
+            uninstrument_vllm()
+        except ImportError:
+            warnings.warn("vllm is installed but uninstrument_vllm could not be imported.")
+    else:
+        warnings.warn("vllm is not installed. It's therefore not uninstrumented.")
+
+    if AGENTOPS_LANGCHAIN_INSTALLED:
+        try:
+            from .agentops_langchain import uninstrument_agentops_langchain
+
+            uninstrument_agentops_langchain()
+        except ImportError:
+            warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
+    else:
+        warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")
+
+    if OPENAI_OTEL_INSTALLED:
+        try:
+            OpenAIInstrumentor().uninstrument()  # type: ignore
+        except Exception:
+            pass

gatewaysdk/instrumentation/agentops.py

@@ -0,0 +1,314 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Callable, no_type_check
+
+import requests
+from agentops.client.api import V3Client, V4Client
+from agentops.client.api.types import AuthTokenResponse
+from agentops.sdk.exporters import AuthenticatedOTLPExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.metrics.export import MetricExportResult
+
+from gatewaysdk.utils.otlp import LightningStoreOTLPExporter
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "instrument_agentops",
+    "uninstrument_agentops",
+]
+
+# Module-level storage for originals
+_original_handle_chat_attributes: Callable[..., Any] | None = None
+_original_handle_response: Callable[..., Any] | None = None
+_agentops_service_enabled = False
+
+
+def enable_agentops_service(enabled: bool = True) -> None:
+    """
+    Enable or disable communication with the AgentOps service.
+
+    By default, AgentOps exporters and clients will run in local mode
+    and will NOT attempt to communicate with the remote AgentOps service.
+
+    Args:
+        enabled: If True, enable all AgentOps exporters and clients.
+            All exporters and clients will operate in normal mode and send data
+            to the [AgentOps service](https://www.agentops.ai).
+    """
+    global _agentops_service_enabled
+    _agentops_service_enabled = enabled
+    logger.info(f"AgentOps service enabled is set to {enabled}.")
+
+
+def _patch_exporters():
+    import agentops.client.api
+    import agentops.sdk.core
+
+    agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter  # type: ignore
+    agentops.sdk.core.OTLPMetricExporter = BypassableOTLPMetricExporter
+    if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
+        agentops.sdk.core.OTLPSpanExporter = BypassableOTLPSpanExporter  # type: ignore
+    agentops.client.api.V3Client = BypassableV3Client
+    agentops.client.api.V4Client = BypassableV4Client
+
+
+def _unpatch_exporters():
+    import agentops.client.api
+    import agentops.sdk.core
+
+    agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter  # type: ignore
+    agentops.sdk.core.OTLPMetricExporter = OTLPMetricExporter
+    if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
+        agentops.sdk.core.OTLPSpanExporter = OTLPSpanExporter  # type: ignore
+    agentops.client.api.V3Client = V3Client
+    agentops.client.api.V4Client = V4Client
+
+
+def _unwrap_legacy_response(response: Any) -> Any:
+    if hasattr(response, "parse") and callable(response.parse):
+        return response.parse()
+    return response
+
+
+def _patch_new_agentops():
+    import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat
+    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes  # type: ignore
+
+    global _original_handle_chat_attributes
+
+    if _original_handle_chat_attributes is not None:
+        logger.warning("AgentOps already patched. Skipping.")
+        return True
+
+    _original_handle_chat_attributes = handle_chat_attributes  # type: ignore
+
+    @no_type_check
+    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
+        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
+
+        # In some cases, response is a openai._legacy_response.LegacyAPIResponse (e.g., LiteLLM, or LangChain),
+        # This is created by client.with_raw_response.create()
+        return_value = _unwrap_legacy_response(return_value)
+
+        if (
+            return_value is not None
+            and hasattr(return_value, "prompt_token_ids")
+            and return_value.prompt_token_ids is not None
+        ):
+            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
+        if (
+            return_value is not None
+            and hasattr(return_value, "response_token_ids")
+            and return_value.response_token_ids is not None
+        ):
+            attributes["response_token_ids"] = list(return_value.response_token_ids[0])
+
+        # For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
+        if (
+            return_value is not None
+            and hasattr(return_value, "choices")
+            and return_value.choices
+            and isinstance(return_value.choices, list)
+            and len(return_value.choices) > 0
+        ):
+            first_choice = return_value.choices[0]
+            # Token IDs from "choices[0].token_ids"
+            if "response_token_ids" not in attributes:
+                if hasattr(first_choice, "token_ids") and first_choice.token_ids is not None:
+                    attributes["response_token_ids"] = list(first_choice.token_ids)
+                # newer versions of OpenAI client SDK
+                elif (
+                    hasattr(first_choice, "provider_specific_fields")
+                    and first_choice.provider_specific_fields.get("token_ids") is not None
+                ):
+                    attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])
+
+            # log probability
+            # This is temporary. We need a unified convention for classifying and naming logprobs.
+            if hasattr(first_choice, "logprobs") and first_choice.logprobs is not None:
+                if hasattr(first_choice.logprobs, "content") and first_choice.logprobs.content is not None:
+                    attributes["logprobs.content"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.content]
+                    )
+                if hasattr(first_choice.logprobs, "refusal") and first_choice.logprobs.refusal is not None:
+                    attributes["logprobs.refusal"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.refusal]
+                    )
+
+        return attributes
+
+    agentops.instrumentation.providers.openai.wrappers.chat.handle_chat_attributes = _handle_chat_attributes_with_tokens
+    agentops.instrumentation.providers.openai.stream_wrapper.handle_chat_attributes = (
+        _handle_chat_attributes_with_tokens
+    )
+    logger.info("Patched newer version of agentops using handle_chat_attributes")
+    return True
+
+
+def _unpatch_new_agentops():
+    import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat
+
+    global _original_handle_chat_attributes
+    if _original_handle_chat_attributes is not None:
+        agentops.instrumentation.providers.openai.wrappers.chat.handle_chat_attributes = (
+            _original_handle_chat_attributes
+        )
+        agentops.instrumentation.providers.openai.stream_wrapper.handle_chat_attributes = (
+            _original_handle_chat_attributes
+        )
+        _original_handle_chat_attributes = None
+        logger.info("Unpatched newer version of agentops using handle_chat_attributes")
+
+
+def _patch_old_agentops():
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
+    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw  # type: ignore
+
+    global _original_handle_response
+    _original_handle_response = _handle_response  # type: ignore
+
+    @dont_throw  # type: ignore
+    def _handle_response_with_tokens(response, span, *args, **kwargs):  # type: ignore
+        _original_handle_response(response, span, *args, **kwargs)  # type: ignore
+        if hasattr(response, "prompt_token_ids"):  # type: ignore
+            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))  # type: ignore
+        if hasattr(response, "response_token_ids"):  # type: ignore
+            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))  # type: ignore
+
+        # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
+        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):  # type: ignore
+            json_data = response.http_response.json()  # type: ignore
+            if isinstance(json_data, dict):
+                if "prompt_token_ids" in json_data:
+                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))  # type: ignore
+                if "response_token_ids" in json_data:
+                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))  # type: ignore
+
+    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens  # type: ignore
+    logger.info("Patched earlier version of agentops using _handle_response")
+    return True
+
+
+def _unpatch_old_agentops():
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
+
+    global _original_handle_response
+    if _original_handle_response is not None:
+        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response  # type: ignore
+        _original_handle_response = None
+        logger.info("Unpatched earlier version of agentops using _handle_response")
+
+
+def instrument_agentops():
+    """
+    Instrument agentops to capture token IDs.
+    Automatically detects and uses the appropriate patching method based on the installed agentops version.
+    """
+    _patch_exporters()
+
+    # Try newest version first (tested for 0.4.16)
+    try:
+        return _patch_new_agentops()
+    except ImportError as e:
+        logger.debug(f"Couldn't patch newer version of agentops: {str(e)}")
+
+    # Note: 0.4.15 needs another patching method, but it's too shortlived to be worth handling separately.
+
+    # Try older version (tested for 0.4.13)
+    try:
+        return _patch_old_agentops()
+    except ImportError as e:
+        logger.warning(f"Couldn't patch older version of agentops: {str(e)}")
+        logger.error("Failed to instrument agentops - neither patching method was successful")
+        return False
+
+
+def uninstrument_agentops():
+    """Uninstrument agentops to stop capturing token IDs."""
+    _unpatch_exporters()
+
+    try:
+        _unpatch_new_agentops()
+    except Exception:
+        pass
+    try:
+        _unpatch_old_agentops()
+    except Exception:
+        pass
+
+
+class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, AuthenticatedOTLPExporter):
+    """
+    AuthenticatedOTLPExporter with switchable service control.
+
+    When `_agentops_service_enabled` is False, skip export and return success.
+    """
+
+    def should_bypass(self) -> bool:
+        return not _agentops_service_enabled
+
+
+class BypassableOTLPMetricExporter(OTLPMetricExporter):
+    """
+    OTLPMetricExporter with switchable service control.
+    When `_agentops_service_enabled` is False, skip export and return success.
+    """
+
+    def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
+        if _agentops_service_enabled:
+            return super().export(*args, **kwargs)  # type: ignore[reportUnknownMemberType]
+        else:
+            logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
+            return MetricExportResult.SUCCESS
+
+
+class BypassableOTLPSpanExporter(LightningStoreOTLPExporter):
+    """
+    OTLPSpanExporter with switchable service control.
+    When `_agentops_service_enabled` is False, skip export and return success.
+
+    This is used instead of BypassableAuthenticatedOTLPExporter on legacy AgentOps versions.
+    """
+
+    def should_bypass(self) -> bool:
+        return not _agentops_service_enabled
+
+
+class BypassableV3Client(V3Client):
+    """
+    V3Client with toggleable authentication calls.
+    Returns dummy auth response when `_agentops_service_enabled` is False.
+    """
+
+    # Temporary synchronous override of fetch_auth_token for mock purposes.
+    def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse:  # type: ignore[override]
+        if _agentops_service_enabled:
+            return super().fetch_auth_token(*args, **kwargs)  # type: ignore[override]
+        else:
+            logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
+            return AuthTokenResponse(token="dummy", project_id="dummy")
+
+
+class BypassableV4Client(V4Client):
+    """
+    V4Client with toggleable post requests.
+    Returns dummy response when `_agentops_service_enabled` is False.
+    """
+
+    def post(self, *args: Any, **kwargs: Any) -> requests.Response:
+        if _agentops_service_enabled:
+            return super().post(*args, **kwargs)
+        else:
+            logger.debug("SwitchableV4Client is switched off, skipping post request.")
+            response = requests.Response()
+            response.status_code = 200
+            response._content = b"{}"
+            return response

gatewaysdk/instrumentation/agentops_langchain.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, Dict
+
+from agentops import instrumentation
+from agentops.integration.callbacks.langchain import LangchainCallbackHandler
+
+original_on_chain_start = LangchainCallbackHandler.on_chain_start
+langgraph_entry = None
+
+__all__ = [
+    "instrument_agentops_langchain",
+    "uninstrument_agentops_langchain",
+]
+
+
+def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
+    if "name" in kwargs:
+        if serialized is None:  # type: ignore
+            serialized = {}
+        serialized = serialized.copy()
+        serialized["name"] = kwargs["name"]
+    if "run_id" in kwargs:
+        if serialized is None:  # type: ignore
+            serialized = {}
+        serialized = serialized.copy()
+        if "id" not in serialized:
+            serialized["id"] = kwargs["run_id"]
+    return original_on_chain_start(self, serialized, inputs, **kwargs)
+
+
+def instrument_agentops_langchain():
+    """Bypass AgentOp's native support for Langchain."""
+    global langgraph_entry
+    langgraph_entry = instrumentation.AGENTIC_LIBRARIES.pop("langgraph", None)
+    LangchainCallbackHandler.on_chain_start = on_chain_start
+
+
+def uninstrument_agentops_langchain():
+    """Restore AgentOp's native support for Langchain."""
+    global langgraph_entry
+    if langgraph_entry is not None:
+        instrumentation.AGENTIC_LIBRARIES["langgraph"] = langgraph_entry
+        langgraph_entry = None
+    LangchainCallbackHandler.on_chain_start = original_on_chain_start