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

aiqa/client.py

@@ -7,7 +7,9 @@ from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
-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 +22,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 +30,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
 
 class AIQAClient:
     """
@@ -93,7 +95,7 @@ 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:
@@ -150,7 +152,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 +163,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,8 +173,8 @@ 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:
             client.enabled = False
@@ -231,17 +233,18 @@ def _attach_aiqa_processor(provider: TracerProvider) -> None:
         # 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")
+                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"),
+            # max_buffer_spans will be read from AIQA_MAX_BUFFER_SPANS env var by the exporter
         )
         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

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.4.7" # 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,69 @@
+"""
+Shared HTTP utilities for AIQA client.
+Provides common functions for building headers, handling errors, and accessing environment variables.
+"""
+
+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.
+    
+    Args:
+        api_key: Optional API key. If not provided, will try to get from AIQA_API_KEY env var.
+    
+    Returns:
+        Dictionary with Content-Type and optionally Authorization header.
+    """
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"ApiKey {api_key}"
+    elif os.getenv("AIQA_API_KEY"):
+        headers["Authorization"] = f"ApiKey {os.getenv('AIQA_API_KEY')}"
+    return headers
+
+
+def get_server_url(server_url: Optional[str] = None) -> str:
+    """
+    Get server URL from parameter or environment variable, with trailing slash removed.
+    
+    Args:
+        server_url: Optional server URL. If not provided, will get from AIQA_SERVER_URL env var.
+    
+    Returns:
+        Server URL with trailing slash removed, or empty string if not set.
+    """
+    url = server_url or os.getenv("AIQA_SERVER_URL", "")
+    return url.rstrip("/")
+
+
+def get_api_key(api_key: Optional[str] = None) -> str:
+    """
+    Get API key from parameter or environment variable.
+    
+    Args:
+        api_key: Optional API key. If not provided, will get from AIQA_API_KEY env var.
+    
+    Returns:
+        API key or empty string if not set.
+    """
+    return api_key or os.getenv("AIQA_API_KEY", "")
+
+
+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}"

aiqa/object_serialiser.py

@@ -7,10 +7,36 @@ import json
 import os
 import dataclasses
 import logging
+from .constants import LOG_TAG
 from datetime import datetime, date, time
 from typing import Any, Callable, Set
+from json.encoder import JSONEncoder
 
-logger = logging.getLogger("aiqa")
+logger = logging.getLogger(LOG_TAG)
+
+def sanitize_string_for_utf8(text: str) -> str:
+    """
+    Sanitize a string to remove surrogate characters that can't be encoded to UTF-8.
+    Surrogate characters (U+D800 to U+DFFF) are invalid in UTF-8 and can cause encoding errors.
+    
+    Args:
+        text: The string to sanitize
+        
+    Returns:
+        A string with surrogate characters replaced by the Unicode replacement character (U+FFFD)
+    """
+    if text == None:
+        return None
+    if not isinstance(text, str): # paranoia
+        text = str(text)
+    try:
+        # Try encoding to UTF-8 to check if there are any issues
+        text.encode('utf-8')
+        return text
+    except UnicodeEncodeError:
+        # If encoding fails, replace surrogates with replacement character
+        # This handles surrogates that can't be encoded
+        return text.encode('utf-8', errors='replace').decode('utf-8', errors='replace')
 
 def toNumber(value: str|int|None) -> int:
     """Convert string to number. handling units like g, m, k, (also mb kb gb though these should be avoided)"""
@@ -105,7 +131,7 @@ def serialize_for_span(value: Any) -> Any:
     """
     Serialize a value for span attributes.
     OpenTelemetry only accepts primitives (bool, str, bytes, int, float) or sequences of those.
-    Complex types (dicts, lists, objects) are converted to JSON strings.
+    Complex types (dicts, objects) are converted to JSON strings.
     
     Handles objects by attempting to convert them to dicts, with safeguards against:
     - Circular references
@@ -118,14 +144,17 @@ def serialize_for_span(value: Any) -> Any:
     
     # For sequences, check if all elements are primitives
     if isinstance(value, (list, tuple)):
-        # If all elements are primitives, return as list
-        if all(isinstance(item, (str, int, float, bool, bytes, type(None))) for item in value):
-            return list(value)
-        # Otherwise serialize to JSON string
-        try:
-            return safe_json_dumps(value)
-        except Exception:
-            return str(value)
+        # Use short-circuiting loop instead of all() for better performance on large lists
+        # Only iterate until we find a non-primitive
+        for item in value:
+            if not isinstance(item, (str, int, float, bool, bytes, type(None))):
+                # Found non-primitive, serialize to JSON string
+                try:
+                    return safe_json_dumps(value)
+                except Exception:
+                    return str(value)
+        # All elements are primitives, return as list
+        return list(value)
     
     # For dicts and other complex types, serialize to JSON string
     try:
@@ -140,10 +169,13 @@ def safe_str_repr(value: Any) -> str:
     Safely convert a value to string representation.
     Handles objects with __repr__ that might raise exceptions.
     Uses AIQA_MAX_OBJECT_STR_CHARS environment variable (default: 100000) to limit length.
+    Also sanitizes surrogate characters to prevent UTF-8 encoding errors.
     """
     try:
         # Try __repr__ first (usually more informative)
         repr_str = repr(value)
+        # Sanitize surrogate characters that can't be encoded to UTF-8
+        repr_str = sanitize_string_for_utf8(repr_str)
         # Limit length to avoid huge strings
         if len(repr_str) > AIQA_MAX_OBJECT_STR_CHARS:
             return repr_str[:AIQA_MAX_OBJECT_STR_CHARS] + "... (truncated)"
@@ -158,7 +190,7 @@ def safe_str_repr(value: Any) -> str:
 
 def object_to_dict(obj: Any, visited: Set[int], max_depth: int = 10, current_depth: int = 0) -> Any:
     """
-    Convert an object to a dictionary representation.
+    Convert an object to a dictionary representation. Applies data filters to the object.
     
     Args:
         obj: The object to convert
@@ -172,7 +204,7 @@ def object_to_dict(obj: Any, visited: Set[int], max_depth: int = 10, current_dep
     if current_depth > max_depth:
         return "<max depth exceeded>"
     
-    obj_id = id(obj)
+    obj_id = id(obj) # note: id cannot raise exception
     if obj_id in visited:
         return "<circular reference>"
     
@@ -185,53 +217,42 @@ def object_to_dict(obj: Any, visited: Set[int], max_depth: int = 10, current_dep
         return obj
     
     # Handle datetime objects
-    if isinstance(obj, datetime):
-        return obj.isoformat()
-    if isinstance(obj, date):
-        return obj.isoformat()
-    if isinstance(obj, time):
-        return obj.isoformat()
+    if isinstance(obj, datetime) or isinstance(obj, date) or isinstance(obj, time):
+        try:
+            return obj.isoformat()
+        except Exception: # paranoia if isoformat() fails (e.g., invalid datetime state, custom implementation bug)
+            return safe_str_repr(obj)
     
     # Handle dict
     if isinstance(obj, dict):
         visited.add(obj_id)
-        try:
-            result = {}
-            for k, v in obj.items():
-                try:
-                    key_str = str(k) if not isinstance(k, (str, int, float, bool)) else k
-                    filtered_value = _apply_data_filters(key_str, v)
-                    result[key_str] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
-                except Exception as e:
-                    # If one key-value pair fails, log and use string representation for the value
-                    key_str = str(k) if not isinstance(k, (str, int, float, bool)) else k
-                    logger.debug(f"Failed to convert dict value for key '{key_str}': {e}")
-                    result[key_str] = safe_str_repr(v)
-            visited.remove(obj_id)
-            return result
-        except Exception as e:
-            visited.discard(obj_id)
-            logger.debug(f"Failed to convert dict to dict: {e}")
-            return safe_str_repr(obj)
+        result = {}
+        for k, v in obj.items():
+            try:
+                key_str = str(k) if not isinstance(k, (str, int, float, bool)) else k
+                filtered_value = _apply_data_filters(key_str, v)
+                result[key_str] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
+            except Exception as e:
+                # If one key-value pair fails, log and use string representation for the value
+                key_str = str(k) if not isinstance(k, (str, int, float, bool)) else k
+                logger.debug(f"Failed to convert dict value for key '{key_str}': {e}")
+                result[key_str] = safe_str_repr(v)
+        visited.remove(obj_id)
+        return result
     
     # Handle list/tuple
     if isinstance(obj, (list, tuple)):
         visited.add(obj_id)
-        try:
-            result = []
-            for item in obj:
-                try:
-                    result.append(object_to_dict(item, visited, max_depth, current_depth + 1))
-                except Exception as e:
-                    # If one item fails, log and use its string representation
-                    logger.debug(f"Failed to convert list item {type(item).__name__} to dict: {e}")
-                    result.append(safe_str_repr(item))
-            visited.remove(obj_id)
-            return result
-        except Exception as e:
-            visited.discard(obj_id)
-            logger.debug(f"Failed to convert list/tuple to dict: {e}")
-            return safe_str_repr(obj)
+        result = []
+        for item in obj:
+            try:
+                result.append(object_to_dict(item, visited, max_depth, current_depth + 1))
+            except Exception as e:
+                # If one item fails, log and use its string representation
+                logger.debug(f"Failed to convert list item {type(item).__name__} to dict: {e}")
+                result.append(safe_str_repr(item))
+        visited.remove(obj_id)
+        return result
     
     # Handle dataclasses
     if dataclasses.is_dataclass(obj):
@@ -258,18 +279,11 @@ def object_to_dict(obj: Any, visited: Set[int], max_depth: int = 10, current_dep
     if hasattr(obj, "__dict__"):
         visited.add(obj_id)
         try:
-            result = {}
-            for key, value in obj.__dict__.items():
-                # Skip private attributes that start with __
-                if not (isinstance(key, str) and key.startswith("__")):
-                    filtered_value = _apply_data_filters(key, value)
-                    result[key] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
-            visited.remove(obj_id)
-            return result
-        except Exception as e:
+            obj_dict = obj.__dict__
+            return object_to_dict(obj_dict, visited, max_depth, current_depth) # Note: Don't count using __dict__ as a recursion depth +1 step
+        except Exception as e: # paranoia: object_to_dict should never raise an exception
             visited.discard(obj_id)
-            # Log the error for debugging, but still return string representation
-            logger.debug(f"Failed to convert object {type(obj).__name__} to dict: {e}")
+            logger.debug(f"Failed to convert object {type(obj).__name__} with __dict__ to dict: {e}")
             return safe_str_repr(obj)
     
     # Handle objects with __slots__
@@ -311,6 +325,36 @@ def object_to_dict(obj: Any, visited: Set[int], max_depth: int = 10, current_dep
     return safe_str_repr(obj)
 
 
+class SizeLimitedJSONEncoder(JSONEncoder):
+    """
+    Custom JSON encoder that stops serialization early when max_size_chars is reached.
+    Tracks output length incrementally and stops yielding chunks when limit is exceeded.
+    """
+    def __init__(self, max_size_chars: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_size_chars = max_size_chars
+        self.current_length = 0
+        self._truncated = False
+    
+    def iterencode(self, o, _one_shot=False):
+        """
+        Encode the object incrementally, checking size after each chunk.
+        Stops early if max_size_chars is exceeded.
+        """
+        self.current_length = 0
+        self._truncated = False
+        
+        # Use _one_shot optimization when possible (faster for simple objects)
+        # The parent class will determine if _one_shot is safe
+        for chunk in super().iterencode(o, _one_shot):
+            self.current_length += len(chunk)
+            if self.current_length > self.max_size_chars:
+                self._truncated = True
+                # Stop yielding chunks when limit is exceeded
+                break
+            yield chunk
+
+
 def safe_json_dumps(value: Any) -> str:
     """
     Safely serialize a value to JSON string with safeguards against:
@@ -329,68 +373,45 @@ def safe_json_dumps(value: Any) -> str:
     max_size_chars = AIQA_MAX_OBJECT_STR_CHARS
     visited: Set[int] = set()
     
-    # Convert the entire structure to ensure circular references are detected
+    # Convert the entire structure to json-friendy form, and ensure circular references are detected
     # across the whole object graph
     try:
         converted = object_to_dict(value, visited)
     except Exception as e:
-        # If conversion fails, try with a fresh visited set and json default handler
-        logger.debug(f"object_to_dict failed for {type(value).__name__}, trying json.dumps with default handler: {e}")
-        try:
-            json_str = json.dumps(value, default=json_default_handler_factory(set()))
-            if len(json_str) > max_size_chars:
-                return f"<object {type(value)} too large: {len(json_str)} chars (limit: {max_size_chars} chars) begins: {json_str[:100]}... conversion error: {e}>"
-            return json_str
-        except Exception as e2:
-            logger.debug(f"json.dumps with default handler also failed for {type(value).__name__}: {e2}")
-            return safe_str_repr(value)
+        # Note: object_to_dict is very defensive but can still raise in rare edge cases:
+        # - Objects with corrupted type metadata causing isinstance()/hasattr() to fail
+        # - Malformed dataclasses causing dataclasses.fields() to raise
+        # - Objects where accessing __dict__ or __slots__ triggers descriptors that raise
+        logger.debug(f"object_to_dict failed for {type(value).__name__}, using safe_str_repr. Error: {e}")
+        return safe_str_repr(value)
     
-    # Try JSON serialization of the converted structure
+    # Try JSON serialization of the converted structure with size-limited encoder
+    # After object_to_dict(), converted is a plain dict/list with circular refs already
+    # converted to "<circular reference>" strings. We use check_circular=True (default)
+    # as an additional safety net, though it's redundant since object_to_dict() already
+    # handled circular refs. We don't need a default handler here since converted
+    # should be JSON-serializable.
     try:
-        json_str = json.dumps(converted, default=json_default_handler_factory(set()))
-        # Check size
-        if len(json_str) > max_size_chars:
+        encoder = SizeLimitedJSONEncoder(
+            max_size_chars=max_size_chars,
+            check_circular=True,  # Safety net for dict/list circular refs (redundant but harmless)
+            ensure_ascii=False
+        )
+        # Use iterencode to get chunks and check size incrementally
+        chunks = []
+        for chunk in encoder.iterencode(converted, _one_shot=True):
+            chunks.append(chunk)
+            if encoder._truncated:
+                # Hit the limit, stop early
+                json_str = ''.join(chunks)
+                return f"<object {type(value)} too large: {len(json_str)} chars (limit: {max_size_chars} chars) begins: {json_str[:100]}...>"
+        json_str = ''.join(chunks)
+        # Check if truncation occurred (encoder may have stopped after last chunk)
+        if encoder._truncated or len(json_str) > max_size_chars:
             return f"<object {type(value)} too large: {len(json_str)} chars (limit: {max_size_chars} chars) begins: {json_str[:100]}...>"
         return json_str
     except Exception as e:
-        logger.debug(f"json.dumps total fail for {type(value).__name__}: {e2}")
+        logger.debug(f"json.dumps total fail for {type(value).__name__}: {e}")
         # Final fallback
         return safe_str_repr(value)
 
-
-def json_default_handler_factory(visited: Set[int]) -> Callable[[Any], Any]:
-    """
-    Create a JSON default handler with a shared visited set for circular reference detection.
-    """
-    def handler(obj: Any) -> Any:
-        # Handle datetime objects
-        if isinstance(obj, datetime):
-            return obj.isoformat()
-        if isinstance(obj, date):
-            return obj.isoformat()
-        if isinstance(obj, time):
-            return obj.isoformat()
-        
-        # Handle bytes
-        if isinstance(obj, bytes):
-            try:
-                return obj.decode('utf-8')
-            except UnicodeDecodeError:
-                return f"<bytes: {len(obj)} bytes>"
-        
-        # Try object conversion with the shared visited set
-        try:
-            return object_to_dict(obj, visited)
-        except Exception:
-            return safe_str_repr(obj)
-    
-    return handler
-
-
-def json_default_handler(obj: Any) -> Any:
-    """
-    Default handler for JSON serialization of non-serializable objects.
-    This is a fallback that creates its own visited set.
-    """
-    return json_default_handler_factory(set())(obj)
-