mantisdk 0.1.0__py3-none-any.whl

mantisdk/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 mantisdk.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

mantisdk/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

mantisdk/instrumentation/litellm.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""LiteLLM instrumentations.
+
+Patches LiteLLM's OpenTelemetry integration to add Mantisdk-specific attributes
+to spans, including session_id, tags, environment, and call_type from context.
+
+The call_type is read from a ContextVar that is set by algorithm-specific
+decorators (e.g., @gepa.judge, @gepa.agent).
+
+[Related documentation](https://docs.litellm.ai/docs/observability/agentops_integration).
+"""
+
+import json
+from typing import Any, Optional
+
+from litellm.integrations.opentelemetry import OpenTelemetry
+
+from mantisdk.types.tracing import get_current_call_type
+
+__all__ = [
+    "instrument_litellm",
+    "uninstrument_litellm",
+]
+
+original_set_attributes = OpenTelemetry.set_attributes  # type: ignore
+
+
+def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
+    """Enhanced set_attributes that adds Mantisdk tracing metadata to spans.
+    
+    Extracts session_id, tags, and environment from kwargs["metadata"] and sets
+    them as OTEL span attributes for visibility in Insight/Langfuse.
+    
+    Also reads the current call_type from context (set by decorators like
+    @gepa.judge) and adds it as a span attribute.
+    """
+    original_set_attributes(self, span, kwargs, response_obj)
+    
+    # Add token IDs if available
+    if response_obj is not None and response_obj.get("prompt_token_ids"):
+        span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
+    if response_obj is not None and response_obj.get("response_token_ids"):
+        span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))
+    
+    # Read call_type from context (set by @gepa.judge, @gepa.agent, etc.)
+    call_type = get_current_call_type()
+    if call_type:
+        span.set_attribute("mantis.call_type", call_type)
+    
+    # Extract Mantisdk tracing metadata from kwargs
+    metadata = kwargs.get("metadata", {}) if kwargs else {}
+    if not metadata:
+        return
+    
+    # Set session_id as span attribute (standard Langfuse/Insight attribute)
+    session_id = metadata.get("session_id")
+    if session_id:
+        span.set_attribute("session.id", session_id)
+    
+    # Set tags as span attribute (JSON array for Langfuse compatibility)
+    tags = metadata.get("tags")
+    if tags and isinstance(tags, list):
+        # Langfuse expects tags as a JSON array string or individual attributes
+        span.set_attribute("tags", json.dumps(tags))
+        # Also set individual tag attributes for filtering
+        for i, tag in enumerate(tags):
+            span.set_attribute(f"tag.{i}", str(tag))
+    
+    # Set environment as span attribute
+    environment = metadata.get("environment")
+    if environment:
+        span.set_attribute("environment", environment)
+
+
+def instrument_litellm():
+    """Instrument litellm to capture token IDs."""
+    OpenTelemetry.set_attributes = patched_set_attributes
+
+
+def uninstrument_litellm():
+    """Uninstrument litellm to stop capturing token IDs."""
+    OpenTelemetry.set_attributes = original_set_attributes

mantisdk/instrumentation/vllm.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import warnings
+from typing import Any, List
+
+import vllm.entrypoints.openai.protocol
+from vllm.entrypoints.openai.protocol import ChatCompletionResponse
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
+
+__all__ = [
+    "instrument_vllm",
+    "uninstrument_vllm",
+]
+
+
+class ChatCompletionResponsePatched(ChatCompletionResponse):
+    prompt_token_ids: List[int] | None = None
+    response_token_ids: List[int] | None = None
+
+
+original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full_generator
+
+
+async def chat_completion_full_generator(
+    self: Any,
+    request: Any,
+    result_generator: Any,
+    request_id: str,
+    model_name: str,
+    conversation: Any,
+    tokenizer: Any,
+    request_metadata: Any,
+) -> Any:
+    prompt_token_ids: List[int] | None = None
+    response_token_ids: List[List[int]] | None = None
+
+    async def _generate_inceptor():
+        nonlocal prompt_token_ids, response_token_ids
+        async for res in result_generator:
+            yield res
+            prompt_token_ids = res.prompt_token_ids
+            response_token_ids = [output.token_ids for output in res.outputs]
+
+    response = await original_chat_completion_full_generator(
+        self,
+        request,
+        _generate_inceptor(),
+        request_id,
+        model_name,
+        conversation,
+        tokenizer,
+        request_metadata,
+    )
+    response = response.model_copy(
+        update={
+            "prompt_token_ids": prompt_token_ids,
+            "response_token_ids": response_token_ids,
+        }
+    )
+
+    return response
+
+
+def instrument_vllm():
+    """Instrument vLLM to capture token IDs generated by engine.
+
+    This instrumentation has been merged to upstream vLLM since v0.10.2.
+    """
+    if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
+        warnings.warn("vllm is already instrumented. Skip the instrumentation.")
+        return
+
+    vllm.entrypoints.openai.protocol.ChatCompletionResponse = ChatCompletionResponsePatched
+    OpenAIServingChat.chat_completion_full_generator = chat_completion_full_generator
+
+
+def uninstrument_vllm():
+    """Uninstrument vLLM to stop capturing token IDs generated by engine."""
+    OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator