isage-data 0.2.1.8__cp311-none-any.whl

sage/data/sources/agent_sft/tests/test_agent_sft_loader.py

@@ -0,0 +1,316 @@
+"""
+Unit tests for Agent SFT DataLoader
+
+Tests batch sampling, tool coverage, split validation, and turn structure.
+"""
+
+import pytest
+
+from sage.data.sources.agent_sft.dataloader import AgentSFTDataLoader
+from sage.data.sources.agent_sft.schemas import AgentSFTDialog
+
+
+class TestAgentSFTDataLoader:
+    """Test suite for AgentSFTDataLoader."""
+
+    @pytest.fixture
+    def loader(self):
+        """Create a loader instance for testing."""
+        return AgentSFTDataLoader()
+
+    def test_loader_initialization(self, loader):
+        """Test that loader initializes correctly."""
+        assert loader is not None
+        assert loader.data_path.exists()
+
+    def test_load_dialogs(self, loader):
+        """Test loading dialogs from file."""
+        dialogs = loader._load_dialogs()
+        assert len(dialogs) > 0
+        assert all(isinstance(d, AgentSFTDialog) for d in dialogs)
+
+    def test_split_indexing(self, loader):
+        """Test that splits are correctly indexed."""
+        by_split = loader._index_by_split()
+
+        assert "train" in by_split
+        assert "dev" in by_split
+        assert "test" in by_split
+
+        # Check split proportions (approximately)
+        total = sum(len(v) for v in by_split.values())
+        train_ratio = len(by_split["train"]) / total
+
+        assert train_ratio > 0.75  # Should be around 80%
+        assert len(by_split["dev"]) > 0
+        assert len(by_split["test"]) > 0
+
+    def test_iter_dialogs(self, loader):
+        """Test iteration over dialogs."""
+        # Test train split
+        train_dialogs = list(loader.iter_dialogs("train"))
+        assert len(train_dialogs) > 0
+
+        # Test dev split
+        dev_dialogs = list(loader.iter_dialogs("dev"))
+        assert len(dev_dialogs) > 0
+
+        # Test test split
+        test_dialogs = list(loader.iter_dialogs("test"))
+        assert len(test_dialogs) > 0
+
+        # All should be AgentSFTDialog instances
+        assert all(isinstance(d, AgentSFTDialog) for d in train_dialogs)
+
+    def test_iter_dialogs_invalid_split(self, loader):
+        """Test that invalid split raises ValueError."""
+        with pytest.raises(ValueError, match="Invalid split"):
+            list(loader.iter_dialogs("invalid"))
+
+    def test_sample_batch(self, loader):
+        """Test batch sampling."""
+        batch_size = 8
+        batch = loader.sample_batch(batch_size=batch_size, split="train")
+
+        assert len(batch) == batch_size
+        assert all(isinstance(d, AgentSFTDialog) for d in batch)
+        assert all(d.split == "train" for d in batch)
+
+    def test_sample_batch_no_shuffle(self, loader):
+        """Test deterministic batch sampling without shuffle."""
+        batch1 = loader.sample_batch(batch_size=5, split="train", shuffle=False)
+        batch2 = loader.sample_batch(batch_size=5, split="train", shuffle=False)
+
+        # Without shuffle, same batch should be returned
+        assert [d.dialog_id for d in batch1] == [d.dialog_id for d in batch2]
+
+    def test_sample_batch_oversized(self, loader):
+        """Test sampling when batch_size exceeds available data."""
+        # Request more than available in dev set
+        batch = loader.sample_batch(batch_size=10000, split="dev")
+
+        # Should return all available
+        dev_dialogs = list(loader.iter_dialogs("dev"))
+        assert len(batch) == len(dev_dialogs)
+
+    def test_get_tools_coverage(self, loader):
+        """Test tool usage coverage analysis."""
+        coverage = loader.get_tools_coverage()
+
+        assert isinstance(coverage, dict)
+        assert len(coverage) > 0
+
+        # All keys should be valid tool IDs
+        for tool_id in coverage.keys():
+            assert isinstance(tool_id, str)
+            assert "_" in tool_id  # Should follow pattern
+
+        # All values should be positive integers
+        assert all(isinstance(count, int) and count > 0 for count in coverage.values())
+
+    def test_get_stats(self, loader):
+        """Test dataset statistics computation."""
+        stats = loader.get_stats()
+
+        assert stats.total_dialogs > 0
+        assert stats.train_count > 0
+        assert stats.dev_count > 0
+        assert stats.test_count > 0
+        assert stats.total_dialogs == stats.train_count + stats.dev_count + stats.test_count
+
+        assert stats.avg_turns > 0
+        assert stats.unique_tools > 0
+        assert stats.avg_tools_per_dialog > 0
+
+        assert isinstance(stats.tool_coverage, dict)
+
+    def test_get_dialog(self, loader):
+        """Test fetching specific dialog by ID."""
+        # Get a dialog from train set
+        train_dialogs = list(loader.iter_dialogs("train"))
+        if train_dialogs:
+            dialog_id = train_dialogs[0].dialog_id
+
+            # Fetch it
+            fetched = loader.get_dialog(dialog_id)
+            assert fetched is not None
+            assert fetched.dialog_id == dialog_id
+
+        # Test non-existent ID
+        non_existent = loader.get_dialog("sft_999999")
+        assert non_existent is None
+
+    def test_filter_by_difficulty(self, loader):
+        """Test filtering dialogs by difficulty level."""
+        hard_dialogs = loader.filter_by_difficulty("hard", split="train")
+
+        # Should return a list
+        assert isinstance(hard_dialogs, list)
+
+        # All should have hard difficulty
+        for dialog in hard_dialogs:
+            assert dialog.metadata.get("difficulty") == "hard"
+
+    def test_filter_by_tool(self, loader):
+        """Test filtering dialogs by tool usage."""
+        # Get a tool that exists in the dataset
+        coverage = loader.get_tools_coverage()
+        if coverage:
+            tool_id = list(coverage.keys())[0]
+
+            filtered = loader.filter_by_tool(tool_id, split="train")
+
+            # Should return a list
+            assert isinstance(filtered, list)
+
+            # All should use the specified tool
+            for dialog in filtered:
+                assert tool_id in dialog.target_tools
+
+    def test_dialog_turn_structure(self, loader):
+        """Test that dialogs have proper turn structure."""
+        for dialog in loader.iter_dialogs("train"):
+            # Check turn count
+            assert 6 <= len(dialog.turns) <= 12, f"Dialog {dialog.dialog_id} has {len(dialog.turns)} turns"
+
+            # Check that each turn has required fields
+            for turn in dialog.turns:
+                assert turn.role in ["user", "assistant", "tool"]
+                assert turn.content is not None
+
+                # Tool turns should have tool_id
+                if turn.role == "tool":
+                    assert turn.tool_id is not None
+                    assert turn.result is not None
+
+            # Only check first 10 dialogs for performance
+            if dialog.dialog_id == "sft_000010":
+                break
+
+    def test_tool_id_format(self, loader):
+        """Test that all tool IDs follow the correct format."""
+        import re
+        pattern = re.compile(r"^[a-z]+(_[a-z]+)*_[0-9]{3}$")
+
+        coverage = loader.get_tools_coverage()
+        for tool_id in coverage.keys():
+            assert pattern.match(tool_id), f"Invalid tool_id format: {tool_id}"
+
+    def test_dialog_id_format(self, loader):
+        """Test that all dialog IDs follow the correct format."""
+        import re
+        pattern = re.compile(r"^sft_\d{6}$")
+
+        for dialog in loader.iter_dialogs("train"):
+            assert pattern.match(dialog.dialog_id), f"Invalid dialog_id: {dialog.dialog_id}"
+
+            # Only check first 10 for performance
+            if dialog.dialog_id == "sft_000010":
+                break
+
+    def test_split_assignment(self, loader):
+        """Test that split field matches actual split assignment."""
+        for split_name in ["train", "dev", "test"]:
+            dialogs = list(loader.iter_dialogs(split_name))
+
+            # All dialogs in this split should have the correct split field
+            for dialog in dialogs:
+                assert dialog.split == split_name
+
+    def test_lazy_loading(self, loader):
+        """Test that data is loaded lazily."""
+        # Initially, internal cache should be None
+        new_loader = AgentSFTDataLoader()
+        assert new_loader._dialogs is None
+
+        # After first access, should be loaded
+        new_loader._load_dialogs()
+        assert new_loader._dialogs is not None
+
+        # Second call should use cache
+        cached = new_loader._load_dialogs()
+        assert cached is new_loader._dialogs
+
+
+class TestAgentSFTDialogSchema:
+    """Test suite for AgentSFTDialog schema validation."""
+
+    def test_valid_dialog(self):
+        """Test that a valid dialog passes validation."""
+        from sage.data.sources.agent_sft.schemas import AgentSFTDialog
+
+        dialog_data = {
+            "dialog_id": "sft_000001",
+            "goal": "Test goal",
+            "turns": [
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+                {"role": "tool", "tool_id": "test_tool_001", "content": "Tool executed", "result": "{}"},
+                {"role": "user", "content": "Another message"},
+                {"role": "assistant", "content": "Another response"},
+                {"role": "tool", "tool_id": "test_tool_001", "content": "Tool executed", "result": "{}"},
+            ],
+            "target_tools": ["test_tool_001"],
+            "split": "train"
+        }
+
+        dialog = AgentSFTDialog(**dialog_data)
+        assert dialog.dialog_id == "sft_000001"
+        assert len(dialog.turns) == 6
+
+    def test_invalid_dialog_id(self):
+        """Test that invalid dialog_id raises ValidationError."""
+        from pydantic import ValidationError
+
+        from sage.data.sources.agent_sft.schemas import AgentSFTDialog
+
+        with pytest.raises(ValidationError):
+            AgentSFTDialog(
+                dialog_id="invalid_id",
+                goal="Test",
+                turns=[{"role": "user", "content": "Test"}] * 6,
+                target_tools=["test_tool_001"]
+            )
+
+    def test_invalid_turn_count(self):
+        """Test that dialogs with <6 or >12 turns fail validation."""
+        from pydantic import ValidationError
+
+        from sage.data.sources.agent_sft.schemas import AgentSFTDialog
+
+        # Too few turns
+        with pytest.raises(ValidationError, match="6-12 turns"):
+            AgentSFTDialog(
+                dialog_id="sft_000001",
+                goal="Test",
+                turns=[{"role": "user", "content": "Test"}] * 3,
+                target_tools=["test_tool_001"]
+            )
+
+        # Too many turns
+        with pytest.raises(ValidationError, match="6-12 turns"):
+            AgentSFTDialog(
+                dialog_id="sft_000001",
+                goal="Test",
+                turns=[{"role": "user", "content": "Test"}] * 15,
+                target_tools=["test_tool_001"]
+            )
+
+    def test_empty_target_tools(self):
+        """Test that empty target_tools fails validation."""
+        from pydantic import ValidationError
+
+        from sage.data.sources.agent_sft.schemas import AgentSFTDialog
+
+        with pytest.raises(ValidationError, match="cannot be empty"):
+            AgentSFTDialog(
+                dialog_id="sft_000001",
+                goal="Test",
+                turns=[{"role": "user", "content": "Test"}] * 6,
+                target_tools=[]
+            )
+
+
+if __name__ == "__main__":
+    # Run tests with pytest
+    pytest.main([__file__, "-v"])

sage/data/sources/agent_tools/__init__.py

@@ -0,0 +1,6 @@
+"""Agent Tools data source for SAGE benchmarks."""
+
+from .dataloader import AgentToolsDataLoader
+from .schemas import AgentToolRecord, CategoryTaxonomy, DatasetStats
+
+__all__ = ["AgentToolsDataLoader", "AgentToolRecord", "CategoryTaxonomy", "DatasetStats"]

sage/data/sources/agent_tools/dataset.yaml

@@ -0,0 +1,9 @@
+name: "agent_tools"
+description: "1200 curated agent tools with categories and metadata for agent benchmarking"
+type: "tools"
+format: "jsonl"
+version: "1.0.0"
+maintainer: "SAGE Agent Benchmark Team"
+tags: ["tools", "catalog", "agent", "benchmark"]
+license: "CC-BY-4.0"
+size: "~15MB"

sage/data/sources/agent_tools/test_integration.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""
+Test DataManager integration with agent_tools data source.
+
+This script verifies that the agent_tools data source can be accessed
+through SAGE's DataManager.
+"""
+
+import sys
+from pathlib import Path
+
+# Add sage-benchmark to path for testing
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
+
+try:
+    from sage.data.manager import DataManager
+    from sage.data.sources.agent_tools import AgentToolsDataLoader
+
+    print("✅ Successfully imported DataManager and AgentToolsDataLoader\n")
+
+    # Test 1: Direct import
+    print("Test 1: Direct loader instantiation")
+    print("=" * 60)
+    loader = AgentToolsDataLoader()
+    print(f"✓ Loaded {len(loader)} tools across {len(loader.get_categories())} categories\n")
+
+    # Test 2: DataManager access by source
+    print("Test 2: Access via DataManager.get_source()")
+    print("=" * 60)
+    dm = DataManager()
+
+    # Check if agent_tools is discoverable
+    sources = dm.source_registry.discover_sources()
+    print(f"Available sources: {len(sources)}")
+
+    if "agent_tools" in sources:
+        print("✓ agent_tools found in available sources")
+
+        # Load via DataManager
+        agent_tools_module = dm.source_registry.load_source("agent_tools")
+        print(f"✓ Loaded agent_tools module: {agent_tools_module}")
+
+        # Get metadata
+        metadata = dm.source_registry.get_metadata("agent_tools")
+        if metadata:
+            print(f"✓ Metadata: {metadata.name} v{metadata.version}")
+            print(f"  Description: {metadata.description}")
+            print(f"  Type: {metadata.type}")
+            print(f"  Tags: {', '.join(metadata.tags)}")
+
+        # Instantiate loader through module
+        loader_from_dm = agent_tools_module.AgentToolsDataLoader()
+        print(f"✓ Loader from DataManager: {len(loader_from_dm)} tools\n")
+    else:
+        print(f"⚠ agent_tools not in discovered sources: {sources[:10]}")
+        print("  (This is expected if DataManager hasn't refreshed yet)\n")
+
+    # Test 3: Search and retrieval
+    print("Test 3: Search and retrieval operations")
+    print("=" * 60)
+
+    # Search by capability
+    weather_tools = loader.search_by_capability("weather", top_k=3)
+    print(f"✓ Found {len(weather_tools)} tools with 'weather' capability:")
+    for tool in weather_tools:
+        print(f"  - {tool.name} ({tool.tool_id})")
+
+    # Get specific tool
+    first_tool_id = loader.list_tool_ids()[0]
+    tool = loader.get_tool(first_tool_id)
+    print(f"\n✓ Retrieved tool: {tool.name}")
+    print(f"  Category: {tool.category}")
+    print(f"  Capabilities: {', '.join(tool.capabilities[:3])}")
+
+    # Category iteration
+    categories = loader.get_categories()
+    category = categories[0]
+    tools_in_cat = list(loader.iter_category(category))
+    print(f"\n✓ Category '{category}' has {len(tools_in_cat)} tools\n")
+
+    # Test 4: Statistics and taxonomy
+    print("Test 4: Load statistics and taxonomy")
+    print("=" * 60)
+
+    stats = loader.load_stats()
+    print("✓ Dataset Statistics:")
+    print(f"  Total tools: {stats.total_tools}")
+    print(f"  Total categories: {stats.total_categories}")
+    print(f"  Last updated: {stats.last_updated}")
+    print(f"  Version: {stats.version}")
+
+    taxonomy = loader.load_taxonomy()
+    print("\n✓ Category Taxonomy:")
+    print(f"  Categories defined: {len(taxonomy.taxonomy)}")
+    print(f"  Version: {taxonomy.version}")
+    print("  Sample categories:")
+    for cat_def in taxonomy.taxonomy[:5]:
+        print(f"    - {cat_def.path}: {cat_def.description}")
+
+    print("\n" + "=" * 60)
+    print("✅ All integration tests passed!")
+    print("=" * 60)
+
+except Exception as e:
+    print(f"\n❌ Integration test failed: {e}")
+    import traceback
+    traceback.print_exc()
+    sys.exit(1)