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

aiqa/test_experiment_runner.py

@@ -0,0 +1,176 @@
+"""
+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_tracing.py

@@ -0,0 +1,230 @@
+"""
+Unit tests for tracing.py functions.
+"""
+
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+from aiqa.tracing import get_span
+
+
+class TestGetSpan:
+    """Tests for get_span function."""
+
+    def test_get_span_success_with_span_id(self):
+        """Test successful retrieval of span using spanId query."""
+        span_data = {
+            "id": "test-span-123",
+            "name": "test_span",
+            "trace_id": "abc123",
+            "attributes": {"key": "value"},
+        }
+        mock_response_data = {"hits": [span_data]}
+
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                mock_response = MagicMock()
+                mock_response.status_code = 200
+                mock_response.json.return_value = mock_response_data
+
+                mock_get.return_value = mock_response
+
+                result = get_span("test-span-123")
+
+                assert result == span_data
+                mock_get.assert_called_once()
+                call_args = mock_get.call_args
+                assert call_args[0][0] == "http://localhost:3000/span"
+                assert "q" in call_args[1]["params"]
+                assert call_args[1]["params"]["q"] == "spanId:test-span-123"
+
+    def test_get_span_success_with_client_span_id(self):
+        """Test successful retrieval of span using clientSpanId query when spanId fails."""
+        span_data = {
+            "id": "test-span-123",
+            "name": "test_span",
+            "trace_id": "abc123",
+        }
+        mock_response_data = {"hits": [span_data]}
+
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                # First call returns 404 (spanId not found), second call succeeds (clientSpanId)
+                mock_response_404 = MagicMock()
+                mock_response_404.status_code = 404
+
+                mock_response_200 = MagicMock()
+                mock_response_200.status_code = 200
+                mock_response_200.json.return_value = mock_response_data
+
+                mock_get.side_effect = [mock_response_404, mock_response_200]
+
+                result = get_span("test-span-123")
+
+                assert result == span_data
+                assert mock_get.call_count == 2
+                # Check that second call uses clientSpanId
+                second_call = mock_get.call_args_list[1]
+                assert second_call[1]["params"]["q"] == "clientSpanId:test-span-123"
+
+    def test_get_span_not_found(self):
+        """Test that get_span returns None when span is not found."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                # Both queries return 404
+                mock_response_404 = MagicMock()
+                mock_response_404.status_code = 404
+
+                mock_get.return_value = mock_response_404
+
+                result = get_span("nonexistent-span")
+
+                assert result is None
+                assert mock_get.call_count == 2
+
+    def test_get_span_empty_hits(self):
+        """Test that get_span returns None when hits array is empty."""
+        mock_response_data = {"hits": []}
+
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                mock_response = MagicMock()
+                mock_response.status_code = 200
+                mock_response.json.return_value = mock_response_data
+
+                mock_get.return_value = mock_response
+
+                result = get_span("test-span-123")
+
+                assert result is None
+
+    def test_get_span_missing_server_url(self):
+        """Test that get_span raises ValueError when AIQA_SERVER_URL is not set."""
+        with patch.dict(os.environ, {}, clear=True):
+            with pytest.raises(ValueError, match="AIQA_SERVER_URL is not set"):
+                get_span("test-span-123")
+
+    def test_get_span_missing_organisation_id(self):
+        """Test that get_span raises ValueError when organisation ID is not provided."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+            },
+            clear=True,
+        ):
+            with pytest.raises(ValueError, match="Organisation ID is required"):
+                get_span("test-span-123")
+
+    def test_get_span_missing_api_key(self):
+        """Test that get_span raises ValueError when AIQA_API_KEY is not set."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+            clear=True,
+        ):
+            with pytest.raises(ValueError, match="API key is required"):
+                get_span("test-span-123")
+
+    def test_get_span_with_organisation_id_parameter(self):
+        """Test that get_span uses organisation_id parameter when provided."""
+        span_data = {"id": "test-span-123", "name": "test_span"}
+        mock_response_data = {"hits": [span_data]}
+
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+            },
+            clear=True,
+        ):
+            with patch("requests.get") as mock_get:
+                mock_response = MagicMock()
+                mock_response.status_code = 200
+                mock_response.json.return_value = mock_response_data
+
+                mock_get.return_value = mock_response
+
+                result = get_span("test-span-123", organisation_id="param-org")
+
+                assert result == span_data
+                call_args = mock_get.call_args
+                assert call_args[1]["params"]["organisation"] == "param-org"
+
+    def test_get_span_server_error(self):
+        """Test that get_span raises ValueError on server error."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                mock_response = MagicMock()
+                mock_response.status_code = 500
+                mock_response.text = "Internal Server Error"
+
+                mock_get.return_value = mock_response
+
+                with pytest.raises(ValueError, match="Failed to get span: 500"):
+                    get_span("test-span-123")
+
+    def test_get_span_authorization_header(self):
+        """Test that get_span includes Authorization header with API key."""
+        span_data = {"id": "test-span-123"}
+        mock_response_data = {"hits": [span_data]}
+
+        with patch.dict(
+            os.environ,
+            {
+                "AIQA_SERVER_URL": "http://localhost:3000",
+                "AIQA_API_KEY": "test-api-key-123",
+                "AIQA_ORGANISATION_ID": "test-org",
+            },
+        ):
+            with patch("requests.get") as mock_get:
+                mock_response = MagicMock()
+                mock_response.status_code = 200
+                mock_response.json.return_value = mock_response_data
+
+                mock_get.return_value = mock_response
+
+                get_span("test-span-123")
+
+                call_args = mock_get.call_args
+                assert call_args[1]["headers"]["Authorization"] == "ApiKey test-api-key-123"