genai-otel-instrument 0.1.24__py3-none-any.whl

genai_otel/logging_config.py

@@ -0,0 +1,45 @@
+"""Centralized logging configuration"""
+
+import logging
+import os
+import sys
+from logging.handlers import RotatingFileHandler
+from typing import Optional
+
+
+def setup_logging(
+    level: Optional[str] = None, log_dir: str = "logs", log_file_name: str = "genai_otel.log"
+):
+    """Configure logging for the library with configurable log level via environment variable
+    and log rotation.
+    """
+    # Determine log level from environment variable or default to INFO
+    env_log_level = os.environ.get("GENAI_OTEL_LOG_LEVEL")
+    log_level_str = level or env_log_level or "INFO"
+    log_level = getattr(logging, log_level_str.upper(), logging.INFO)
+
+    # Create logs directory if it doesn't exist
+    os.makedirs(log_dir, exist_ok=True)
+    log_file_path = os.path.join(log_dir, log_file_name)
+
+    # Setup handlers
+    handlers = [logging.StreamHandler(sys.stdout)]
+
+    # Add rotating file handler
+    file_handler = RotatingFileHandler(log_file_path, maxBytes=10 * 1024 * 1024, backupCount=10)
+    handlers.append(file_handler)
+
+    # Set library logger
+    logger = logging.getLogger("genai_otel")
+    logger.setLevel(log_level)
+
+    # Clear existing handlers to prevent duplicates in case of multiple calls
+    if logger.handlers:
+        for handler in logger.handlers:
+            handler.close()
+        logger.handlers = []
+
+    for handler in handlers:
+        logger.addHandler(handler)
+
+    return logger

genai_otel/mcp_instrumentors/__init__.py

@@ -0,0 +1,14 @@
+"""Module for OpenTelemetry instrumentors for Model Context Protocol (MCP) tools.
+
+This package contains individual instrumentor classes for various MCP tools,
+including databases, caching layers, message queues, vector databases, and
+generic API clients, enabling automatic tracing and metric collection of their operations.
+"""
+
+# pylint: disable=R0801
+import httpx
+
+from .base import BaseMCPInstrumentor
+from .manager import MCPInstrumentorManager
+
+__all__ = ["BaseMCPInstrumentor", "MCPInstrumentorManager"]

genai_otel/mcp_instrumentors/api_instrumentor.py

@@ -0,0 +1,144 @@
+"""OpenTelemetry instrumentor for generic API calls.
+
+This module provides the `APIInstrumentor` class, which automatically traces
+HTTP requests made using popular libraries like `requests` and `httpx`.
+It enriches spans with relevant attributes, including detected GenAI system
+information based on the URL.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+
+import httpx
+import requests
+import wrapt
+
+from genai_otel.instrumentors.base import BaseInstrumentor
+
+from ..config import OTelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class APIInstrumentor(BaseInstrumentor):
+    """Instrument custom API calls, adding GenAI-specific attributes.
+
+    This instrumentor targets common HTTP client libraries like `requests` and `httpx`.
+    It aims to add relevant attributes to spans, including GenAI system information
+    if detectable from the URL or headers.
+    """
+
+    def __init__(self, config: OTelConfig):
+        """Initializes the APIInstrumentor.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        super().__init__()  # Initialize BaseInstrumentor
+        self.config = config
+
+    def instrument(self, config: OTelConfig):
+        """Instrument requests and httpx libraries for API calls.
+
+        Applies wrappers to `requests.request`, `requests.Session.request`,
+        and `httpx.Client.request` to capture API call details.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        self.config = config  # Store the config
+
+        # CRITICAL: Do NOT wrap requests library when using OTLP HTTP exporters
+        # Wrapping requests.Session breaks OTLP exporters that use it internally
+        # try:
+        #     # Wrap requests.Session.request and requests.request
+        #     wrapt.wrap_function_wrapper(requests, "request", self._wrap_api_call)
+        #     wrapt.wrap_function_wrapper(requests.sessions.Session, "request", self._wrap_api_call)
+        #     logger.info("requests library instrumented for API calls.")
+        # except ImportError:
+        #     logger.debug("requests library not found, skipping instrumentation.")
+        # except Exception as e:
+        #     logger.error("Failed to instrument requests library: %s", e, exc_info=True)
+        #     if self.config.fail_on_error:
+        #         raise
+
+        logger.warning(
+            "requests library instrumentation disabled to prevent OTLP exporter conflicts"
+        )
+
+        try:
+            # Wrap httpx.Client.request
+            wrapt.wrap_function_wrapper(httpx.Client, "request", self._wrap_api_call)
+            logger.info("httpx library instrumented for API calls.")
+        except ImportError:
+            logger.debug("httpx library not found, skipping instrumentation.")
+        except Exception as e:
+            logger.error("Failed to instrument httpx library: %s", e, exc_info=True)
+            if self.config.fail_on_error:
+                raise
+
+    def _wrap_api_call(self, wrapped, instance, args, kwargs):
+        """Wrapper function for API calls using create_span_wrapper.
+
+        This method prepares the arguments for `create_span_wrapper` and applies it.
+        """
+        method = kwargs.get("method", args[0] if args else "unknown").upper()
+        url = kwargs.get("url", args[1] if len(args) > 1 else None)
+        span_name = f"api.call.{method.lower()}"
+        if url:
+            try:
+                parsed_url = urlparse(url)
+                span_name = f"api.call.{method.lower()}.{parsed_url.hostname}"
+            except Exception:
+                pass  # Keep default span name if URL parsing fails
+
+        instrumented_call = self.create_span_wrapper(
+            span_name=span_name, extract_attributes=self._extract_api_attributes
+        )
+        return instrumented_call(wrapped, instance, args, kwargs)
+
+    def _extract_api_attributes(
+        self, instance: Any, args: Any, kwargs: Any
+    ) -> Dict[str, Any]:  # pylint: disable=W0613
+        """Extract attributes from API call arguments for OpenTelemetry spans.
+
+        Args:
+            instance: The instance of the class the method is called on (e.g., requests.Session).
+            args: Positional arguments passed to the method.
+            kwargs: Keyword arguments passed to the method.
+
+        Returns:
+            Dict[str, Any]: A dictionary of attributes to be set on the span.
+        """
+        attrs = {}
+        method = kwargs.get("method", args[0] if args else "unknown").upper()
+        url = kwargs.get("url", args[1] if len(args) > 1 else None)
+
+        if url:
+            try:
+                parsed_url = urlparse(url)
+                if parsed_url.hostname:
+                    attrs["net.peer.name"] = parsed_url.hostname
+                attrs["url.full"] = url
+                attrs["http.method"] = method
+            except Exception as e:
+                logger.warning("Failed to parse URL '%s' for attributes: %s", url, e)
+
+        if url:
+            if "openai.com" in url:
+                attrs["gen_ai.system"] = "openai"
+            elif "anthropic.com" in url:
+                attrs["gen_ai.system"] = "anthropic"
+            elif "google.com" in url:
+                attrs["gen_ai.system"] = "google"
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:  # pylint: disable=W0613
+        """API calls typically don't have direct token usage like LLMs.
+
+        This method is part of the BaseInstrumentor interface but is not implemented
+        for generic API calls as token usage is not a standard concept here.
+        """
+        return None

genai_otel/mcp_instrumentors/base.py

@@ -0,0 +1,105 @@
+"""Base class for MCP instrumentors with shared metrics.
+
+This module provides the `BaseMCPInstrumentor` class which creates and manages
+MCP-specific metrics (requests, duration, payload sizes) that are shared across
+all MCP instrumentors (databases, APIs, vector DBs, etc.).
+"""
+
+import logging
+from typing import Optional
+
+from opentelemetry import metrics
+from opentelemetry.metrics import Counter, Histogram
+
+logger = logging.getLogger(__name__)
+
+# Import semantic conventions
+try:
+    from openlit.semcov import SemanticConvention as SC
+except ImportError:
+    # Fallback if openlit not available
+    class SC:
+        MCP_REQUESTS = "mcp.requests"
+        MCP_CLIENT_OPERATION_DURATION_METRIC = "mcp.client.operation.duration"
+        MCP_REQUEST_SIZE = "mcp.request.size"
+        MCP_RESPONSE_SIZE_METRIC = "mcp.response.size"
+
+
+class BaseMCPInstrumentor:
+    """Base class for MCP instrumentors with shared metrics.
+
+    This class provides MCP-specific metrics that can be used by all MCP instrumentors
+    to track requests, operation duration, and payload sizes.
+
+    Metrics:
+        - mcp.requests: Counter for number of MCP requests
+        - mcp.client.operation.duration: Histogram for operation duration in seconds
+        - mcp.request.size: Histogram for request payload size in bytes
+        - mcp.response.size: Histogram for response payload size in bytes
+    """
+
+    # Class-level shared metrics (created once, shared by all instances)
+    _shared_request_counter: Optional[Counter] = None
+    _shared_duration_histogram: Optional[Histogram] = None
+    _shared_request_size_histogram: Optional[Histogram] = None
+    _shared_response_size_histogram: Optional[Histogram] = None
+    _metrics_initialized = False
+
+    def __init__(self):
+        """Initialize BaseMCPInstrumentor and create shared metrics if needed."""
+        if not BaseMCPInstrumentor._metrics_initialized:
+            self._create_shared_metrics()
+
+        # Instance references to shared metrics
+        self.mcp_request_counter = BaseMCPInstrumentor._shared_request_counter
+        self.mcp_duration_histogram = BaseMCPInstrumentor._shared_duration_histogram
+        self.mcp_request_size_histogram = BaseMCPInstrumentor._shared_request_size_histogram
+        self.mcp_response_size_histogram = BaseMCPInstrumentor._shared_response_size_histogram
+
+    @classmethod
+    def _create_shared_metrics(cls):
+        """Create shared MCP metrics once at class level."""
+        if cls._metrics_initialized:
+            return
+
+        try:
+            meter = metrics.get_meter(__name__)
+
+            # MCP request counter
+            cls._shared_request_counter = meter.create_counter(
+                SC.MCP_REQUESTS,
+                description="Number of MCP requests",
+                unit="1",
+            )
+
+            # MCP operation duration histogram
+            cls._shared_duration_histogram = meter.create_histogram(
+                SC.MCP_CLIENT_OPERATION_DURATION_METRIC,
+                description="MCP operation duration",
+                unit="s",
+            )
+
+            # MCP request size histogram
+            cls._shared_request_size_histogram = meter.create_histogram(
+                SC.MCP_REQUEST_SIZE,
+                description="MCP request payload size",
+                unit="By",
+            )
+
+            # MCP response size histogram
+            cls._shared_response_size_histogram = meter.create_histogram(
+                SC.MCP_RESPONSE_SIZE_METRIC,
+                description="MCP response payload size",
+                unit="By",
+            )
+
+            cls._metrics_initialized = True
+            logger.debug("MCP shared metrics created successfully")
+
+        except Exception as e:
+            logger.warning(f"Failed to create MCP shared metrics: {e}")
+            # Set to None if creation fails
+            cls._shared_request_counter = None
+            cls._shared_duration_histogram = None
+            cls._shared_request_size_histogram = None
+            cls._shared_response_size_histogram = None

genai_otel/mcp_instrumentors/database_instrumentor.py

@@ -0,0 +1,336 @@
+"""OpenTelemetry instrumentor for various database clients.
+
+This module provides the `DatabaseInstrumentor` class, which automatically
+instruments popular Python database libraries such as SQLAlchemy, psycopg2,
+pymongo, and mysql, enabling tracing of database operations within GenAI applications.
+
+This instrumentor uses a hybrid approach:
+1. Built-in OTel instrumentors create spans with full trace context
+2. Custom wrapt wrappers add MCP-specific metrics (duration, payload sizes)
+"""
+
+import json
+import logging
+import time
+
+import wrapt
+from opentelemetry.instrumentation.mysql import MySQLInstrumentor
+from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
+from opentelemetry.instrumentation.pymongo import PymongoInstrumentor
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+
+from ..config import OTelConfig
+from .base import BaseMCPInstrumentor
+
+logger = logging.getLogger(__name__)
+
+# Conditional imports for database libraries
+try:
+    import psycopg2
+    from psycopg2.extensions import cursor as Psycopg2Cursor
+
+    PSYCOPG2_AVAILABLE = True
+except ImportError:
+    psycopg2 = None
+    Psycopg2Cursor = None
+    PSYCOPG2_AVAILABLE = False
+
+try:
+    import pymongo
+    from pymongo.collection import Collection as PymongoCollection
+
+    PYMONGO_AVAILABLE = True
+except ImportError:
+    pymongo = None
+    PymongoCollection = None
+    PYMONGO_AVAILABLE = False
+
+try:
+    import mysql.connector
+    from mysql.connector.cursor import MySQLCursor
+
+    MYSQL_AVAILABLE = True
+except ImportError:
+    mysql = None
+    MySQLCursor = None
+    MYSQL_AVAILABLE = False
+
+
+class DatabaseInstrumentor(BaseMCPInstrumentor):  # pylint: disable=R0903
+    """Instrument various database clients with traces and MCP metrics.
+
+    Uses a hybrid approach:
+    - Built-in OTel instrumentors for spans/traces
+    - Custom wrappers for MCP-specific metrics
+    """
+
+    def __init__(self, config: OTelConfig):
+        super().__init__()
+        self.config = config
+
+    def instrument(self):
+        """Instrument all detected database libraries with traces and MCP metrics.
+
+        Uses hybrid approach:
+        1. Built-in OTel instrumentors for spans/traces
+        2. Custom wrappers for MCP metrics (duration, payload sizes)
+        """
+        instrumented_count = 0
+
+        # Step 1: Use built-in instrumentors for traces/spans
+        # SQLAlchemy
+        try:
+            SQLAlchemyInstrumentor().instrument()
+            logger.info("SQLAlchemy instrumentation enabled")
+            instrumented_count += 1
+        except ImportError:
+            logger.debug("SQLAlchemy not installed, skipping instrumentation.")
+        except Exception as e:
+            logger.warning(f"SQLAlchemy instrumentation failed: {e}")
+
+        # PostgreSQL (psycopg2)
+        try:
+            Psycopg2Instrumentor().instrument()
+            logger.info("PostgreSQL (psycopg2) instrumentation enabled")
+            instrumented_count += 1
+        except ImportError:
+            logger.debug("Psycopg2 not installed, skipping instrumentation.")
+        except Exception as e:
+            logger.warning(f"PostgreSQL instrumentation failed: {e}")
+
+        # MongoDB
+        try:
+            PymongoInstrumentor().instrument()
+            logger.info("MongoDB instrumentation enabled")
+            instrumented_count += 1
+        except ImportError:
+            logger.debug("Pymongo not installed, skipping instrumentation.")
+        except Exception as e:
+            logger.warning(f"MongoDB instrumentation failed: {e}")
+
+        # MySQL
+        try:
+            MySQLInstrumentor().instrument()
+            logger.info("MySQL instrumentation enabled")
+            instrumented_count += 1
+        except ImportError:
+            logger.debug("MySQL-python not installed, skipping instrumentation.")
+        except Exception as e:
+            logger.warning(f"MySQL instrumentation failed: {e}")
+
+        # Step 2: Add custom MCP metrics wrappers
+        if self.mcp_request_counter is not None:
+            # Add metrics collection for databases that are available
+            if PSYCOPG2_AVAILABLE:
+                self._add_psycopg2_metrics()
+            if PYMONGO_AVAILABLE:
+                self._add_pymongo_metrics()
+            if MYSQL_AVAILABLE:
+                self._add_mysql_metrics()
+
+        return instrumented_count
+
+    def _add_psycopg2_metrics(self):
+        """Add MCP metrics collection to psycopg2 cursor execute methods."""
+        try:
+            # Wrap psycopg2 cursor execute methods
+            if hasattr(Psycopg2Cursor, "execute"):
+                wrapt.wrap_function_wrapper(
+                    "psycopg2.extensions", "cursor.execute", self._db_execute_wrapper("psycopg2")
+                )
+            if hasattr(Psycopg2Cursor, "executemany"):
+                wrapt.wrap_function_wrapper(
+                    "psycopg2.extensions",
+                    "cursor.executemany",
+                    self._db_execute_wrapper("psycopg2"),
+                )
+            logger.debug("PostgreSQL MCP metrics enabled")
+        except Exception as e:
+            logger.debug(f"Failed to add PostgreSQL MCP metrics: {e}")
+
+    def _add_pymongo_metrics(self):
+        """Add MCP metrics collection to pymongo collection methods."""
+        try:
+            # Wrap common pymongo collection methods
+            methods_to_wrap = [
+                "find",
+                "find_one",
+                "insert_one",
+                "insert_many",
+                "update_one",
+                "update_many",
+                "delete_one",
+                "delete_many",
+                "count_documents",
+                "aggregate",
+            ]
+            for method_name in methods_to_wrap:
+                if hasattr(PymongoCollection, method_name):
+                    wrapt.wrap_function_wrapper(
+                        "pymongo.collection",
+                        f"Collection.{method_name}",
+                        self._db_operation_wrapper("pymongo", method_name),
+                    )
+            logger.debug("MongoDB MCP metrics enabled")
+        except Exception as e:
+            logger.debug(f"Failed to add MongoDB MCP metrics: {e}")
+
+    def _add_mysql_metrics(self):
+        """Add MCP metrics collection to MySQL cursor execute methods."""
+        try:
+            # Wrap MySQL cursor execute methods
+            if hasattr(MySQLCursor, "execute"):
+                wrapt.wrap_function_wrapper(
+                    "mysql.connector.cursor",
+                    "MySQLCursor.execute",
+                    self._db_execute_wrapper("mysql"),
+                )
+            if hasattr(MySQLCursor, "executemany"):
+                wrapt.wrap_function_wrapper(
+                    "mysql.connector.cursor",
+                    "MySQLCursor.executemany",
+                    self._db_execute_wrapper("mysql"),
+                )
+            logger.debug("MySQL MCP metrics enabled")
+        except Exception as e:
+            logger.debug(f"Failed to add MySQL MCP metrics: {e}")
+
+    def _db_execute_wrapper(self, db_system: str):
+        """Create a wrapper for database execute methods that records MCP metrics.
+
+        Args:
+            db_system: Name of the database system (e.g., "psycopg2", "mysql")
+
+        Returns:
+            Wrapper function compatible with wrapt
+        """
+
+        def wrapper(wrapped, instance, args, kwargs):
+            start_time = time.time()
+            try:
+                result = wrapped(*args, **kwargs)
+                return result
+            finally:
+                # Record duration
+                duration = time.time() - start_time
+                if self.mcp_duration_histogram:
+                    self.mcp_duration_histogram.record(
+                        duration, {"db.system": db_system, "mcp.operation": "execute"}
+                    )
+
+                # Record request count
+                if self.mcp_request_counter:
+                    self.mcp_request_counter.add(
+                        1, {"db.system": db_system, "mcp.operation": "execute"}
+                    )
+
+                # Estimate request size (query + params)
+                try:
+                    query = args[0] if args else ""
+                    params = (
+                        args[1] if len(args) > 1 else kwargs.get("vars") or kwargs.get("params")
+                    )
+                    request_size = len(str(query))
+                    if params:
+                        try:
+                            request_size += len(json.dumps(params, default=str))
+                        except (TypeError, ValueError):
+                            request_size += len(str(params))
+
+                    if self.mcp_request_size_histogram:
+                        self.mcp_request_size_histogram.record(
+                            request_size, {"db.system": db_system}
+                        )
+
+                    # Estimate response size from rowcount
+                    if hasattr(instance, "rowcount") and instance.rowcount > 0:
+                        # Rough estimate: 100 bytes per row
+                        response_size = instance.rowcount * 100
+                        if self.mcp_response_size_histogram:
+                            self.mcp_response_size_histogram.record(
+                                response_size, {"db.system": db_system}
+                            )
+                except Exception as e:
+                    logger.debug(f"Failed to record payload size for {db_system}: {e}")
+
+        return wrapper
+
+    def _db_operation_wrapper(self, db_system: str, operation: str):
+        """Create a wrapper for database operations that records MCP metrics.
+
+        Args:
+            db_system: Name of the database system (e.g., "pymongo")
+            operation: Name of the operation (e.g., "find", "insert_one")
+
+        Returns:
+            Wrapper function compatible with wrapt
+        """
+
+        def wrapper(wrapped, instance, args, kwargs):
+            start_time = time.time()
+            try:
+                result = wrapped(*args, **kwargs)
+                return result
+            finally:
+                # Record duration
+                duration = time.time() - start_time
+                if self.mcp_duration_histogram:
+                    self.mcp_duration_histogram.record(
+                        duration, {"db.system": db_system, "mcp.operation": operation}
+                    )
+
+                # Record request count
+                if self.mcp_request_counter:
+                    self.mcp_request_counter.add(
+                        1, {"db.system": db_system, "mcp.operation": operation}
+                    )
+
+                # Estimate payload sizes
+                try:
+                    # Request size: serialize args and kwargs
+                    request_size = 0
+                    if args:
+                        for arg in args:
+                            if arg is not None:
+                                try:
+                                    request_size += len(json.dumps(arg, default=str))
+                                except (TypeError, ValueError):
+                                    request_size += len(str(arg))
+                    if kwargs:
+                        for val in kwargs.values():
+                            if val is not None:
+                                try:
+                                    request_size += len(json.dumps(val, default=str))
+                                except (TypeError, ValueError):
+                                    request_size += len(str(val))
+
+                    if self.mcp_request_size_histogram and request_size > 0:
+                        self.mcp_request_size_histogram.record(
+                            request_size, {"db.system": db_system, "mcp.operation": operation}
+                        )
+
+                    # Response size: estimate based on result type
+                    response_size = 0
+                    if result is not None:
+                        if isinstance(result, dict):
+                            try:
+                                response_size = len(json.dumps(result, default=str))
+                            except (TypeError, ValueError):
+                                response_size = len(str(result))
+                        elif isinstance(result, (list, tuple)):
+                            response_size = len(result) * 100  # Estimate 100 bytes per item
+                        elif isinstance(result, int):
+                            response_size = 8  # Integer size
+                        elif hasattr(result, "inserted_ids"):
+                            response_size = len(str(result.inserted_ids))
+                        elif hasattr(result, "matched_count"):
+                            response_size = 8
+
+                    if self.mcp_response_size_histogram and response_size > 0:
+                        self.mcp_response_size_histogram.record(
+                            response_size, {"db.system": db_system, "mcp.operation": operation}
+                        )
+                except Exception as e:
+                    logger.debug(f"Failed to record payload size for {db_system}.{operation}: {e}")
+
+        return wrapper

genai_otel/mcp_instrumentors/kafka_instrumentor.py

@@ -0,0 +1,31 @@
+"""OpenTelemetry instrumentor for Apache Kafka clients.
+
+This module provides the `KafkaInstrumentor` class, which automatically
+instruments Kafka producers and consumers, enabling tracing of message
+queue operations within GenAI applications.
+"""
+
+import logging
+
+from opentelemetry.instrumentation.kafka import KafkaInstrumentor as OTelKafkaInstrumentor
+
+from ..config import OTelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class KafkaInstrumentor:  # pylint: disable=R0903
+    """Instrument Kafka producers and consumers"""
+
+    def __init__(self, config: OTelConfig):
+        self.config = config
+
+    def instrument(self):
+        """Instrument Kafka"""
+        try:
+            OTelKafkaInstrumentor().instrument()
+            logger.info("Kafka instrumentation enabled")
+        except ImportError:
+            logger.debug("Kafka-python not installed, skipping instrumentation.")
+        except Exception as e:
+            logger.warning(f"Kafka instrumentation failed: {e}")