mseep-agentops 0.4.18__py3-none-any.whl → 0.4.23__py3-none-any.whl

agentops/instrumentation/providers/ibm_watsonx_ai/instrumentor.py

@@ -0,0 +1,162 @@
+"""IBM watsonx.ai Instrumentation for AgentOps
+
+This module provides instrumentation for the IBM watsonx.ai API, implementing OpenTelemetry
+instrumentation for model requests and responses.
+
+Key endpoints instrumented:
+- Model.generate - Text generation API
+- Model.generate_text_stream - Streaming text generation API
+- Model.chat - Chat completion API
+- Model.chat_stream - Streaming chat completion API
+- Model.tokenize - Tokenization API
+- Model.get_details - Model details API
+"""
+
+from typing import List, Dict, Any
+from wrapt import wrap_function_wrapper
+from opentelemetry.metrics import Meter
+
+from agentops.logging import logger
+from agentops.instrumentation.common import CommonInstrumentor, StandardMetrics, InstrumentorConfig
+from agentops.instrumentation.common.wrappers import WrapConfig
+from agentops.instrumentation.providers.ibm_watsonx_ai.attributes.attributes import (
+    get_generate_attributes,
+    get_tokenize_attributes,
+    get_model_details_attributes,
+    get_chat_attributes,
+)
+from agentops.instrumentation.providers.ibm_watsonx_ai.stream_wrapper import (
+    generate_text_stream_wrapper,
+    chat_stream_wrapper,
+)
+
+# Library info for tracer/meter
+LIBRARY_NAME = "agentops.instrumentation.ibm_watsonx_ai"
+LIBRARY_VERSION = "0.1.0"
+
+# Methods to wrap for instrumentation
+WRAPPED_METHODS: List[WrapConfig] = [
+    WrapConfig(
+        trace_name="watsonx.generate",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="generate",
+        handler=get_generate_attributes,
+    ),
+    WrapConfig(
+        trace_name="watsonx.generate_text_stream",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="generate_text_stream",
+        handler=None,  # Handled by dedicated wrapper
+    ),
+    WrapConfig(
+        trace_name="watsonx.chat",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="chat",
+        handler=get_chat_attributes,
+    ),
+    WrapConfig(
+        trace_name="watsonx.chat_stream",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="chat_stream",
+        handler=None,  # Handled by dedicated wrapper
+    ),
+    WrapConfig(
+        trace_name="watsonx.tokenize",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="tokenize",
+        handler=get_tokenize_attributes,
+    ),
+    WrapConfig(
+        trace_name="watsonx.get_details",
+        package="ibm_watsonx_ai.foundation_models.inference",
+        class_name="ModelInference",
+        method_name="get_details",
+        handler=get_model_details_attributes,
+    ),
+]
+
+
+class WatsonxInstrumentor(CommonInstrumentor):
+    """An instrumentor for IBM watsonx.ai API."""
+
+    def __init__(self):
+        """Initialize the IBM watsonx.ai instrumentor."""
+        # Filter out stream methods that need custom wrapping
+        standard_methods = [
+            wc for wc in WRAPPED_METHODS if wc.method_name not in ["generate_text_stream", "chat_stream"]
+        ]
+
+        config = InstrumentorConfig(
+            library_name=LIBRARY_NAME,
+            library_version=LIBRARY_VERSION,
+            wrapped_methods=standard_methods,
+            metrics_enabled=True,
+            dependencies=["ibm-watsonx-ai >= 1.3.11"],
+        )
+        super().__init__(config)
+
+    def _create_metrics(self, meter: Meter) -> Dict[str, Any]:
+        """Create metrics for IBM watsonx.ai operations.
+
+        Args:
+            meter: The OpenTelemetry meter to use for creating metrics.
+
+        Returns:
+            Dictionary containing the created metrics.
+        """
+        return StandardMetrics.create_standard_metrics(meter)
+
+    def _custom_wrap(self, **kwargs):
+        """Perform custom wrapping for streaming methods."""
+        # Dedicated wrappers for stream methods
+        try:
+            generate_text_stream_config = next(wc for wc in WRAPPED_METHODS if wc.method_name == "generate_text_stream")
+            wrap_function_wrapper(
+                generate_text_stream_config.package,
+                f"{generate_text_stream_config.class_name}.{generate_text_stream_config.method_name}",
+                generate_text_stream_wrapper,
+            )
+            logger.debug(
+                f"Wrapped {generate_text_stream_config.package}.{generate_text_stream_config.class_name}.{generate_text_stream_config.method_name} with dedicated wrapper"
+            )
+        except (StopIteration, AttributeError, ModuleNotFoundError) as e:
+            logger.debug(f"Could not wrap generate_text_stream with dedicated wrapper: {e}")
+
+        try:
+            chat_stream_config = next(wc for wc in WRAPPED_METHODS if wc.method_name == "chat_stream")
+            wrap_function_wrapper(
+                chat_stream_config.package,
+                f"{chat_stream_config.class_name}.{chat_stream_config.method_name}",
+                chat_stream_wrapper,
+            )
+            logger.debug(
+                f"Wrapped {chat_stream_config.package}.{chat_stream_config.class_name}.{chat_stream_config.method_name} with dedicated wrapper"
+            )
+        except (StopIteration, AttributeError, ModuleNotFoundError) as e:
+            logger.debug(f"Could not wrap chat_stream with dedicated wrapper: {e}")
+
+        logger.info("IBM watsonx.ai instrumentation enabled")
+
+    def _custom_unwrap(self, **kwargs):
+        """Remove custom wrapping for streaming methods."""
+        # Unwrap streaming methods manually
+        from opentelemetry.instrumentation.utils import unwrap as otel_unwrap
+
+        for wrap_config in WRAPPED_METHODS:
+            if wrap_config.method_name in ["generate_text_stream", "chat_stream"]:
+                try:
+                    otel_unwrap(wrap_config.package, f"{wrap_config.class_name}.{wrap_config.method_name}")
+                    logger.debug(
+                        f"Unwrapped streaming method {wrap_config.package}.{wrap_config.class_name}.{wrap_config.method_name}"
+                    )
+                except Exception as e:
+                    logger.debug(
+                        f"Failed to unwrap streaming method {wrap_config.package}.{wrap_config.class_name}.{wrap_config.method_name}: {e}"
+                    )
+
+        logger.info("IBM watsonx.ai instrumentation disabled")

agentops/instrumentation/providers/ibm_watsonx_ai/stream_wrapper.py

@@ -0,0 +1,302 @@
+"""Stream wrappers for IBM watsonx.ai.
+
+This module provides stream wrapper classes and functions for IBM watsonx.ai's streaming
+responses, implementing telemetry tracking for streaming content.
+"""
+
+import json
+from opentelemetry.trace import get_tracer, SpanKind
+from agentops.logging import logger
+from agentops.instrumentation.providers.ibm_watsonx_ai import LIBRARY_NAME, LIBRARY_VERSION
+from agentops.instrumentation.providers.ibm_watsonx_ai.attributes.common import (
+    extract_params_attributes,
+    convert_params_to_dict,
+    extract_prompt_from_args,
+    extract_messages_from_args,
+    extract_params_from_args,
+)
+from agentops.semconv import SpanAttributes, LLMRequestTypeValues, CoreAttributes, MessageAttributes
+
+
+class TracedStream:
+    """A wrapper for IBM watsonx.ai's streaming response that adds telemetry."""
+
+    def __init__(self, original_stream, span):
+        """Initialize with the original stream and span."""
+        self.original_stream = original_stream
+        self.span = span
+        self.completion_content = ""
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.model_id = None
+
+    def __iter__(self):
+        """Iterate through chunks, tracking content and attempting to extract token data."""
+        try:
+            for yielded_chunk in self.original_stream:
+                # Initialize data for this chunk
+                generated_text_chunk = ""
+                model_id_chunk = None
+
+                try:
+                    # Attempt to access internal frame local variable 'chunk' for full data
+                    internal_chunk_data_str = getattr(self.original_stream, "gi_frame", {}).f_locals.get("chunk")
+
+                    if isinstance(internal_chunk_data_str, str) and internal_chunk_data_str.startswith("data: "):
+                        try:
+                            # Remove 'data: ' prefix and parse JSON
+                            json_payload_str = internal_chunk_data_str[len("data: ") :]
+                            json_payload = json.loads(json_payload_str)
+
+                            # Determine if it's generate_text_stream or chat_stream structure
+                            if "results" in json_payload:  # Likely generate_text_stream
+                                model_id_chunk = json_payload.get("model_id")
+                                if isinstance(json_payload["results"], list):
+                                    for result in json_payload["results"]:
+                                        if isinstance(result, dict):
+                                            # Use yielded_chunk for generated_text as internal one might be partial
+                                            if isinstance(yielded_chunk, str):
+                                                generated_text_chunk = yielded_chunk
+                                            # Use the first non-zero input token count found
+                                            if self.input_tokens == 0 and result.get("input_token_count", 0) > 0:
+                                                self.input_tokens = result.get("input_token_count", 0)
+                                            # Accumulate output tokens
+                                            self.output_tokens += result.get("generated_token_count", 0)
+
+                            elif "choices" in json_payload:  # Likely chat_stream
+                                # model_id might be at top level or within choices in other APIs, check top first
+                                model_id_chunk = json_payload.get("model_id") or json_payload.get("model")
+                                if isinstance(json_payload["choices"], list) and json_payload["choices"]:
+                                    choice = json_payload["choices"][0]
+                                    if isinstance(choice, dict):
+                                        delta = choice.get("delta", {})
+                                        if isinstance(delta, dict):
+                                            generated_text_chunk = delta.get("content", "")
+
+                                        # Check for finish reason to potentially get final usage
+                                        finish_reason = choice.get("finish_reason")
+                                        if finish_reason == "stop":
+                                            try:
+                                                final_response_data = getattr(
+                                                    self.original_stream, "gi_frame", {}
+                                                ).f_locals.get("parsed_response")
+                                                if (
+                                                    isinstance(final_response_data, dict)
+                                                    and "usage" in final_response_data
+                                                ):
+                                                    usage = final_response_data["usage"]
+                                                    if isinstance(usage, dict):
+                                                        # Update token counts with final values
+                                                        self.input_tokens = usage.get(
+                                                            "prompt_tokens", self.input_tokens
+                                                        )
+                                                        self.output_tokens = usage.get(
+                                                            "completion_tokens", self.output_tokens
+                                                        )
+                                                        # Update span immediately with final counts
+                                                        if self.input_tokens is not None:
+                                                            self.span.set_attribute(
+                                                                SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+                                                                self.input_tokens,
+                                                            )
+                                                        if self.output_tokens is not None:
+                                                            self.span.set_attribute(
+                                                                SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+                                                                self.output_tokens,
+                                                            )
+                                                        if (
+                                                            self.input_tokens is not None
+                                                            and self.output_tokens is not None
+                                                        ):
+                                                            self.span.set_attribute(
+                                                                SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+                                                                self.input_tokens + self.output_tokens,
+                                                            )
+
+                                            except AttributeError as final_attr_err:
+                                                logger.debug(
+                                                    f"Could not access internal generator state for final response: {final_attr_err}"
+                                                )
+                                            except Exception as final_err:
+                                                logger.debug(
+                                                    f"Error accessing or processing final response data: {final_err}"
+                                                )
+
+                        except json.JSONDecodeError as json_err:
+                            logger.debug(f"Failed to parse JSON from internal chunk data: {json_err}")
+                            # Fallback to using the yielded chunk directly
+                            if isinstance(yielded_chunk, dict):  # chat_stream yields dicts
+                                if "choices" in yielded_chunk and yielded_chunk["choices"]:
+                                    delta = yielded_chunk["choices"][0].get("delta", {})
+                                    generated_text_chunk = delta.get("content", "")
+                            elif isinstance(yielded_chunk, str):  # generate_text_stream yields strings
+                                generated_text_chunk = yielded_chunk
+                        except Exception as parse_err:
+                            logger.debug(f"Error processing internal chunk data: {parse_err}")
+                            if isinstance(yielded_chunk, dict):  # Fallback for chat
+                                if "choices" in yielded_chunk and yielded_chunk["choices"]:
+                                    delta = yielded_chunk["choices"][0].get("delta", {})
+                                    generated_text_chunk = delta.get("content", "")
+                            elif isinstance(yielded_chunk, str):  # Fallback for generate
+                                generated_text_chunk = yielded_chunk
+                    else:
+                        # If internal data not found or not in expected format, use yielded chunk
+                        if isinstance(yielded_chunk, dict):  # chat_stream yields dicts
+                            if "choices" in yielded_chunk and yielded_chunk["choices"]:
+                                delta = yielded_chunk["choices"][0].get("delta", {})
+                                generated_text_chunk = delta.get("content", "")
+                        elif isinstance(yielded_chunk, str):  # generate_text_stream yields strings
+                            generated_text_chunk = yielded_chunk
+
+                except AttributeError as attr_err:
+                    logger.debug(f"Could not access internal generator state (gi_frame.f_locals): {attr_err}")
+                    if isinstance(yielded_chunk, dict):  # Fallback for chat
+                        if "choices" in yielded_chunk and yielded_chunk["choices"]:
+                            delta = yielded_chunk["choices"][0].get("delta", {})
+                            generated_text_chunk = delta.get("content", "")
+                    elif isinstance(yielded_chunk, str):  # Fallback for generate
+                        generated_text_chunk = yielded_chunk
+                except Exception as e:
+                    logger.debug(f"Error accessing or processing internal generator state: {e}")
+                    if isinstance(yielded_chunk, dict):  # Fallback for chat
+                        if "choices" in yielded_chunk and yielded_chunk["choices"]:
+                            delta = yielded_chunk["choices"][0].get("delta", {})
+                            generated_text_chunk = delta.get("content", "")
+                    elif isinstance(yielded_chunk, str):  # Fallback for generate
+                        generated_text_chunk = yielded_chunk
+
+                # Accumulate completion content regardless of where it came from
+                self.completion_content += generated_text_chunk
+
+                # Update span attributes within the loop if data is available
+                if model_id_chunk and not self.model_id:
+                    self.model_id = model_id_chunk
+                    self.span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, self.model_id)
+
+                if self.input_tokens is not None:
+                    self.span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, self.input_tokens)
+                if self.output_tokens is not None:
+                    self.span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, self.output_tokens)
+                if self.input_tokens is not None and self.output_tokens is not None:
+                    self.span.set_attribute(
+                        SpanAttributes.LLM_USAGE_TOTAL_TOKENS, self.input_tokens + self.output_tokens
+                    )
+
+                # Yield the original chunk that the user expects
+                yield yielded_chunk
+        finally:
+            # Update final completion content attribute after stream finishes
+            if self.completion_content:
+                self.span.set_attribute(MessageAttributes.COMPLETION_TYPE.format(i=0), "text")
+                self.span.set_attribute(MessageAttributes.COMPLETION_ROLE.format(i=0), "assistant")
+                self.span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), self.completion_content)
+
+            # Final update for token counts
+            if self.input_tokens is not None:
+                self.span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, self.input_tokens)
+            if self.output_tokens is not None:
+                self.span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, self.output_tokens)
+            if self.input_tokens is not None and self.output_tokens is not None:
+                self.span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, self.input_tokens + self.output_tokens)
+
+            # End the span when the stream is exhausted
+            if self.span.is_recording():
+                self.span.end()
+
+
+def generate_text_stream_wrapper(wrapped, instance, args, kwargs):
+    """Wrapper for the Model.generate_text_stream method."""
+    tracer = get_tracer(LIBRARY_NAME, LIBRARY_VERSION)
+    span = tracer.start_span(
+        "watsonx.generate_text_stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value},
+    )
+
+    # Extract prompt using helper function
+    prompt = extract_prompt_from_args(args, kwargs)
+    if prompt:
+        span.set_attribute(MessageAttributes.PROMPT_ROLE.format(i=0), "user")
+        span.set_attribute(MessageAttributes.PROMPT_CONTENT.format(i=0), prompt)
+        span.set_attribute(MessageAttributes.PROMPT_TYPE.format(i=0), "text")
+
+    # Extract parameters using helper function
+    params = extract_params_from_args(args, kwargs)
+    if params:
+        params_dict = convert_params_to_dict(params)
+        if params_dict:
+            try:
+                span_attributes = extract_params_attributes(params_dict)
+                for key, value in span_attributes.items():
+                    span.set_attribute(key, value)
+            except Exception as e:
+                logger.debug(f"Error extracting attributes from params dict: {e}")
+
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    try:
+        stream = wrapped(*args, **kwargs)
+        return TracedStream(stream, span)
+    except Exception as e:
+        span.record_exception(e)
+        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+        span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+        span.end()
+        raise
+
+
+def chat_stream_wrapper(wrapped, instance, args, kwargs):
+    """Wrapper for the Model.chat_stream method."""
+    tracer = get_tracer(LIBRARY_NAME, LIBRARY_VERSION)
+    span = tracer.start_span(
+        "watsonx.chat_stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+    )
+
+    # Extract messages using helper function
+    messages = extract_messages_from_args(args, kwargs)
+    if messages and isinstance(messages, list):
+        for i, message in enumerate(messages):
+            if isinstance(message, dict):
+                role = message.get("role")
+                content = message.get("content")
+                # Handle complex content (list of dicts) vs simple string
+                if isinstance(content, list):
+                    text_content = []
+                    for item in content:
+                        if isinstance(item, dict) and item.get("type") == "text":
+                            text_content.append(item.get("text", ""))
+                    content_str = " ".join(text_content)
+                else:
+                    content_str = str(content)
+
+                if role:
+                    span.set_attribute(MessageAttributes.PROMPT_ROLE.format(i=i), role)
+                if content_str:
+                    span.set_attribute(MessageAttributes.PROMPT_CONTENT.format(i=i), content_str)
+                    span.set_attribute(MessageAttributes.PROMPT_TYPE.format(i=i), "text")
+
+    # Extract parameters using helper function
+    params = extract_params_from_args(args, kwargs)
+    if params:
+        params_dict = convert_params_to_dict(params)
+        if params_dict:
+            try:
+                span_attributes = extract_params_attributes(params_dict)
+                for key, value in span_attributes.items():
+                    span.set_attribute(key, value)
+            except Exception as e:
+                logger.debug(f"Error extracting attributes from params dict: {e}")
+
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    try:
+        stream = wrapped(*args, **kwargs)
+        return TracedStream(stream, span)
+    except Exception as e:
+        span.record_exception(e)
+        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+        span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+        span.end()
+        raise

agentops/instrumentation/providers/mem0/__init__.py

@@ -0,0 +1,45 @@
+"""Mem0 instrumentation library for AgentOps.
+
+This package provides instrumentation for the Mem0 memory management system,
+capturing telemetry data for memory operations.
+"""
+
+import logging
+from agentops.instrumentation.common import LibraryInfo
+
+# Import memory operation wrappers
+from .memory import (
+    mem0_add_wrapper,
+    mem0_search_wrapper,
+    mem0_get_all_wrapper,
+    mem0_get_wrapper,
+    mem0_delete_wrapper,
+    mem0_update_wrapper,
+    mem0_delete_all_wrapper,
+    mem0_history_wrapper,
+)
+
+logger = logging.getLogger(__name__)
+
+# Library information
+_library_info = LibraryInfo(name="agentops.instrumentation.mem0", package_name="mem0ai")
+LIBRARY_NAME = _library_info.name
+LIBRARY_VERSION = "1.0.0"  # Internal version for instrumentation
+
+# Import after defining constants to avoid circular imports
+from agentops.instrumentation.providers.mem0.instrumentor import Mem0Instrumentor  # noqa: E402
+
+__all__ = [
+    "LIBRARY_NAME",
+    "LIBRARY_VERSION",
+    "Mem0Instrumentor",
+    # Memory operation wrappers
+    "mem0_add_wrapper",
+    "mem0_search_wrapper",
+    "mem0_get_all_wrapper",
+    "mem0_get_wrapper",
+    "mem0_delete_wrapper",
+    "mem0_update_wrapper",
+    "mem0_delete_all_wrapper",
+    "mem0_history_wrapper",
+]