genai-otel-instrument 0.1.1.dev0__py3-none-any.whl

genai_otel/__init__.py

@@ -0,0 +1,129 @@
+"""Top-level package for GenAI OpenTelemetry Auto-Instrumentation.
+
+This package provides a comprehensive solution for automatically instrumenting
+Generative AI (GenAI) and Large Language Model (LLM) applications with OpenTelemetry.
+It supports various LLM providers, frameworks, and common data stores (MCP tools).
+"""
+
+import logging
+import os
+import warnings
+
+import httpx
+
+# Suppress known third-party library warnings that we cannot control
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="pydantic")
+warnings.filterwarnings("ignore", message=".*validate_default.*", module="pydantic")
+warnings.filterwarnings("ignore", message=".*NumPy module was reloaded.*", module="replicate")
+
+from .__version__ import __version__
+
+# Package metadata (from pyproject.toml)
+__author__ = "Kshitij Thakkar"
+__email__ = "kshitijthakkar@rocketmail.com"
+__license__ = "Apache-2.0"
+
+# Re-exporting key components for easier access
+from .auto_instrument import setup_auto_instrumentation  # Restoring direct import
+from .config import OTelConfig
+from .cost_calculator import CostCalculator
+from .gpu_metrics import GPUMetricsCollector
+
+# Import instrumentors conditionally to avoid errors if dependencies aren't installed
+from .instrumentors import (
+    AnthropicInstrumentor,
+    AnyscaleInstrumentor,
+    AWSBedrockInstrumentor,
+    AzureOpenAIInstrumentor,
+    CohereInstrumentor,
+    GoogleAIInstrumentor,
+    GroqInstrumentor,
+    HuggingFaceInstrumentor,
+    LangChainInstrumentor,
+    LlamaIndexInstrumentor,
+    MistralAIInstrumentor,
+    OllamaInstrumentor,
+    OpenAIInstrumentor,
+    ReplicateInstrumentor,
+    TogetherAIInstrumentor,
+    VertexAIInstrumentor,
+)
+from .mcp_instrumentors.manager import MCPInstrumentorManager
+
+logger = logging.getLogger(__name__)
+
+
+def instrument(**kwargs):
+    """Public function to initialize and start auto-instrumentation.
+
+    Loads configuration from environment variables or provided keyword arguments,
+    then sets up OpenTelemetry tracing and metrics.
+
+    Args:
+        **kwargs: Configuration parameters that can override environment variables.
+                  See OTelConfig for available parameters (e.g., service_name, endpoint).
+
+    Example:
+        >>> from genai_otel import instrument
+        >>> instrument(service_name="my-app", endpoint="http://localhost:4318")
+
+    Environment Variables:
+        OTEL_SERVICE_NAME: Name of the service (default: "genai-app")
+        OTEL_EXPORTER_OTLP_ENDPOINT: OTLP endpoint (default: "http://localhost:4318")
+        GENAI_ENABLE_GPU_METRICS: Enable GPU metrics (default: "true")
+        GENAI_ENABLE_COST_TRACKING: Enable cost tracking (default: "true")
+        GENAI_ENABLE_MCP_INSTRUMENTATION: Enable MCP instrumentation (default: "true")
+        GENAI_FAIL_ON_ERROR: Fail if instrumentation errors occur (default: "false")
+        OTEL_EXPORTER_OTLP_HEADERS: OTLP headers in format "key1=val1,key2=val2"
+        GENAI_LOG_LEVEL: Logging level (default: "INFO")
+        GENAI_LOG_FILE: Log file path (optional)
+    """
+    try:
+        # Create config object, allowing kwargs to override env vars
+        config = OTelConfig(**kwargs)
+        setup_auto_instrumentation(config)
+        logger.info("GenAI OpenTelemetry instrumentation initialized successfully")
+    except Exception as e:
+        # Log the error and potentially re-raise based on fail_on_error
+        logger.error("Failed to initialize instrumentation: %s", e, exc_info=True)
+        fail_on_error = kwargs.get(
+            "fail_on_error", os.getenv("GENAI_FAIL_ON_ERROR", "false").lower() == "true"
+        )
+        if fail_on_error:
+            raise
+
+
+__all__ = [
+    # Version info
+    "__version__",
+    "__author__",
+    "__email__",
+    "__license__",
+    # Core functions
+    "instrument",
+    "setup_auto_instrumentation",  # Re-added to __all__
+    # Configuration
+    "OTelConfig",
+    # Utilities
+    "CostCalculator",
+    "GPUMetricsCollector",
+    # Instrumentors
+    "OpenAIInstrumentor",
+    "AnthropicInstrumentor",
+    "GoogleAIInstrumentor",
+    "AWSBedrockInstrumentor",
+    "AzureOpenAIInstrumentor",
+    "CohereInstrumentor",
+    "MistralAIInstrumentor",
+    "TogetherAIInstrumentor",
+    "GroqInstrumentor",
+    "LangChainInstrumentor",
+    "LlamaIndexInstrumentor",
+    "HuggingFaceInstrumentor",
+    "OllamaInstrumentor",
+    "VertexAIInstrumentor",
+    "ReplicateInstrumentor",
+    "AnyscaleInstrumentor",
+    # MCP Manager
+    "MCPInstrumentorManager",
+]

genai_otel/__version__.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.1.dev0'
+__version_tuple__ = version_tuple = (0, 1, 1, 'dev0')
+
+__commit_id__ = commit_id = None

genai_otel/auto_instrument.py

@@ -0,0 +1,413 @@
+"""Module for setting up OpenTelemetry auto-instrumentation for GenAI applications."""
+
+# isort: skip_file
+
+import logging
+import sys
+
+from opentelemetry import metrics, trace
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import ConsoleMetricExporter, PeriodicExportingMetricReader
+from opentelemetry.sdk.metrics.view import View
+from opentelemetry.sdk.metrics._internal.aggregation import ExplicitBucketHistogramAggregation
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+
+from .config import OTelConfig
+from .gpu_metrics import GPUMetricsCollector
+from .mcp_instrumentors import MCPInstrumentorManager
+from .metrics import (
+    _GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS,
+    _GEN_AI_SERVER_TBT,
+    _GEN_AI_SERVER_TFTT,
+    _MCP_CLIENT_OPERATION_DURATION_BUCKETS,
+    _MCP_PAYLOAD_SIZE_BUCKETS,
+)
+
+# Import semantic conventions
+try:
+    from openlit.semcov import SemanticConvention as SC
+except ImportError:
+    # Fallback if openlit not available
+    class SC:
+        GEN_AI_CLIENT_OPERATION_DURATION = "gen_ai.client.operation.duration"
+        GEN_AI_SERVER_TTFT = "gen_ai.server.ttft"
+        GEN_AI_SERVER_TBT = "gen_ai.server.tbt"
+
+
+# Import instrumentors - fix the import path based on your actual structure
+try:
+    from .instrumentors import (
+        AnthropicInstrumentor,
+        AnyscaleInstrumentor,
+        AWSBedrockInstrumentor,
+        AzureOpenAIInstrumentor,
+        CohereInstrumentor,
+        GoogleAIInstrumentor,
+        GroqInstrumentor,
+        HuggingFaceInstrumentor,
+        LangChainInstrumentor,
+        LlamaIndexInstrumentor,
+        MistralAIInstrumentor,
+        OllamaInstrumentor,
+        OpenAIInstrumentor,
+        ReplicateInstrumentor,
+        TogetherAIInstrumentor,
+        VertexAIInstrumentor,
+    )
+except ImportError:
+    # Fallback for testing or if instrumentors are in different structure
+    from genai_otel.instrumentors import (
+        AnthropicInstrumentor,
+        AnyscaleInstrumentor,
+        AWSBedrockInstrumentor,
+        AzureOpenAIInstrumentor,
+        CohereInstrumentor,
+        GoogleAIInstrumentor,
+        GroqInstrumentor,
+        HuggingFaceInstrumentor,
+        LangChainInstrumentor,
+        LlamaIndexInstrumentor,
+        MistralAIInstrumentor,
+        OllamaInstrumentor,
+        OpenAIInstrumentor,
+        ReplicateInstrumentor,
+        TogetherAIInstrumentor,
+        VertexAIInstrumentor,
+    )
+
+logger = logging.getLogger(__name__)
+
+# Optional OpenInference instrumentors (requires Python >= 3.10)
+try:
+    from openinference.instrumentation.litellm import LiteLLMInstrumentor  # noqa: E402
+    from openinference.instrumentation.mcp import MCPInstrumentor  # noqa: E402
+    from openinference.instrumentation.smolagents import (  # noqa: E402
+        SmolagentsInstrumentor,
+    )
+
+    OPENINFERENCE_AVAILABLE = True
+except ImportError:
+    LiteLLMInstrumentor = None
+    MCPInstrumentor = None
+    SmolagentsInstrumentor = None
+    OPENINFERENCE_AVAILABLE = False
+
+# Defines the available instrumentors. This is now at the module level for easier mocking in tests.
+INSTRUMENTORS = {
+    "openai": OpenAIInstrumentor,
+    "anthropic": AnthropicInstrumentor,
+    "google.generativeai": GoogleAIInstrumentor,
+    "boto3": AWSBedrockInstrumentor,
+    "azure.ai.openai": AzureOpenAIInstrumentor,
+    "cohere": CohereInstrumentor,
+    "mistralai": MistralAIInstrumentor,
+    "together": TogetherAIInstrumentor,
+    "groq": GroqInstrumentor,
+    "ollama": OllamaInstrumentor,
+    "vertexai": VertexAIInstrumentor,
+    "replicate": ReplicateInstrumentor,
+    "anyscale": AnyscaleInstrumentor,
+    "langchain": LangChainInstrumentor,
+    "llama_index": LlamaIndexInstrumentor,
+    "transformers": HuggingFaceInstrumentor,
+}
+
+# Add OpenInference instrumentors if available (requires Python >= 3.10)
+if OPENINFERENCE_AVAILABLE:
+    INSTRUMENTORS.update(
+        {
+            "smolagents": SmolagentsInstrumentor,
+            "mcp": MCPInstrumentor,
+            "litellm": LiteLLMInstrumentor,
+        }
+    )
+
+
+# Global list to store OTLP exporter sessions that should not be instrumented
+_OTLP_EXPORTER_SESSIONS = []
+
+
+def setup_auto_instrumentation(config: OTelConfig):
+    """
+    Set up OpenTelemetry with auto-instrumentation for LLM frameworks and MCP tools.
+
+    Args:
+        config: OTelConfig instance with configuration parameters.
+    """
+    global _OTLP_EXPORTER_SESSIONS
+    logger.info("Starting auto-instrumentation setup...")
+
+    # Configure OpenTelemetry SDK (TracerProvider, MeterProvider, etc.)
+    import os
+
+    service_instance_id = os.getenv("OTEL_SERVICE_INSTANCE_ID")
+    environment = os.getenv("OTEL_ENVIRONMENT")
+    resource_attributes = {"service.name": config.service_name}
+    if service_instance_id:
+        resource_attributes["service.instance.id"] = service_instance_id
+    if environment:
+        resource_attributes["environment"] = environment
+    resource = Resource.create(resource_attributes)
+
+    # Configure Tracing
+    tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(tracer_provider)
+    from opentelemetry.propagate import set_global_textmap
+    from opentelemetry.trace.propagation.tracecontext import (
+        TraceContextTextMapPropagator,
+    )
+
+    set_global_textmap(TraceContextTextMapPropagator())
+
+    logger.debug(f"OTelConfig endpoint: {config.endpoint}")
+    if config.endpoint:
+        # Convert timeout to float safely
+        timeout_str = os.getenv("OTEL_EXPORTER_OTLP_TIMEOUT", "10.0")
+        try:
+            timeout = float(timeout_str)
+        except (ValueError, TypeError):
+            logger.warning(f"Invalid timeout value '{timeout_str}', using default 10.0")
+            timeout = 10.0
+
+        # CRITICAL FIX: Set endpoint in environment variable so exporters can append correct paths
+        # The exporters only call _append_trace_path() when reading from env vars
+        from urllib.parse import urlparse
+
+        # Set the base endpoint in environment variable
+        os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = config.endpoint
+
+        parsed = urlparse(config.endpoint)
+        base_url = f"{parsed.scheme}://{parsed.netloc}"
+
+        # Build list of URLs to exclude from instrumentation
+        excluded_urls = [
+            base_url,
+            config.endpoint,
+            f"{base_url}/v1/traces",
+            f"{base_url}/v1/metrics",
+            config.endpoint.rstrip("/") + "/v1/traces",
+            config.endpoint.rstrip("/") + "/v1/metrics",
+        ]
+
+        # Add to environment variable (comma-separated)
+        existing = os.environ.get("OTEL_PYTHON_REQUESTS_EXCLUDED_URLS", "")
+        if existing:
+            excluded_urls.append(existing)
+        os.environ["OTEL_PYTHON_REQUESTS_EXCLUDED_URLS"] = ",".join(excluded_urls)
+        logger.info(f"Excluded OTLP endpoints from instrumentation: {base_url}")
+
+        # Set timeout in environment variable
+        os.environ["OTEL_EXPORTER_OTLP_TIMEOUT"] = str(timeout)
+
+        # Create exporters WITHOUT passing endpoint (let them read from env vars)
+        # This ensures they call _append_trace_path() correctly
+        span_exporter = OTLPSpanExporter(
+            headers=config.headers,
+        )
+        tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
+        logger.info(
+            f"OpenTelemetry tracing configured with OTLP endpoint: {span_exporter._endpoint}"
+        )
+
+        # Configure Metrics with Views for histogram buckets
+        metric_exporter = OTLPMetricExporter(
+            headers=config.headers,
+        )
+        metric_reader = PeriodicExportingMetricReader(exporter=metric_exporter)
+
+        # Create Views to configure histogram buckets for GenAI operation duration
+        duration_view = View(
+            instrument_name=SC.GEN_AI_CLIENT_OPERATION_DURATION,
+            aggregation=ExplicitBucketHistogramAggregation(
+                boundaries=_GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS
+            ),
+        )
+
+        # Create Views for MCP metrics histograms
+        mcp_duration_view = View(
+            instrument_name="mcp.client.operation.duration",
+            aggregation=ExplicitBucketHistogramAggregation(
+                boundaries=_MCP_CLIENT_OPERATION_DURATION_BUCKETS
+            ),
+        )
+
+        mcp_request_size_view = View(
+            instrument_name="mcp.request.size",
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_MCP_PAYLOAD_SIZE_BUCKETS),
+        )
+
+        mcp_response_size_view = View(
+            instrument_name="mcp.response.size",
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_MCP_PAYLOAD_SIZE_BUCKETS),
+        )
+
+        # Create Views for streaming metrics (Phase 3.4)
+        ttft_view = View(
+            instrument_name=SC.GEN_AI_SERVER_TTFT,
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_GEN_AI_SERVER_TFTT),
+        )
+
+        tbt_view = View(
+            instrument_name=SC.GEN_AI_SERVER_TBT,
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_GEN_AI_SERVER_TBT),
+        )
+
+        meter_provider = MeterProvider(
+            resource=resource,
+            metric_readers=[metric_reader],
+            views=[
+                duration_view,
+                mcp_duration_view,
+                mcp_request_size_view,
+                mcp_response_size_view,
+                ttft_view,
+                tbt_view,
+            ],
+        )
+        metrics.set_meter_provider(meter_provider)
+        logger.info(
+            f"OpenTelemetry metrics configured with OTLP endpoint: {metric_exporter._endpoint}"
+        )
+    else:
+        # Configure Console Exporters if no OTLP endpoint is set
+        span_exporter = ConsoleSpanExporter()
+        tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
+        logger.info("No OTLP endpoint configured, traces will be exported to console.")
+
+        metric_exporter = ConsoleMetricExporter()
+        metric_reader = PeriodicExportingMetricReader(exporter=metric_exporter)
+
+        # Create Views to configure histogram buckets (same as OTLP path)
+        duration_view = View(
+            instrument_name=SC.GEN_AI_CLIENT_OPERATION_DURATION,
+            aggregation=ExplicitBucketHistogramAggregation(
+                boundaries=_GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS
+            ),
+        )
+
+        # Create Views for MCP metrics histograms
+        mcp_duration_view = View(
+            instrument_name="mcp.client.operation.duration",
+            aggregation=ExplicitBucketHistogramAggregation(
+                boundaries=_MCP_CLIENT_OPERATION_DURATION_BUCKETS
+            ),
+        )
+
+        mcp_request_size_view = View(
+            instrument_name="mcp.request.size",
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_MCP_PAYLOAD_SIZE_BUCKETS),
+        )
+
+        mcp_response_size_view = View(
+            instrument_name="mcp.response.size",
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_MCP_PAYLOAD_SIZE_BUCKETS),
+        )
+
+        # Create Views for streaming metrics (Phase 3.4)
+        ttft_view = View(
+            instrument_name=SC.GEN_AI_SERVER_TTFT,
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_GEN_AI_SERVER_TFTT),
+        )
+
+        tbt_view = View(
+            instrument_name=SC.GEN_AI_SERVER_TBT,
+            aggregation=ExplicitBucketHistogramAggregation(boundaries=_GEN_AI_SERVER_TBT),
+        )
+
+        meter_provider = MeterProvider(
+            resource=resource,
+            metric_readers=[metric_reader],
+            views=[
+                duration_view,
+                mcp_duration_view,
+                mcp_request_size_view,
+                mcp_response_size_view,
+                ttft_view,
+                tbt_view,
+            ],
+        )
+        metrics.set_meter_provider(meter_provider)
+        logger.info("No OTLP endpoint configured, metrics will be exported to console.")
+
+    # OpenInference instrumentors that use different API (no config parameter)
+    # Only include if OpenInference is available (Python >= 3.10)
+    OPENINFERENCE_INSTRUMENTORS = (
+        {"smolagents", "mcp", "litellm"} if OPENINFERENCE_AVAILABLE else set()
+    )
+
+    # Auto-instrument LLM libraries based on the configuration
+    for name in config.enabled_instrumentors:
+        if name in INSTRUMENTORS:
+            try:
+                instrumentor_class = INSTRUMENTORS[name]
+                instrumentor = instrumentor_class()
+
+                # OpenInference instrumentors don't take config parameter
+                if name in OPENINFERENCE_INSTRUMENTORS:
+                    instrumentor.instrument()
+                else:
+                    instrumentor.instrument(config=config)
+
+                logger.info(f"{name} instrumentation enabled")
+            except Exception as e:
+                logger.error(f"Failed to instrument {name}: {e}", exc_info=True)
+                if config.fail_on_error:
+                    raise
+        else:
+            logger.warning(f"Unknown instrumentor '{name}' requested.")
+
+    # Auto-instrument MCP tools (databases, APIs, etc.)
+    # NOTE: OTLP endpoints are excluded via OTEL_PYTHON_REQUESTS_EXCLUDED_URLS set above
+    if config.enable_mcp_instrumentation:
+        try:
+            mcp_manager = MCPInstrumentorManager(config)
+            mcp_manager.instrument_all(config.fail_on_error)
+            logger.info("MCP tools instrumentation enabled and set up.")
+        except Exception as e:
+            logger.error(f"Failed to set up MCP tools instrumentation: {e}", exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    # Start GPU metrics collection if enabled
+    if config.enable_gpu_metrics:
+        try:
+            meter_provider = metrics.get_meter_provider()
+            gpu_collector = GPUMetricsCollector(
+                meter_provider.get_meter("genai.gpu"),
+                config,
+                interval=config.gpu_collection_interval,
+            )
+            gpu_collector.start()
+            logger.info(
+                f"GPU metrics collection started (interval: {config.gpu_collection_interval}s)."
+            )
+        except Exception as e:
+            logger.error(f"Failed to start GPU metrics collection: {e}", exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    logger.info("Auto-instrumentation setup complete")
+
+
+def instrument(**kwargs):
+    """
+    Convenience wrapper for setup_auto_instrumentation that accepts kwargs.
+
+    Set up OpenTelemetry with auto-instrumentation for LLM frameworks and MCP tools.
+
+    Args:
+        **kwargs: Keyword arguments to configure OTelConfig. These will override
+                  environment variables.
+
+    Example:
+        >>> instrument(service_name="my-app", endpoint="http://localhost:4318")
+    """
+    # Load configuration from environment variables or use provided kwargs
+    config = OTelConfig(**kwargs)
+
+    # Call the main setup function
+    setup_auto_instrumentation(config)

genai_otel/cli.py

@@ -0,0 +1,92 @@
+"""CLI tool for running instrumented applications"""
+
+import argparse
+import logging
+import os
+import runpy
+import sys
+
+from genai_otel import instrument
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    """Main entry point for the genai-instrument CLI tool.
+
+    Parses command-line arguments, initializes OpenTelemetry instrumentation,
+    and then executes the specified command/script with its arguments.
+
+    Supports two usage patterns:
+    1. genai-instrument python script.py [args...]
+    2. genai-instrument script.py [args...]
+
+    In both cases, the Python script is executed in the same process to ensure
+    instrumentation hooks are active.
+    """
+    parser = argparse.ArgumentParser(
+        description=("Run a Python script with GenAI OpenTelemetry instrumentation.")
+    )
+    parser.add_argument(
+        "command",
+        nargs=argparse.REMAINDER,
+        help="The command to run (python script.py or script.py)",
+    )
+
+    args = parser.parse_args()
+
+    if not args.command:
+        parser.print_help()
+        sys.exit(1)
+
+    # Load configuration from environment variables
+    # The `instrument` function will handle loading config.
+    try:
+        # Initialize instrumentation. This reads env vars like OTEL_SERVICE_NAME, etc.
+        # If GENAI_FAIL_ON_ERROR is true and setup fails, it will raise an exception.
+        instrument()
+    except Exception as e:
+        logger.error(f"Failed to initialize instrumentation: {e}", exc_info=True)
+        sys.exit(1)  # Exit if instrumentation setup fails and fail_on_error is true
+
+    # Parse the command to extract the Python script and its arguments
+    script_path = None
+    script_args = []
+
+    # Check if command starts with 'python' or 'python3' or 'python.exe'
+    if args.command[0].lower() in [
+        "python",
+        "python3",
+        "python.exe",
+        "python3.exe",
+    ] or os.path.basename(args.command[0]).lower().startswith("python"):
+        # Format: genai-instrument python script.py [args...]
+        if len(args.command) < 2:
+            logger.error("No Python script specified after 'python' command")
+            sys.exit(1)
+        script_path = args.command[1]
+        script_args = args.command[2:]
+    elif args.command[0].endswith(".py"):
+        # Format: genai-instrument script.py [args...]
+        script_path = args.command[0]
+        script_args = args.command[1:]
+    else:
+        logger.error(
+            f"Invalid command format. Expected 'python script.py' or 'script.py', got: {' '.join(args.command)}"
+        )
+        sys.exit(1)
+
+    # Set sys.argv to simulate running the script directly
+    # This ensures the target script receives the correct arguments
+    sys.argv = [script_path] + script_args
+
+    # Run the target script in the same process using runpy
+    # This ensures instrumentation hooks are active in the script
+    try:
+        runpy.run_path(script_path, run_name="__main__")
+    except FileNotFoundError:
+        logger.error(f"Script not found: {script_path}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error running script {script_path}: {e}", exc_info=True)
+        sys.exit(1)