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

aiqa/__init__.py

@@ -6,7 +6,7 @@ The client initializes automatically when WithTracing is first used.
 
 Set environment variables:
     AIQA_SERVER_URL: URL of the AIQA server
-    AIQA_API_KEY: API key for authentication
+    AIQA_API_KEY: API key for authentication (required)
     AIQA_COMPONENT_TAG: Optional component identifier
     AIQA_STARTUP_DELAY_SECONDS: Optional delay before first flush (default: 10s)
 

aiqa/client.py

@@ -2,12 +2,16 @@
 import os
 import logging
 from functools import lru_cache
-from typing import Optional, TYPE_CHECKING, Any
+from typing import Optional, TYPE_CHECKING, Any, Dict
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+import requests
 
-logger = logging.getLogger("AIQA")
+from .constants import AIQA_TRACER_NAME, LOG_TAG
+
+logger = logging.getLogger(LOG_TAG)
 
 # Compatibility import for TraceIdRatioBased sampler
 # In older OpenTelemetry versions it was TraceIdRatioBasedSampler
@@ -20,7 +24,7 @@ except ImportError:
         from opentelemetry.sdk.trace.sampling import TraceIdRatioBasedSampler as TraceIdRatioBased
     except ImportError:
         logger.warning(
-            "Could not import TraceIdRatioBased or TraceIdRatioBasedSampler from "
+            f"Could not import TraceIdRatioBased or TraceIdRatioBasedSampler from "
             "opentelemetry.sdk.trace.sampling. AIQA tracing may not work correctly. "
             "Please ensure opentelemetry-sdk>=1.24.0 is installed. "
             "Try: pip install --upgrade opentelemetry-sdk"
@@ -28,7 +32,7 @@ except ImportError:
         # Set to None so we can check later
         TraceIdRatioBased = None
 
-from .constants import AIQA_TRACER_NAME
+from .http_utils import get_server_url, get_api_key, build_headers, format_http_error
 
 class AIQAClient:
     """
@@ -93,19 +97,19 @@ class AIQAClient:
         This will also set enabled=False to prevent further tracing attempts.
         """
         try:
-            logger.info("AIQA tracing shutting down")
+            logger.info(f"AIQA tracing shutting down")
             # Disable tracing to prevent attempts to use shut-down system
             self.enabled = False
             if self._provider:
                 self._provider.shutdown()
-            if self._exporter:
-                self._exporter.shutdown()
+            # OTLP exporter doesn't have a separate shutdown method - it's handled by the provider
         except Exception as e:
             logger.error(f"Error shutting down tracing: {e}")
             # Still disable even if shutdown had errors
             self.enabled = False
 
 
+
 # Global singleton instance
 client: AIQAClient = AIQAClient()
 
@@ -150,7 +154,7 @@ def get_aiqa_client() -> AIQAClient:
         # Optional: Initialize explicitly (usually not needed)
         client = get_aiqa_client()
         if client.enabled:
-            print("Tracing is enabled")
+            print(f"Tracing is enabled")
         
         @WithTracing
         def my_function():
@@ -161,7 +165,7 @@ def get_aiqa_client() -> AIQAClient:
         _init_tracing()
     except Exception as e:
         logger.error(f"Failed to initialize AIQA tracing: {e}")
-        logger.warning("AIQA tracing is disabled. Your application will continue to run without tracing.")
+        logger.warning(f"AIQA tracing is disabled. Your application will continue to run without tracing.")
     return client
 
 def _init_tracing() -> None:
@@ -171,14 +175,13 @@ def _init_tracing() -> None:
         return
     
     try:
-        server_url = os.getenv("AIQA_SERVER_URL")
-        api_key = os.getenv("AIQA_API_KEY")
+        server_url = get_server_url()
+        api_key = get_api_key()
         
-        if not server_url or not api_key:
+        if not api_key:
             client.enabled = False
-            missing_vars = [var for var, val in [("AIQA_SERVER_URL", server_url), ("AIQA_API_KEY", api_key)] if not val]
             logger.warning(
-                f"AIQA tracing is disabled: missing required environment variables: {', '.join(missing_vars)}"
+                f"AIQA tracing is disabled: missing required environment variables: AIQA_API_KEY"
             )
             client._initialized = True
             return
@@ -225,23 +228,43 @@ def _init_tracing() -> None:
 
 def _attach_aiqa_processor(provider: TracerProvider) -> None:
     """Attach AIQA span processor to the provider. Idempotent - safe to call multiple times."""
-    from .aiqa_exporter import AIQASpanExporter
-    
     try:
         # Check if already attached
         for p in provider._active_span_processor._span_processors:
-            if isinstance(getattr(p, "exporter", None), AIQASpanExporter):
-                logger.debug("AIQA span processor already attached, skipping")
+            if isinstance(getattr(p, "exporter", None), OTLPSpanExporter):
+                logger.debug(f"AIQA span processor already attached, skipping")
                 return
 
-        exporter = AIQASpanExporter(
-            server_url=os.getenv("AIQA_SERVER_URL"),
-            api_key=os.getenv("AIQA_API_KEY"),
+        server_url = get_server_url()
+        api_key = get_api_key()
+        
+        # Build headers for authentication
+        # OTLP exporter sets its own Content-Type, so we only need Authorization
+        auth_headers = {}
+        if api_key:
+            auth_headers["Authorization"] = f"ApiKey {api_key}"
+        elif os.getenv("AIQA_API_KEY"):
+            auth_headers["Authorization"] = f"ApiKey {os.getenv('AIQA_API_KEY')}"
+        
+        # OTLP HTTP exporter requires the full endpoint URL including /v1/traces
+        # Ensure server_url doesn't have trailing slash or /v1/traces, then append /v1/traces
+        base_url = server_url.rstrip('/')
+        if base_url.endswith('/v1/traces'):
+            endpoint = base_url
+        else:
+            endpoint = f"{base_url}/v1/traces"
+        
+        # Create OTLP exporter with authentication headers only
+        # The exporter will set Content-Type and other headers automatically
+        exporter = OTLPSpanExporter(
+            endpoint=endpoint,
+            headers=auth_headers if auth_headers else None,
         )
+        
         provider.add_span_processor(BatchSpanProcessor(exporter))
         global client
         client.exporter = exporter
-        logger.debug("AIQA span processor attached successfully")
+        logger.debug(f"AIQA span processor attached successfully")
     except Exception as e:
         logger.error(f"Error attaching AIQA span processor: {e}")
         # Re-raise to let _init_tracing handle it - it will log and continue
@@ -263,4 +286,66 @@ def get_aiqa_tracer() -> trace.Tracer:
     except Exception as e:
         # Log issue but still return a tracer
         logger.info(f"Issue getting AIQA tracer with version: {e}, using fallback")
-        return trace.get_tracer(AIQA_TRACER_NAME)
+        return trace.get_tracer(AIQA_TRACER_NAME)
+
+
+def get_organisation(
+    organisation_id: str,
+    server_url: Optional[str] = None,
+    api_key: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Get organisation information based on API key via an API call.
+    
+    Args:
+        organisation_id: ID of the organisation to retrieve
+        server_url: Optional server URL (defaults to AIQA_SERVER_URL env var)
+        api_key: Optional API key (defaults to AIQA_API_KEY env var)
+    
+    Returns:
+        Organisation object as a dictionary
+    """
+    url = get_server_url(server_url)
+    key = get_api_key(api_key)
+    headers = build_headers(key)
+    
+    response = requests.get(
+        f"{url}/organisation/{organisation_id}",
+        headers=headers,
+    )
+    
+    if not response.ok:
+        raise Exception(format_http_error(response, "get organisation"))
+    
+    return response.json()
+
+
+def get_api_key_info(
+    api_key_id: str,
+    server_url: Optional[str] = None,
+    api_key: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Get API key information via an API call.
+    
+    Args:
+        api_key_id: ID of the API key to retrieve
+        server_url: Optional server URL (defaults to AIQA_SERVER_URL env var)
+        api_key: Optional API key (defaults to AIQA_API_KEY env var)
+    
+    Returns:
+        ApiKey object as a dictionary
+    """
+    url = get_server_url(server_url)
+    key = get_api_key(api_key)
+    headers = build_headers(key)
+    
+    response = requests.get(
+        f"{url}/api-key/{api_key_id}",
+        headers=headers,
+    )
+    
+    if not response.ok:
+        raise Exception(format_http_error(response, "get api key info"))
+    
+    return response.json()

aiqa/constants.py

@@ -3,4 +3,6 @@ Constants used across the AIQA client package.
 """
 
 AIQA_TRACER_NAME = "aiqa-tracer"
-VERSION = "0.4.3" # automatically updated by set-version-json.sh
+VERSION = "0.5.2" # automatically updated by set-version-json.sh
+
+LOG_TAG = "AIQA" # Used in all logging output to identify AIQA messages

aiqa/experiment_runner.py

@@ -4,6 +4,8 @@ ExperimentRunner - runs experiments on datasets and scores results
 
 import os
 import time
+from .constants import LOG_TAG
+from .http_utils import build_headers, get_server_url, get_api_key, format_http_error
 from typing import Any, Dict, List, Optional, Callable, Awaitable, Union
 import requests
 
@@ -35,18 +37,15 @@ class ExperimentRunner:
         """
         self.dataset_id = dataset_id
         self.experiment_id = experiment_id
-        self.server_url = (server_url or os.getenv("AIQA_SERVER_URL", "")).rstrip("/")
-        self.api_key = api_key or os.getenv("AIQA_API_KEY", "")
+        self.server_url = get_server_url(server_url)
+        self.api_key = get_api_key(api_key)
         self.organisation = organisation_id
         self.experiment: Optional[Dict[str, Any]] = None
         self.scores: List[Dict[str, Any]] = []
 
     def _get_headers(self) -> Dict[str, str]:
         """Build HTTP headers for API requests."""
-        headers = {"Content-Type": "application/json"}
-        if self.api_key:
-            headers["Authorization"] = f"ApiKey {self.api_key}"
-        return headers
+        return build_headers(self.api_key)
 
     def get_dataset(self) -> Dict[str, Any]:
         """
@@ -61,10 +60,7 @@ class ExperimentRunner:
         )
 
         if not response.ok:
-            error_text = response.text if hasattr(response, "text") else "Unknown error"
-            raise Exception(
-                f"Failed to fetch dataset: {response.status_code} {response.reason} - {error_text}"
-            )
+            raise Exception(format_http_error(response, "fetch dataset"))
 
         return response.json()
 
@@ -92,10 +88,7 @@ class ExperimentRunner:
         )
 
         if not response.ok:
-            error_text = response.text if hasattr(response, "text") else "Unknown error"
-            raise Exception(
-                f"Failed to fetch example inputs: {response.status_code} {response.reason} - {error_text}"
-            )
+            raise Exception(format_http_error(response, "fetch example inputs"))
 
         data = response.json()
         return data.get("hits", [])
@@ -130,7 +123,7 @@ class ExperimentRunner:
             "summary_results": {},
         }
 
-        print("Creating experiment")
+        print(f"Creating experiment")
         response = requests.post(
             f"{self.server_url}/experiment",
             json=experiment_setup,
@@ -138,10 +131,7 @@ class ExperimentRunner:
         )
 
         if not response.ok:
-            error_text = response.text if hasattr(response, "text") else "Unknown error"
-            raise Exception(
-                f"Failed to create experiment: {response.status_code} {response.reason} - {error_text}"
-            )
+            raise Exception(format_http_error(response, "create experiment"))
 
         experiment = response.json()
         self.experiment_id = experiment["id"]
@@ -186,10 +176,7 @@ class ExperimentRunner:
         )
 
         if not response.ok:
-            error_text = response.text if hasattr(response, "text") else "Unknown error"
-            raise Exception(
-                f"Failed to score and store: {response.status_code} {response.reason} - {error_text}"
-            )
+            raise Exception(format_http_error(response, "score and store"))
 
         json_result = response.json()
         print(f"scoreAndStore response: {json_result}")
@@ -270,8 +257,7 @@ class ExperimentRunner:
             input_data = example["spans"][0].get("attributes", {}).get("input")
 
         if not input_data:
-            print(
-                f"Warning: Example has no input field or spans with input attribute: {example}"
+            print(f"Warning: Example has no input field or spans with input attribute: {example}"
             )
             # Run engine anyway -- this could make sense if it's all about the parameters
 
@@ -326,10 +312,7 @@ class ExperimentRunner:
         )
 
         if not response.ok:
-            error_text = response.text if hasattr(response, "text") else "Unknown error"
-            raise Exception(
-                f"Failed to fetch summary results: {response.status_code} {response.reason} - {error_text}"
-            )
+            raise Exception(format_http_error(response, "fetch summary results"))
 
         experiment2 = response.json()
         return experiment2.get("summary_results", {})

aiqa/http_utils.py

@@ -0,0 +1,143 @@
+"""
+Shared HTTP utilities for AIQA client.
+Provides common functions for building headers, handling errors, and accessing environment variables.
+Supports AIQA-specific env vars (AIQA_SERVER_URL, AIQA_API_KEY) with fallback to OTLP standard vars.
+"""
+
+import os
+from typing import Dict, Optional
+
+
+def build_headers(api_key: Optional[str] = None) -> Dict[str, str]:
+    """
+    Build HTTP headers for AIQA API requests.
+    
+    Checks AIQA_API_KEY first, then falls back to OTEL_EXPORTER_OTLP_HEADERS if not set.
+    
+    Args:
+        api_key: Optional API key. If not provided, will try to get from AIQA_API_KEY env var,
+                 then from OTEL_EXPORTER_OTLP_HEADERS.
+    
+    Returns:
+        Dictionary with Content-Type, Accept-Encoding, and optionally Authorization header.
+    """
+    headers = {
+        "Content-Type": "application/json",
+        "Accept-Encoding": "gzip, deflate, br",  # Request compression (aiohttp handles decompression automatically)
+    }
+    
+    # Check parameter first
+    if api_key:
+        headers["Authorization"] = f"ApiKey {api_key}"
+        return headers
+    
+    # Check AIQA_API_KEY env var
+    aiqa_api_key = os.getenv("AIQA_API_KEY")
+    if aiqa_api_key:
+        headers["Authorization"] = f"ApiKey {aiqa_api_key}"
+        return headers
+    
+    # Fallback to OTLP headers (format: "key1=value1,key2=value2")
+    otlp_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS")
+    if otlp_headers:
+        # Parse comma-separated key=value pairs
+        for header_pair in otlp_headers.split(","):
+            header_pair = header_pair.strip()
+            if "=" in header_pair:
+                key, value = header_pair.split("=", 1)
+                key = key.strip()
+                value = value.strip()
+                if key.lower() == "authorization":
+                    headers["Authorization"] = value
+                else:
+                    headers[key] = value
+    
+    return headers
+
+
+def get_server_url(server_url: Optional[str] = None) -> str:
+    """
+    Get server URL from parameter or environment variable, with trailing slash removed.
+    
+    Checks AIQA_SERVER_URL first, then falls back to OTEL_EXPORTER_OTLP_ENDPOINT if not set.
+    
+    Args:
+        server_url: Optional server URL. If not provided, will get from AIQA_SERVER_URL env var,
+                    then from OTEL_EXPORTER_OTLP_ENDPOINT.
+    
+    Returns:
+        Server URL with trailing slash removed. Defaults to https://server-aiqa.winterwell.com if not set.
+    """
+    # Check parameter first
+    if server_url:
+        return server_url.rstrip("/")
+    
+    # Check AIQA_SERVER_URL env var
+    url = os.getenv("AIQA_SERVER_URL")
+    if url:
+        return url.rstrip("/")
+    
+    # Fallback to OTLP endpoint
+    url = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+    if url:
+        return url.rstrip("/")
+    
+    # Default fallback
+    return "https://server-aiqa.winterwell.com"
+
+
+def get_api_key(api_key: Optional[str] = None) -> str:
+    """
+    Get API key from parameter or environment variable.
+    
+    Checks AIQA_API_KEY first, then falls back to OTEL_EXPORTER_OTLP_HEADERS if not set.
+    
+    Args:
+        api_key: Optional API key. If not provided, will get from AIQA_API_KEY env var,
+                 then from OTEL_EXPORTER_OTLP_HEADERS (looking for Authorization header).
+    
+    Returns:
+        API key or empty string if not set.
+    """
+    # Check parameter first
+    if api_key:
+        return api_key
+    
+    # Check AIQA_API_KEY env var
+    aiqa_api_key = os.getenv("AIQA_API_KEY")
+    if aiqa_api_key:
+        return aiqa_api_key
+    
+    # Fallback to OTLP headers (look for Authorization header)
+    otlp_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS")
+    if otlp_headers:
+        for header_pair in otlp_headers.split(","):
+            header_pair = header_pair.strip()
+            if "=" in header_pair:
+                key, value = header_pair.split("=", 1)
+                key = key.strip()
+                value = value.strip()
+                if key.lower() == "authorization":
+                    # Extract API key from "ApiKey <key>" or just return the value
+                    if value.startswith("ApiKey "):
+                        return value[7:]
+                    return value
+    
+    return ""
+
+
+def format_http_error(response, operation: str) -> str:
+    """
+    Format an HTTP error message from a response object.
+    
+    Args:
+        response: Response object with status_code, reason, and text attributes
+        operation: Description of the operation that failed (e.g., "fetch dataset")
+    
+    Returns:
+        Formatted error message string.
+    """
+    error_text = response.text if hasattr(response, "text") else "Unknown error"
+    status_code = getattr(response, "status_code", getattr(response, "status", "unknown"))
+    reason = getattr(response, "reason", "")
+    return f"Failed to {operation}: {status_code} {reason} - {error_text}"