aiqa-client 0.4.3__py3-none-any.whl → 0.4.7__py3-none-any.whl

aiqa/tracing_llm_utils.py

@@ -0,0 +1,191 @@
+# Functions for extracting and setting LLM-specific attributes on a span.
+import logging
+from .constants import LOG_TAG
+from opentelemetry import trace
+from typing import Any
+
+logger = logging.getLogger(LOG_TAG)
+
+
+def _is_attribute_set(span: trace.Span, attribute_name: str) -> bool:
+    """
+    Check if an attribute is already set on a span.
+    Returns True if the attribute exists, False otherwise.
+    Safe against exceptions.
+    """
+    try:
+        # Try multiple ways to access span attributes (SDK spans may store them differently)
+        # Check public 'attributes' property
+        if hasattr(span, "attributes"):
+            attrs = span.attributes
+            if attrs and attribute_name in attrs:
+                return True
+        
+        # Check private '_attributes' (common in OpenTelemetry SDK)
+        if hasattr(span, "_attributes"):
+            attrs = span._attributes
+            if attrs and attribute_name in attrs:
+                return True
+        
+        # If we can't find the attribute, assume not set (conservative approach)
+        return False
+    except Exception:
+        # If anything goes wrong, assume not set (conservative approach)
+        return False
+
+def _extract_and_set_token_usage(span: trace.Span, result: Any) -> None:
+    """
+    Extract OpenAI API style token usage from result and add to span attributes
+    using OpenTelemetry semantic conventions for gen_ai.
+    
+    Looks for usage dict with prompt_tokens, completion_tokens, and total_tokens.
+    Sets gen_ai.usage.input_tokens, gen_ai.usage.output_tokens, and gen_ai.usage.total_tokens.
+    Only sets attributes that are not already set.
+    
+    This function detects token usage from OpenAI API response patterns:
+    - OpenAI Chat Completions API: The 'usage' object contains 'prompt_tokens', 'completion_tokens', and 'total_tokens'.
+      See https://platform.openai.com/docs/api-reference/chat/object (usage field)
+    - OpenAI Completions API: The 'usage' object contains 'prompt_tokens', 'completion_tokens', and 'total_tokens'.
+      See https://platform.openai.com/docs/api-reference/completions/object (usage field)
+    
+    This function is safe against exceptions and will not derail tracing or program execution.
+    """
+    try:
+        if not span.is_recording():
+            return
+        
+        usage = None
+        
+        # Check if result is a dict with 'usage' key
+        try:
+            if isinstance(result, dict):
+                usage = result.get("usage")
+                # Also check if result itself is a usage dict (OpenAI format)
+                if usage is None and all(key in result for key in ("prompt_tokens", "completion_tokens", "total_tokens")):
+                    usage = result
+                # Also check if result itself is a usage dict (Bedrock format)
+                elif usage is None and all(key in result for key in ("input_tokens", "output_tokens")):
+                    usage = result
+            
+            # Check if result has a 'usage' attribute (e.g., OpenAI response object)
+            elif hasattr(result, "usage"):
+                usage = result.usage
+        except Exception:
+            # If accessing result properties fails, just return silently
+            return
+        
+        # Extract token usage if found
+        if isinstance(usage, dict):
+            try:
+                # Support both OpenAI format (prompt_tokens/completion_tokens) and Bedrock format (input_tokens/output_tokens)
+                prompt_tokens = usage.get("prompt_tokens") or usage.get("PromptTokens")
+                completion_tokens = usage.get("completion_tokens") or usage.get("CompletionTokens")
+                input_tokens = usage.get("input_tokens") or usage.get("InputTokens")
+                output_tokens = usage.get("output_tokens") or usage.get("OutputTokens")
+                total_tokens = usage.get("total_tokens") or usage.get("TotalTokens")
+                
+                # Use Bedrock format if OpenAI format not available
+                if prompt_tokens is None:
+                    prompt_tokens = input_tokens
+                if completion_tokens is None:
+                    completion_tokens = output_tokens
+                
+                # Calculate total_tokens if not provided but we have input and output
+                if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
+                    total_tokens = prompt_tokens + completion_tokens
+                
+                # Only set attributes that are not already set
+                if prompt_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.input_tokens"):
+                    span.set_attribute("gen_ai.usage.input_tokens", prompt_tokens)
+                if completion_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.output_tokens"):
+                    span.set_attribute("gen_ai.usage.output_tokens", completion_tokens)
+                if total_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.total_tokens"):
+                    span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
+            except Exception:
+                # If setting attributes fails, log but don't raise
+                logger.debug(f"Failed to set token usage attributes on span")
+    except Exception:
+        # Catch any other exceptions to ensure this never derails tracing
+        logger.debug(f"Error in _extract_and_set_token_usage")
+
+
+def _extract_and_set_provider_and_model(span: trace.Span, result: Any) -> None:
+    """
+    Extract provider and model information from result and add to span attributes
+    using OpenTelemetry semantic conventions for gen_ai.
+    
+    Looks for 'model', 'provider', 'provider_name' fields in the result.
+    Sets gen_ai.provider.name and gen_ai.request.model.
+    Only sets attributes that are not already set.
+    
+    This function detects model information from common API response patterns:
+    - OpenAI Chat Completions API: The 'model' field is at the top level of the response.
+      See https://platform.openai.com/docs/api-reference/chat/object
+    - OpenAI Completions API: The 'model' field is at the top level of the response.
+      See https://platform.openai.com/docs/api-reference/completions/object
+    
+    This function is safe against exceptions and will not derail tracing or program execution.
+    """
+    try:
+        if not span.is_recording():
+            return
+        
+        model = None
+        provider = None
+        
+        # Check if result is a dict
+        try:
+            if isinstance(result, dict):
+                model = result.get("model") or result.get("Model")
+                provider = result.get("provider") or result.get("Provider") or result.get("provider_name") or result.get("providerName")
+            
+            # Check if result has attributes (e.g., OpenAI response object)
+            elif hasattr(result, "model"):
+                model = result.model
+            if hasattr(result, "provider"):
+                provider = result.provider
+            elif hasattr(result, "provider_name"):
+                provider = result.provider_name
+            elif hasattr(result, "providerName"):
+                provider = result.providerName
+            
+            # Check nested structures (e.g., response.data.model)
+            if model is None and hasattr(result, "data"):
+                data = result.data
+                if isinstance(data, dict):
+                    model = data.get("model") or data.get("Model")
+                elif hasattr(data, "model"):
+                    model = data.model
+            
+            # Check for model in choices (OpenAI pattern)
+            if model is None and isinstance(result, dict):
+                choices = result.get("choices")
+                if choices and isinstance(choices, list) and len(choices) > 0:
+                    first_choice = choices[0]
+                    if isinstance(first_choice, dict):
+                        model = first_choice.get("model")
+                    elif hasattr(first_choice, "model"):
+                        model = first_choice.model
+        except Exception:
+            # If accessing result properties fails, just return silently
+            return
+        
+        # Set attributes if found and not already set
+        try:
+            if model is not None and not _is_attribute_set(span, "gen_ai.request.model"):
+                # Convert to string if needed
+                model_str = str(model) if model is not None else None
+                if model_str:
+                    span.set_attribute("gen_ai.request.model", model_str)
+            
+            if provider is not None and not _is_attribute_set(span, "gen_ai.provider.name"):
+                # Convert to string if needed
+                provider_str = str(provider) if provider is not None else None
+                if provider_str:
+                    span.set_attribute("gen_ai.provider.name", provider_str)
+        except Exception:
+            # If setting attributes fails, log but don't raise
+            logger.debug(f"Failed to set provider/model attributes on span")
+    except Exception:
+        # Catch any other exceptions to ensure this never derails tracing
+        logger.debug(f"Error in _extract_and_set_provider_and_model")

{aiqa_client-0.4.3.dist-info → aiqa_client-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqa-client
-Version: 0.4.3
+Version: 0.4.7
 Summary: OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server
 Author-email: AIQA <info@aiqa.dev>
 License: MIT

aiqa_client-0.4.7.dist-info/RECORD

@@ -0,0 +1,15 @@
+aiqa/__init__.py,sha256=8MQBrnisjeYNrwrbTheUafEWS09GtIF7ff0fBZ1Jb24,1710
+aiqa/aiqa_exporter.py,sha256=PAEwnrqTiII_OY1q6bskPob7rKGoYOYaE7ismU1pIv4,38630
+aiqa/client.py,sha256=lcENe5LlyfH8v312ElcX_HtVuOoyIMzzJnmeKrbjXYw,10063
+aiqa/constants.py,sha256=rUI3WuY1fKB_Isaok4C9vYer2XZYEgAVxAIe13pJi14,226
+aiqa/experiment_runner.py,sha256=XAZsjVP70UH_QTk5ANSOQYAhmozuGXwKB5qWWHs-zeE,11186
+aiqa/http_utils.py,sha256=m4fu3NI9CSAfdz4yz3S-nYLoAOmUhDGas4ZcpuMLog8,2241
+aiqa/object_serialiser.py,sha256=DBv7EyXIwfwjwXHDsIwdZNFmQffRb5fKAE0r8qhoqgc,16958
+aiqa/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiqa/tracing.py,sha256=1DLiQ-HHRgVV-mLIdkzeBeWD0bLsNCV4kh1yRlurvps,44080
+aiqa/tracing_llm_utils.py,sha256=rNx6v6Wh_Mhv-_DPU9_aWS7YQcO46oiv0YPdBK1KVL8,9338
+aiqa_client-0.4.7.dist-info/licenses/LICENSE.txt,sha256=kIzkzLuzG0HHaWYm4F4W5FeJ1Yxut3Ec6bhLWyw798A,1062
+aiqa_client-0.4.7.dist-info/METADATA,sha256=a8uGQEoyu0rmY4N2UcS7FyTC6sIyGSb8Qn_-cxJq8Yc,7705
+aiqa_client-0.4.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqa_client-0.4.7.dist-info/top_level.txt,sha256=nwcsuVVSuWu27iLxZd4n1evVzv1W6FVTrSnCXCc-NQs,5
+aiqa_client-0.4.7.dist-info/RECORD,,

aiqa/test_experiment_runner.py

@@ -1,176 +0,0 @@
-"""
-Example usage of the ExperimentRunner class.
-"""
-
-import asyncio
-import os
-from dotenv import load_dotenv
-from aiqa import ExperimentRunner
-
-# Load environment variables
-load_dotenv()
-
-
-# A dummy test engine that returns a dummy response
-async def my_engine(input_data):
-    """
-    Example engine function that simulates an API call.
-    Note: For run(), the engine only takes input_data.
-    For run_example(), you can use an engine that takes (input_data, parameters).
-    """
-    # Imitate an OpenAI API response
-    # Sleep for random about 0.5 - 1 seconds
-    import random
-
-    sleep_time = random.random() * 0.5 + 0.5
-    await asyncio.sleep(sleep_time)
-    return {
-        "choices": [
-            {
-                "message": {
-                    "content": f"hello {input_data}",
-                },
-            },
-        ],
-    }
-
-
-async def scorer(output, example):
-    """
-    Example scorer function that scores the output.
-    In a real scenario, you would use the metrics from the dataset.
-    Note: For run(), the scorer only takes (output, example).
-    For run_example(), you can use a scorer that takes (output, example, parameters).
-    """
-    # This is a simple example - in practice, you'd use the metrics from the dataset
-    # and call the scoring functions accordingly
-    scores = {}
-    # Add your scoring logic here
-    return scores
-
-
-async def example_basic_usage():
-    """
-    Basic example of using ExperimentRunner.
-    """
-    if not os.getenv("AIQA_API_KEY"):
-        print("Warning: AIQA_API_KEY environment variable is not set. Example may fail.")
-
-    dataset_id = "your-dataset-id-here"
-    organisation_id = "your-organisation-id-here"
-
-    experiment_runner = ExperimentRunner(
-        dataset_id=dataset_id,
-        organisation_id=organisation_id,
-    )
-
-    # Get metrics from the dataset
-    dataset = experiment_runner.get_dataset()
-    metrics = dataset.get("metrics", [])
-    print(f"Found {len(metrics)} metrics in dataset: {[m['name'] for m in metrics]}")
-
-    # Create scorer that scores all metrics from the dataset
-    # (In practice, you'd implement this based on your metrics)
-    async def dataset_scorer(output, example):
-        # Use the metrics from the dataset to score
-        # This is a placeholder - implement based on your actual metrics
-        return await scorer(output, example)
-
-    # Get example inputs
-    example_inputs = experiment_runner.get_example_inputs()
-    print(f"Processing {len(example_inputs)} examples")
-
-    # Run experiments on each example
-    for example in example_inputs:
-        result = await experiment_runner.run_example(example, my_engine, dataset_scorer)
-        if result and len(result) > 0:
-            print(f"Scored example {example['id']}: {result}")
-        else:
-            print(f"No results for example {example['id']}")
-
-    # Get summary results
-    summary_results = experiment_runner.get_summary_results()
-    print(f"Summary results: {summary_results}")
-
-
-async def example_with_experiment_setup():
-    """
-    Example of creating an experiment with custom setup.
-    """
-    dataset_id = "your-dataset-id-here"
-    organisation_id = "your-organisation-id-here"
-
-    experiment_runner = ExperimentRunner(
-        dataset_id=dataset_id,
-        organisation_id=organisation_id,
-    )
-
-    # Create experiment with custom parameters
-    experiment = experiment_runner.create_experiment(
-        {
-            "name": "My Custom Experiment",
-            "parameters": {
-                "model": "gpt-4",
-                "temperature": 0.7,
-            },
-            "comparison_parameters": [
-                {"temperature": 0.5},
-                {"temperature": 0.9},
-            ],
-        }
-    )
-
-    print(f"Created experiment: {experiment['id']}")
-
-    # Now run the experiment
-    await experiment_runner.run(my_engine, scorer)
-
-
-async def example_stepwise():
-    """
-    Example of running experiments step by step (more control).
-    """
-    dataset_id = "your-dataset-id-here"
-    organisation_id = "your-organisation-id-here"
-
-    experiment_runner = ExperimentRunner(
-        dataset_id=dataset_id,
-        organisation_id=organisation_id,
-    )
-
-    # Get the dataset
-    dataset = experiment_runner.get_dataset()
-    metrics = dataset.get("metrics", [])
-    print(f"Found {len(metrics)} metrics in dataset")
-
-    # Create scorer for run_example (takes parameters)
-    async def my_scorer(output, example, parameters):
-        # Implement your scoring logic here
-        # Note: run_example() passes parameters, so this scorer can use them
-        return {"score": 0.8}  # Placeholder
-
-    # Get examples
-    examples = experiment_runner.get_example_inputs(limit=100)
-    print(f"Processing {len(examples)} examples")
-
-    # Process each example individually
-    for example in examples:
-        try:
-            result = await experiment_runner.run_example(example, my_engine, my_scorer)
-            print(f"Example {example['id']} completed: {result}")
-        except Exception as e:
-            print(f"Example {example['id']} failed: {e}")
-
-    # Get final summary
-    summary = experiment_runner.get_summary_results()
-    print(f"Final summary: {summary}")
-
-
-if __name__ == "__main__":
-    # Uncomment the example you want to run:
-    # asyncio.run(example_basic_usage())
-    # asyncio.run(example_with_experiment_setup())
-    # asyncio.run(example_stepwise())
-    print("Please uncomment one of the examples above to run it.")
-    print("Make sure to set your dataset_id and organisation_id in the example functions.")
-

aiqa/test_startup_reliability.py

@@ -1,249 +0,0 @@
-"""
-Test startup reliability - simulates ECS deployment scenarios where rapid initialization
-and network issues could cause deployment failures.
-
-These tests verify that:
-1. Exporter initialization doesn't block or create threads immediately
-2. Thread creation is lazy (only on first export)
-3. Network failures during startup don't cause hangs
-4. Multiple rapid initializations don't cause issues
-"""
-
-import os
-import time
-import threading
-import pytest
-from unittest.mock import patch, MagicMock
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
-
-from aiqa.client import get_aiqa_client, AIQAClient
-from aiqa.aiqa_exporter import AIQASpanExporter
-
-
-class TestStartupReliability:
-    """Tests for startup reliability in ECS-like scenarios."""
-
-    def test_exporter_initialization_does_not_create_thread_immediately(self):
-        """Verify that creating an exporter doesn't immediately start a thread."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            exporter = AIQASpanExporter(startup_delay_seconds=0.1)
-            
-            # Thread should not be created immediately
-            assert exporter.flush_timer is None
-            assert not exporter._auto_flush_started
-            
-            # Cleanup
-            exporter.shutdown()
-
-    def test_thread_created_lazily_on_first_export(self):
-        """Verify thread is only created when first span is exported."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            exporter = AIQASpanExporter(startup_delay_seconds=0.1)
-            
-            # Thread should not exist yet
-            assert exporter.flush_timer is None
-            
-            # Create a mock span and export it
-            from opentelemetry.sdk.trace import ReadableSpan
-            from opentelemetry.trace import SpanContext, TraceFlags
-            
-            mock_span = MagicMock(spec=ReadableSpan)
-            mock_span.get_span_context.return_value = SpanContext(
-                trace_id=1, span_id=1, is_remote=False, trace_flags=TraceFlags(0x01)
-            )
-            mock_span.name = "test_span"
-            mock_span.kind = 1
-            mock_span.start_time = 1000000000
-            mock_span.end_time = 2000000000
-            mock_span.status.status_code = 1
-            mock_span.attributes = {}
-            mock_span.links = []
-            mock_span.events = []
-            mock_span.resource.attributes = {}
-            mock_span.parent = None
-            
-            # Export should trigger thread creation
-            result = exporter.export([mock_span])
-            
-            # Give thread a moment to start
-            time.sleep(0.2)
-            
-            # Now thread should exist
-            assert exporter._auto_flush_started
-            assert exporter.flush_timer is not None
-            assert exporter.flush_timer.is_alive()
-            
-            # Cleanup
-            exporter.shutdown()
-            if exporter.flush_timer:
-                exporter.flush_timer.join(timeout=2.0)
-
-    def test_rapid_multiple_initializations(self):
-        """Test that multiple rapid initializations don't cause issues (simulates health checks)."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            # Simulate rapid health check calls
-            clients = []
-            for _ in range(10):
-                client = get_aiqa_client()
-                clients.append(client)
-                time.sleep(0.01)  # Very short delay
-            
-            # All should be the same singleton
-            assert all(c is clients[0] for c in clients)
-            
-            # Should not have created multiple threads
-            if clients[0].exporter:
-                assert clients[0].exporter._auto_flush_started or clients[0].exporter.flush_timer is None
-
-    def test_initialization_with_unreachable_server(self):
-        """Test that initialization doesn't hang when server is unreachable."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://unreachable-server:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            # Should not block or raise
-            client = get_aiqa_client()
-            assert client is not None
-            assert client._initialized
-            
-            # Exporter should exist but thread shouldn't be started yet
-            if client.exporter:
-                # Thread creation is lazy, so it might not exist
-                assert client.exporter.flush_timer is None or not client.exporter._auto_flush_started
-
-    def test_startup_delay_respected(self):
-        """Verify that startup delay prevents immediate flush attempts."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            exporter = AIQASpanExporter(startup_delay_seconds=0.5)
-            
-            # Create and export a span to trigger thread creation
-            from opentelemetry.sdk.trace import ReadableSpan
-            from opentelemetry.trace import SpanContext, TraceFlags
-            
-            mock_span = MagicMock(spec=ReadableSpan)
-            mock_span.get_span_context.return_value = SpanContext(
-                trace_id=1, span_id=1, is_remote=False, trace_flags=TraceFlags(0x01)
-            )
-            mock_span.name = "test_span"
-            mock_span.kind = 1
-            mock_span.start_time = 1000000000
-            mock_span.end_time = 2000000000
-            mock_span.status.status_code = 1
-            mock_span.attributes = {}
-            mock_span.links = []
-            mock_span.events = []
-            mock_span.resource.attributes = {}
-            mock_span.parent = None
-            
-            exporter.export([mock_span])
-            
-            # Thread should be created
-            time.sleep(0.1)
-            assert exporter._auto_flush_started
-            
-            # But flush should not have happened yet (within delay period)
-            # We can't easily test this without mocking time, but we verify thread exists
-            assert exporter.flush_timer is not None
-            
-            # Cleanup
-            exporter.shutdown()
-            if exporter.flush_timer:
-                exporter.flush_timer.join(timeout=2.0)
-
-    def test_concurrent_initialization(self):
-        """Test concurrent initialization from multiple threads (simulates ECS health checks)."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            clients = []
-            errors = []
-            
-            def init_client():
-                try:
-                    client = get_aiqa_client()
-                    clients.append(client)
-                except Exception as e:
-                    errors.append(e)
-            
-            # Start multiple threads initializing simultaneously
-            threads = [threading.Thread(target=init_client) for _ in range(5)]
-            for t in threads:
-                t.start()
-            for t in threads:
-                t.join(timeout=5.0)
-            
-            # Should have no errors
-            assert len(errors) == 0
-            
-            # All should be the same singleton
-            assert len(set(id(c) for c in clients)) == 1
-
-    def test_shutdown_before_thread_starts(self):
-        """Test that shutdown works even if thread was never started."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            exporter = AIQASpanExporter(startup_delay_seconds=1.0)
-            
-            # Thread should not exist
-            assert exporter.flush_timer is None
-            
-            # Shutdown should work without errors
-            exporter.shutdown()
-            
-            # Should still be able to call shutdown again
-            exporter.shutdown()
-
-    def test_initialization_timeout(self):
-        """Test that initialization completes quickly even with network issues."""
-        with patch.dict(
-            os.environ,
-            {
-                "AIQA_SERVER_URL": "http://localhost:3000",
-                "AIQA_API_KEY": "test-api-key",
-            },
-        ):
-            start_time = time.time()
-            client = get_aiqa_client()
-            elapsed = time.time() - start_time
-            
-            # Initialization should be fast (< 1 second)
-            assert elapsed < 1.0
-            assert client is not None
-