aiqa-client 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

aiqa/experiment_runner.py

@@ -0,0 +1,336 @@
+"""
+ExperimentRunner - runs experiments on datasets and scores results
+"""
+
+import os
+import time
+from typing import Any, Dict, List, Optional, Callable, Awaitable, Union
+import requests
+
+
+class ExperimentRunner:
+    """
+    The ExperimentRunner is the main class for running experiments on datasets.
+    It can create an experiment, run it, and score the results.
+    Handles setting up environment variables and passing parameters to the engine function.
+    """
+
+    def __init__(
+        self,
+        dataset_id: str,
+        experiment_id: Optional[str] = None,
+        server_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        organisation_id: Optional[str] = None,
+    ):
+        """
+        Initialize the ExperimentRunner.
+
+        Args:
+            dataset_id: ID of the dataset to run experiments on
+            experiment_id: Usually unset, and a fresh experiment is created with a random ID
+            server_url: URL of the AIQA server (defaults to AIQA_SERVER_URL env var)
+            api_key: API key for authentication (defaults to AIQA_API_KEY env var)
+            organisation_id: Organisation ID for the experiment
+        """
+        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.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
+
+    def get_dataset(self) -> Dict[str, Any]:
+        """
+        Fetch the dataset to get its metrics.
+
+        Returns:
+            The dataset object with metrics and other information
+        """
+        response = requests.get(
+            f"{self.server_url}/dataset/{self.dataset_id}",
+            headers=self._get_headers(),
+        )
+
+        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}"
+            )
+
+        return response.json()
+
+    def get_example_inputs(self, limit: int = 10000) -> List[Dict[str, Any]]:
+        """
+        Fetch example inputs from the dataset.
+
+        Args:
+            limit: Maximum number of examples to fetch (default: 10000)
+
+        Returns:
+            List of example objects
+        """
+        params = {
+            "dataset_id": self.dataset_id,
+            "limit": str(limit),
+        }
+        if self.organisation:
+            params["organisation"] = self.organisation
+
+        response = requests.get(
+            f"{self.server_url}/example",
+            params=params,
+            headers=self._get_headers(),
+        )
+
+        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}"
+            )
+
+        data = response.json()
+        return data.get("hits", [])
+
+    def create_experiment(
+        self, experiment_setup: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Create an experiment if one does not exist.
+
+        Args:
+            experiment_setup: Optional setup for the experiment object. You may wish to set:
+                - name (recommended for labelling the experiment)
+                - parameters
+                - comparison_parameters
+
+        Returns:
+            The created experiment object
+        """
+        if not self.organisation or not self.dataset_id:
+            raise Exception("Organisation and dataset ID are required to create an experiment")
+
+        if not experiment_setup:
+            experiment_setup = {}
+
+        # Fill in if not set
+        experiment_setup = {
+            **experiment_setup,
+            "organisation": self.organisation,
+            "dataset": self.dataset_id,
+            "results": [],
+            "summary_results": {},
+        }
+
+        print("Creating experiment")
+        response = requests.post(
+            f"{self.server_url}/experiment",
+            json=experiment_setup,
+            headers=self._get_headers(),
+        )
+
+        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}"
+            )
+
+        experiment = response.json()
+        self.experiment_id = experiment["id"]
+        self.experiment = experiment
+        return experiment
+
+    def score_and_store(
+        self,
+        example: Dict[str, Any],
+        result: Any,
+        scores: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Ask the server to score an example result. Stores the score for later summary calculation.
+
+        Args:
+            example: The example object
+            result: The output from running the engine on the example
+            scores: Optional pre-computed scores
+
+        Returns:
+            The score result from the server
+        """
+        # Do we have an experiment ID? If not, we need to create the experiment first
+        if not self.experiment_id:
+            self.create_experiment()
+
+        if scores is None:
+            scores = {}
+
+        print(f"Scoring and storing example: {example['id']}")
+        print(f"Scores: {scores}")
+
+        response = requests.post(
+            f"{self.server_url}/experiment/{self.experiment_id}/example/{example['id']}/scoreAndStore",
+            json={
+                "output": result,
+                "traceId": example.get("traceId"),
+                "scores": scores,
+            },
+            headers=self._get_headers(),
+        )
+
+        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}"
+            )
+
+        json_result = response.json()
+        print(f"scoreAndStore response: {json_result}")
+        return json_result
+
+    async def run(
+        self,
+        engine: Callable[[Any], Union[Any, Awaitable[Any]]],
+        scorer: Optional[
+            Callable[[Any, Dict[str, Any]], Awaitable[Dict[str, Any]]]
+        ] = None,
+    ) -> None:
+        """
+        Run an engine function on all examples and score the results.
+
+        Args:
+            engine: Function that takes input, returns output (can be async)
+            scorer: Optional function that scores the output given the example
+        """
+        examples = self.get_example_inputs()
+
+        # Wrap engine to match run_example signature (input, parameters)
+        def wrapped_engine(input_data, parameters):
+            return engine(input_data)
+
+        # Wrap scorer to match run_example signature (output, example, parameters)
+        async def wrapped_scorer(output, example, parameters):
+            if scorer:
+                return await scorer(output, example)
+            return {}
+
+        for example in examples:
+            scores = await self.run_example(example, wrapped_engine, wrapped_scorer)
+            if scores:
+                self.scores.append(
+                    {
+                        "example": example,
+                        "result": scores,
+                        "scores": scores,
+                    }
+                )
+
+    async def run_example(
+        self,
+        example: Dict[str, Any],
+        call_my_code: Callable[[Any, Dict[str, Any]], Union[Any, Awaitable[Any]]],
+        score_this_output: Optional[
+            Callable[[Any, Dict[str, Any], Dict[str, Any]], Awaitable[Dict[str, Any]]]
+        ] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Run the engine on an example with the given parameters (looping over comparison parameters),
+        and score the result. Also calls scoreAndStore to store the result in the server.
+
+        Args:
+            example: The example to run
+            call_my_code: Function that takes input and parameters, returns output (can be async)
+            score_this_output: Optional function that scores the output given the example and parameters
+
+        Returns:
+            One set of scores for each comparison parameter set. If no comparison parameters,
+            returns an array of one.
+        """
+        # Ensure experiment exists
+        if not self.experiment:
+            self.create_experiment()
+        if not self.experiment:
+            raise Exception("Failed to create experiment")
+
+        # Make the parameters
+        parameters_fixed = self.experiment.get("parameters") or {}
+        # If comparison_parameters is empty/undefined, default to [{}] so we run at least once
+        parameters_loop = self.experiment.get("comparison_parameters") or [{}]
+
+        # Handle both spans array and input field
+        input_data = example.get("input")
+        if not input_data and example.get("spans") and len(example["spans"]) > 0:
+            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}"
+            )
+            # Run engine anyway -- this could make sense if it's all about the parameters
+
+        all_scores: List[Dict[str, Any]] = []
+        # This loop should not be parallelized - it should run sequentially, one after the other
+        # to avoid creating interference between the runs.
+        for parameters in parameters_loop:
+            parameters_here = {**parameters_fixed, **parameters}
+            print(f"Running with parameters: {parameters_here}")
+
+            # Set env vars from parameters_here
+            for key, value in parameters_here.items():
+                if value:
+                    os.environ[key] = str(value)
+
+            start = time.time() * 1000  # milliseconds
+            output = call_my_code(input_data, parameters_here)
+            # Handle async functions
+            if hasattr(output, "__await__"):
+                import asyncio
+
+                output = await output
+            end = time.time() * 1000  # milliseconds
+            duration = int(end - start)
+
+            print(f"Output: {output}")
+
+            scores: Dict[str, Any] = {}
+            if score_this_output:
+                scores = await score_this_output(output, example, parameters_here)
+
+            scores["duration"] = duration
+
+            # TODO: this call as async and wait for all to complete before returning
+            print(f"Call scoreAndStore ... for example: {example['id']} with scores: {scores}")
+            result = self.score_and_store(example, output, scores)
+            print(f"scoreAndStore returned: {result}")
+            all_scores.append(result)
+
+        return all_scores
+
+    def get_summary_results(self) -> Dict[str, Any]:
+        """
+        Get summary results from the experiment.
+
+        Returns:
+            Dictionary of metric names to summary statistics
+        """
+        response = requests.get(
+            f"{self.server_url}/experiment/{self.experiment_id}",
+            headers=self._get_headers(),
+        )
+
+        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}"
+            )
+
+        experiment2 = response.json()
+        return experiment2.get("summary_results", {})
+

aiqa/object_serialiser.py

@@ -0,0 +1,361 @@
+"""
+Object serialization utilities for converting Python objects to JSON-safe formats.
+Handles objects, dataclasses, circular references, and size limits.
+"""
+
+import json
+import os
+import dataclasses
+from datetime import datetime, date, time
+from typing import Any, Callable, Set
+
+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)"""
+    if value is None:
+        return 0
+    if isinstance(value, int):
+        return value    
+    if value.endswith("b"): # drop the b
+        value = value[:-1]
+    if value.endswith("g"):
+        return int(value[:-1]) * 1024 * 1024 * 1024
+    elif value.endswith("m"):
+        return int(value[:-1]) * 1024 * 1024
+    elif value.endswith("k"):
+        return int(value[:-1]) * 1024
+    return int(value)
+
+
+# Configurable limit for object string representation (in characters)
+AIQA_MAX_OBJECT_STR_CHARS = toNumber(os.getenv("AIQA_MAX_OBJECT_STR_CHARS", "1m"))
+
+# Data filters configuration
+def _get_enabled_filters() -> Set[str]:
+    """Get set of enabled filter names from AIQA_DATA_FILTERS env var."""
+    filters_env = os.getenv("AIQA_DATA_FILTERS", "RemovePasswords, RemoveJWT, RemoveAuthHeaders, RemoveAPIKeys")
+    if not filters_env or filters_env.lower() == "false":
+        return set()
+    return {f.strip() for f in filters_env.split(",") if f.strip()}
+
+_ENABLED_FILTERS = _get_enabled_filters()
+
+def _is_jwt_token(value: Any) -> bool:
+    """Check if a value looks like a JWT token (starts with 'eyJ' and has 3 parts separated by dots)."""
+    if not isinstance(value, str):
+        return False
+    # JWT tokens have format: header.payload.signature (3 parts separated by dots)
+    # They typically start with 'eyJ' (base64 encoded '{"')
+    parts = value.split('.')
+    return len(parts) == 3 and value.startswith('eyJ') and all(len(p) > 0 for p in parts)
+
+def _is_api_key(value: Any) -> bool:
+    """Check if a value looks like an API key based on common patterns."""
+    if not isinstance(value, str):
+        return False
+    value = value.strip()
+    # Common API key prefixes:
+    api_key_prefixes = [
+        'sk-',    # OpenAI secret key
+        'pk-',    # possibly public key
+        'AKIA',   # AWS access key
+        'ghp_',   # GitHub personal access token
+        'gho_',   # GitHub OAuth token
+        'ghu_',   # GitHub unidentified token
+        'ghs_',   # GitHub SAML token
+        'ghr_'    # GitHub refresh token
+    ]
+    return any(value.startswith(prefix) for prefix in api_key_prefixes)
+
+def _apply_data_filters(key: str, value: Any) -> Any:
+    """Apply data filters to a key-value pair based on enabled filters."""
+    if not value:  # Don't filter falsy values
+        return value
+    
+    key_lower = str(key).lower()
+    
+    # RemovePasswords filter: if key contains "password", replace value with "****"
+    if "RemovePasswords" in _ENABLED_FILTERS and "password" in key_lower:
+        return "****"
+    
+    # RemoveJWT filter: if value looks like a JWT token, replace with "****"
+    if "RemoveJWT" in _ENABLED_FILTERS and _is_jwt_token(value):
+        return "****"
+    
+    # RemoveAuthHeaders filter: if key is "authorization" (case-insensitive), replace value with "****"
+    if "RemoveAuthHeaders" in _ENABLED_FILTERS and key_lower == "authorization":
+        return "****"
+    
+    # RemoveAPIKeys filter: if key contains API key patterns or value looks like an API key
+    if "RemoveAPIKeys" in _ENABLED_FILTERS:
+        # Check key patterns
+        api_key_key_patterns = ['api_key', 'apikey', 'api-key', 'apikey']
+        if any(pattern in key_lower for pattern in api_key_key_patterns):
+            return "****"
+        # Check value patterns
+        if _is_api_key(value):
+            return "****"
+    
+    return value
+
+
+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.
+    
+    Handles objects by attempting to convert them to dicts, with safeguards against:
+    - Circular references
+    - Unconvertible parts
+    - Large objects (size limits)
+    """
+    # Keep primitives as is (including None)
+    if value is None or isinstance(value, (str, int, float, bool, bytes)):
+        return value
+    
+    # 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)
+    
+    # For dicts and other complex types, serialize to JSON string
+    try:
+        return safe_json_dumps(value)
+    except Exception:
+        # If JSON serialization fails, convert to string
+        return safe_str_repr(value)
+
+
+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.
+    """
+    try:
+        # Try __repr__ first (usually more informative)
+        repr_str = repr(value)
+        # Limit length to avoid huge strings
+        if len(repr_str) > AIQA_MAX_OBJECT_STR_CHARS:
+            return repr_str[:AIQA_MAX_OBJECT_STR_CHARS] + "... (truncated)"
+        return repr_str
+    except Exception:
+        # Fallback to type name
+        try:
+            return f"<{type(value).__name__} object>"
+        except Exception:
+            return "<unknown object>"
+
+
+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.
+    
+    Args:
+        obj: The object to convert
+        visited: Set of object IDs to detect circular references
+        max_depth: Maximum recursion depth
+        current_depth: Current recursion depth
+    
+    Returns:
+        Dictionary representation of the object, or a string if conversion fails
+    """
+    if current_depth > max_depth:
+        return "<max depth exceeded>"
+    
+    obj_id = id(obj)
+    if obj_id in visited:
+        return "<circular reference>"
+    
+    # Handle None
+    if obj is None:
+        return None
+    
+    # Handle primitives
+    if isinstance(obj, (str, int, float, bool, bytes)):
+        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()
+    
+    # Handle dict
+    if isinstance(obj, dict):
+        visited.add(obj_id)
+        try:
+            result = {}
+            for k, v in obj.items():
+                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)
+            visited.remove(obj_id)
+            return result
+        except Exception:
+            visited.discard(obj_id)
+            return safe_str_repr(obj)
+    
+    # Handle list/tuple
+    if isinstance(obj, (list, tuple)):
+        visited.add(obj_id)
+        try:
+            result = [object_to_dict(item, visited, max_depth, current_depth + 1) for item in obj]
+            visited.remove(obj_id)
+            return result
+        except Exception:
+            visited.discard(obj_id)
+            return safe_str_repr(obj)
+    
+    # Handle dataclasses
+    if dataclasses.is_dataclass(obj):
+        visited.add(obj_id)
+        try:
+            result = {}
+            for field in dataclasses.fields(obj):
+                value = getattr(obj, field.name, None)
+                filtered_value = _apply_data_filters(field.name, value)
+                result[field.name] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
+            visited.remove(obj_id)
+            return result
+        except Exception:
+            visited.discard(obj_id)
+            return safe_str_repr(obj)
+    
+    # Handle objects with __dict__
+    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:
+            visited.discard(obj_id)
+            return safe_str_repr(obj)
+    
+    # Handle objects with __slots__
+    if hasattr(obj, "__slots__"):
+        visited.add(obj_id)
+        try:
+            result = {}
+            for slot in obj.__slots__:
+                if hasattr(obj, slot):
+                    value = getattr(obj, slot, None)
+                    filtered_value = _apply_data_filters(slot, value)
+                    result[slot] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
+            visited.remove(obj_id)
+            return result
+        except Exception:
+            visited.discard(obj_id)
+            return safe_str_repr(obj)
+    
+    # Fallback: try to get a few common attributes
+    try:
+        result = {}
+        for attr in ["name", "id", "value", "type", "status"]:
+            if hasattr(obj, attr):
+                value = getattr(obj, attr, None)
+                filtered_value = _apply_data_filters(attr, value)
+                result[attr] = object_to_dict(filtered_value, visited, max_depth, current_depth + 1)
+        if result:
+            return result
+    except Exception:
+        pass
+    
+    # Final fallback: string representation
+    return safe_str_repr(obj)
+
+
+def safe_json_dumps(value: Any) -> str:
+    """
+    Safely serialize a value to JSON string with safeguards against:
+    - Circular references
+    - Large objects (size limits)
+    - Unconvertible parts
+    
+    Args:
+        value: The value to serialize
+
+    Uses AIQA_MAX_OBJECT_STR_CHARS environment variable (default: 1000000) to limit length.
+    
+    Returns:
+        JSON string representation
+    """    
+    max_size_chars = AIQA_MAX_OBJECT_STR_CHARS
+    visited: Set[int] = set()
+    
+    # Convert the entire structure to ensure circular references are detected
+    # across the whole object graph
+    try:
+        converted = object_to_dict(value, visited)
+    except Exception:
+        # If conversion fails, try with a fresh visited set and json default handler
+        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:
+            return safe_str_repr(value)
+    
+    # Try JSON serialization of the converted structure
+    try:
+        json_str = json.dumps(converted, default=json_default_handler_factory(set()))
+        # Check size
+        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]}...>"
+        return json_str
+    except Exception:
+        # 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)
+