isage-data 0.2.1.8__cp311-none-any.whl

sage/data/sources/agent_tools/tests/test_agent_tools_loader.py

@@ -0,0 +1,306 @@
+"""
+Unit tests for Agent Tools DataLoader
+
+Tests cover:
+- Data loading and validation
+- Tool ID format validation
+- Search and retrieval operations
+- Category indexing
+- Deduplication checks
+- Coverage metrics
+"""
+
+import re
+
+import pytest
+
+from sage.data.sources.agent_tools import AgentToolRecord, AgentToolsDataLoader
+
+
+class TestAgentToolsDataLoader:
+    """Test suite for AgentToolsDataLoader."""
+
+    @pytest.fixture
+    def loader(self):
+        """Create a loader instance for testing."""
+        return AgentToolsDataLoader()
+
+    def test_loader_initialization(self, loader):
+        """Test that loader initializes successfully."""
+        assert loader is not None
+        assert len(loader) > 0
+        assert loader.get_total_tools() > 0
+
+    def test_minimum_tool_count(self, loader):
+        """Test that we have at least 1000 tools."""
+        assert len(loader) >= 1000, f"Expected >= 1000 tools, got {len(loader)}"
+
+    def test_tool_id_format(self, loader):
+        """Test that all tool_ids match required regex pattern."""
+        pattern = re.compile(r"^[a-z]+(_[a-z]+)*_[0-9]{3}$")
+
+        invalid_ids = []
+        for tool_id in loader.list_tool_ids():
+            if not pattern.match(tool_id):
+                invalid_ids.append(tool_id)
+
+        assert len(invalid_ids) == 0, f"Invalid tool_ids found: {invalid_ids[:10]}"
+
+    def test_tool_id_uniqueness(self, loader):
+        """Test that all tool_ids are unique."""
+        tool_ids = loader.list_tool_ids()
+        assert len(tool_ids) == len(set(tool_ids)), "Duplicate tool_ids found"
+
+    def test_tool_name_uniqueness(self, loader):
+        """Test that all tool names are unique."""
+        names = [tool.name for tool in loader.tools.values()]
+        duplicates = [name for name in names if names.count(name) > 1]
+        assert len(duplicates) == 0, f"Duplicate names found: {set(duplicates)}"
+
+    def test_get_tool(self, loader):
+        """Test getting a tool by ID."""
+        # Get first tool_id
+        tool_ids = loader.list_tool_ids()
+        assert len(tool_ids) > 0
+
+        tool = loader.get_tool(tool_ids[0])
+        assert isinstance(tool, AgentToolRecord)
+        assert tool.tool_id == tool_ids[0]
+
+    def test_get_tool_invalid_id(self, loader):
+        """Test that getting invalid tool_id raises KeyError."""
+        with pytest.raises(KeyError):
+            loader.get_tool("nonexistent_tool_999")
+
+    def test_capabilities_non_empty(self, loader):
+        """Test that all tools have non-empty capabilities."""
+        for tool in loader.tools.values():
+            assert len(tool.capabilities) > 0, f"Tool {tool.tool_id} has empty capabilities"
+
+    def test_category_format(self, loader):
+        """Test that all categories follow path format."""
+        category_pattern = re.compile(r"^[a-z]+(/[a-z_]+)*$")
+
+        invalid_categories = []
+        for tool in loader.tools.values():
+            if not category_pattern.match(tool.category):
+                invalid_categories.append(tool.category)
+
+        assert len(invalid_categories) == 0, f"Invalid categories: {set(invalid_categories)}"
+
+    def test_category_index(self, loader):
+        """Test that category index is built correctly."""
+        categories = loader.get_categories()
+        assert len(categories) > 0
+
+        # Test that each category has tools
+        for category in categories:
+            tools_in_cat = list(loader.iter_category(category))
+            assert len(tools_in_cat) > 0, f"Category {category} has no tools"
+
+    def test_search_by_capability(self, loader):
+        """Test capability-based search."""
+        # Search for common capability
+        results = loader.search_by_capability("forecast", top_k=10)
+        assert len(results) > 0, "No tools found with 'forecast' capability"
+
+        # Verify results have matching capability
+        for tool in results:
+            assert any("forecast" in cap for cap in tool.capabilities), \
+                f"Tool {tool.tool_id} doesn't have 'forecast' in capabilities"
+
+    def test_search_top_k_limit(self, loader):
+        """Test that search respects top_k parameter."""
+        results = loader.search_by_capability("search", top_k=5)
+        assert len(results) <= 5, "Search returned more than top_k results"
+
+    def test_iter_category(self, loader):
+        """Test category iteration."""
+        categories = loader.get_categories()
+        assert len(categories) > 0
+
+        # Test first category
+        category = categories[0]
+        count = 0
+        for tool in loader.iter_category(category):
+            assert tool.category == category
+            count += 1
+
+        assert count > 0, f"No tools found in category {category}"
+
+    def test_iter_category_invalid(self, loader):
+        """Test that iterating invalid category raises ValueError."""
+        with pytest.raises(ValueError):
+            list(loader.iter_category("invalid/category"))
+
+    def test_load_taxonomy(self, loader):
+        """Test loading category taxonomy."""
+        taxonomy = loader.load_taxonomy()
+        assert taxonomy is not None
+        assert len(taxonomy.taxonomy) > 0
+        assert taxonomy.version is not None
+
+    def test_load_stats(self, loader):
+        """Test loading dataset statistics."""
+        stats = loader.load_stats()
+        assert stats is not None
+        assert stats.total_tools > 0
+        assert stats.total_categories > 0
+        assert len(stats.category_distribution) > 0
+
+    def test_stats_accuracy(self, loader):
+        """Test that stats match actual data."""
+        stats = loader.load_stats()
+        assert stats.total_tools == len(loader)
+
+        # Check category distribution
+        for category, count in stats.category_distribution.items():
+            actual_count = len(list(loader.iter_category(category)))
+            assert actual_count == count, \
+                f"Category {category}: stats={count}, actual={actual_count}"
+
+    def test_reliability_scores(self, loader):
+        """Test that reliability scores are valid."""
+        for tool in loader.tools.values():
+            if tool.reliability_score is not None:
+                assert 0.0 <= tool.reliability_score <= 1.0, \
+                    f"Tool {tool.tool_id} has invalid reliability: {tool.reliability_score}"
+
+    def test_latency_values(self, loader):
+        """Test that latency values are non-negative."""
+        for tool in loader.tools.values():
+            if tool.latency_ms_p50 is not None:
+                assert tool.latency_ms_p50 >= 0, \
+                    f"Tool {tool.tool_id} has negative latency: {tool.latency_ms_p50}"
+
+    def test_filter_tools(self, loader):
+        """Test multi-criteria filtering."""
+        # Filter by category
+        results = loader.filter_tools(category="environment/weather")
+        assert all(t.category == "environment/weather" for t in results)
+
+        # Filter by reliability
+        results = loader.filter_tools(min_reliability=0.95)
+        assert all(
+            t.reliability_score is None or t.reliability_score >= 0.95
+            for t in results
+        )
+
+        # Filter by latency
+        results = loader.filter_tools(max_latency=200)
+        assert all(
+            t.latency_ms_p50 is None or t.latency_ms_p50 <= 200
+            for t in results
+        )
+
+    def test_get_tool_by_name(self, loader):
+        """Test getting tool by name."""
+        # Get first tool
+        tool_ids = loader.list_tool_ids()
+        original_tool = loader.get_tool(tool_ids[0])
+
+        # Get by name
+        tool_by_name = loader.get_tool_by_name(original_tool.name)
+        assert tool_by_name.tool_id == original_tool.tool_id
+
+    def test_search_by_name(self, loader):
+        """Test name-based search."""
+        results = loader.search_by_name("weather", top_k=5)
+        assert len(results) > 0
+        assert all("weather" in t.name.lower() for t in results)
+
+    def test_capability_index(self, loader):
+        """Test that capability index is comprehensive."""
+        capabilities = loader.get_capabilities()
+        assert len(capabilities) > 0
+
+        # Each capability should have tools
+        for cap in capabilities[:10]:  # Test first 10
+            tools_with_cap = [
+                t for t in loader.tools.values()
+                if cap in t.capabilities
+            ]
+            assert len(tools_with_cap) > 0, f"No tools with capability: {cap}"
+
+    def test_category_stats(self, loader):
+        """Test category statistics calculation."""
+        categories = loader.get_categories()
+        category = categories[0]
+
+        stats = loader.get_category_stats(category)
+        assert "total_tools" in stats
+        assert "avg_reliability" in stats
+        assert "avg_latency_ms" in stats
+        assert stats["total_tools"] > 0
+
+
+class TestAgentToolRecord:
+    """Test suite for AgentToolRecord schema validation."""
+
+    def test_valid_tool_id(self):
+        """Test valid tool_id formats."""
+        valid_ids = [
+            "weather_query_001",
+            "calendar_event_create_042",
+            "environment_air_quality_015",
+            "a_001",
+            "test_tool_999"
+        ]
+
+        for tool_id in valid_ids:
+            tool = AgentToolRecord(
+                tool_id=tool_id,
+                name="Test Tool",
+                category="test/category",
+                capabilities=["test"]
+            )
+            assert tool.tool_id == tool_id
+
+    def test_invalid_tool_id(self):
+        """Test that invalid tool_ids raise validation error."""
+        invalid_ids = [
+            "WeatherQuery_001",  # uppercase
+            "weather-query_001",  # hyphen
+            "weather_query_1",    # not 3 digits
+            "weather_query_1234", # too many digits
+            "001_weather_query",  # digits not at end
+        ]
+
+        from pydantic import ValidationError
+
+        for tool_id in invalid_ids:
+            with pytest.raises(ValidationError):
+                AgentToolRecord(
+                    tool_id=tool_id,
+                    name="Test Tool",
+                    category="test/category",
+                    capabilities=["test"]
+                )
+
+    def test_empty_capabilities(self):
+        """Test that empty capabilities raise error."""
+        from pydantic import ValidationError
+
+        with pytest.raises(ValidationError):
+            AgentToolRecord(
+                tool_id="test_tool_001",
+                name="Test Tool",
+                category="test/category",
+                capabilities=[]
+            )
+
+    def test_invalid_category(self):
+        """Test that invalid category format raises error."""
+        from pydantic import ValidationError
+
+        with pytest.raises(ValidationError):
+            AgentToolRecord(
+                tool_id="test_tool_001",
+                name="Test Tool",
+                category="Invalid-Category",  # hyphen not allowed
+                capabilities=["test"]
+            )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

sage/data/sources/bbh/__init__.py

@@ -0,0 +1,5 @@
+"""BBH dataset source wrapper."""
+
+from .dataloader import BBHDataLoader
+
+__all__ = ["BBHDataLoader"]

sage/data/sources/bbh/dataset.yaml

@@ -0,0 +1,9 @@
+name: "bbh"
+description: "BIG-Bench Hard - 27 challenging reasoning tasks for language models"
+type: "text"
+format: "json"
+maintainer: "sage-team"
+tags: ["reasoning", "benchmark", "big-bench"]
+size: "~5MB"
+license: "Apache-2.0"
+version: "1.0.0"

sage/data/sources/control_plane_benchmark/__init__.py

@@ -0,0 +1,41 @@
+"""
+Control Plane Benchmark Dataset Module
+
+This module provides data loaders for the Control Plane scheduling benchmark,
+including LLM workloads, hybrid workloads (LLM + Embedding), and test prompts.
+
+Usage:
+    from sage.data.sources.control_plane_benchmark import ControlPlaneBenchmarkDataLoader
+
+    loader = ControlPlaneBenchmarkDataLoader()
+    
+    # List available workloads
+    print(loader.list_workloads())
+    print(loader.list_workloads(category="hybrid"))
+    
+    # Load a workload configuration
+    workload = loader.load_workload("llm_medium")
+    print(f"Requests: {workload.request_count}, Rate: {workload.rate_per_second}")
+    
+    # Load test prompts
+    llm_prompts = loader.load_prompts("llm")
+    embed_texts = loader.load_prompts("embedding")
+"""
+
+from .dataloader import (
+    ControlPlaneBenchmarkDataLoader,
+    EmbeddingText,
+    HybridWorkloadConfig,
+    LLMPrompt,
+    LLMWorkloadConfig,
+    WorkloadConfig,
+)
+
+__all__ = [
+    "ControlPlaneBenchmarkDataLoader",
+    "WorkloadConfig",
+    "LLMWorkloadConfig",
+    "HybridWorkloadConfig",
+    "LLMPrompt",
+    "EmbeddingText",
+]

sage/data/sources/control_plane_benchmark/dataset.yaml

@@ -0,0 +1,101 @@
+name: "control_plane_benchmark"
+description: "Benchmark data for Control Plane scheduling policy evaluation"
+type: "benchmark"
+format: "jsonl"
+version: "0.1.0"
+maintainer: "SAGE Control Plane Team"
+tags: ["control_plane", "scheduling", "benchmark", "llm", "embedding", "hybrid"]
+license: "Apache-2.0"
+size: "~5MB"
+
+categories:
+  - name: "llm_workloads"
+    description: "Pure LLM request workload configurations"
+    files:
+      - name: "light"
+        description: "Light workload: 100 requests, 10 req/s, single model"
+      - name: "medium"
+        description: "Medium workload: 1000 requests, 100 req/s, multi-model"
+      - name: "heavy"
+        description: "Heavy workload: 5000 requests, 500 req/s, multi-priority"
+
+  - name: "hybrid_workloads"
+    description: "Mixed LLM + Embedding workload configurations"
+    files:
+      - name: "balanced"
+        description: "Balanced mix: 50% LLM, 50% Embedding"
+      - name: "llm_heavy"
+        description: "LLM heavy: 80% LLM, 20% Embedding"
+      - name: "embed_heavy"
+        description: "Embedding heavy: 20% LLM, 80% Embedding"
+      - name: "burst"
+        description: "Burst pattern with variable mix ratios"
+
+  - name: "prompts"
+    description: "Test data for benchmark requests"
+    files:
+      - name: "llm_prompts"
+        description: "200+ LLM test prompts with expected token counts"
+      - name: "embed_texts"
+        description: "200+ Embedding test texts with batch configurations"
+
+statistics:
+  llm_workloads: 3
+  hybrid_workloads: 4
+  llm_prompts: 200
+  embed_texts: 200
+
+schema:
+  llm_workload:
+    required:
+      - workload_id
+      - request_count
+      - rate_per_second
+    optional:
+      - arrival_pattern
+      - model_distribution
+      - priority_distribution
+      - prompt_len_range
+      - output_len_range
+      - slo_deadlines
+
+  hybrid_workload:
+    required:
+      - workload_id
+      - request_count
+      - rate_per_second
+      - llm_ratio
+      - embedding_ratio
+    optional:
+      - arrival_pattern
+      - llm_model_distribution
+      - embedding_model
+      - embedding_batch_sizes
+      - priority_distribution
+
+  llm_prompt:
+    required:
+      - prompt_id
+      - text
+    optional:
+      - expected_tokens
+      - priority
+      - category
+      - difficulty
+
+  embed_text:
+    required:
+      - text_id
+      - texts
+    optional:
+      - model
+      - batch_size
+      - category
+
+citation: |
+  @dataset{control_plane_benchmark_2025,
+    title={Control Plane Benchmark: Evaluating LLM Scheduling Policies},
+    author={SAGE Team},
+    year={2025},
+    publisher={SAGE Framework}
+  }

sage/data/sources/gpqa/__init__.py

@@ -0,0 +1,5 @@
+"""GPQA dataset source wrapper."""
+
+from .dataloader import GPQADataLoader
+
+__all__ = ["GPQADataLoader"]

sage/data/sources/gpqa/dataset.yaml

@@ -0,0 +1,10 @@
+name: "gpqa"
+description: "GPQA (Graduate-Level Google-Proof Q&A) - Expert-level questions in Physics, Chemistry, Biology"
+type: "text"
+format: "huggingface-dataset"
+maintainer: "sage-team"
+tags: ["gpqa", "expert-level", "reasoning", "science"]
+size: "~5MB (cached)"
+license: "MIT"
+version: "1.0.0"
+source_url: "https://huggingface.co/datasets/Idavidrein/gpqa"

sage/data/sources/libamm_benchmark/__init__.py

@@ -0,0 +1,10 @@
+"""LibAMM benchmark dataset source.
+
+This provides access to the LibAMM benchmark datasets located in the
+libamm-benchmark directory.
+"""
+
+# Note: The actual data is in ../libamm-benchmark/
+# This wrapper provides a consistent interface through the sources layer
+
+__all__ = []

sage/data/sources/libamm_benchmark/dataset.yaml

@@ -0,0 +1,9 @@
+name: "libamm_benchmark"
+description: "LibAMM benchmark datasets for approximate matrix multiplication evaluation"
+type: "matrix"
+format: "binary/text"
+maintainer: "libamm-team"
+tags: ["matrix", "benchmark", "approximate-multiplication", "ann"]
+size: "~325MB"
+license: "Research Use"
+version: "1.0.0"

sage/data/sources/locomo/__init__.py

@@ -0,0 +1,5 @@
+"""LoCoMo dataset source wrapper."""
+
+from .dataloader import LocomoDataLoader
+
+__all__ = ["LocomoDataLoader"]

sage/data/sources/locomo/dataset.yaml

@@ -0,0 +1,10 @@
+name: "locomo"
+description: "Long-Context Memory (LoCoMo) dataset for long-context conversation and QA evaluation"
+type: "text"
+format: "json"
+maintainer: "sage-team"
+tags: ["long-context", "memory", "conversation", "qa"]
+size: "~2.68MB"
+license: "MIT"
+version: "1.0.0"
+source_url: "https://github.com/your-repo/locomo"