isage-benchmark-agent 0.1.0.1__cp311-none-any.whl

sage/benchmark/benchmark_agent/experiments/tool_selection_exp.py

@@ -0,0 +1,250 @@
+"""
+Tool Selection Experiment
+
+Experiment runner for evaluating tool selection capabilities.
+Tests ability to select relevant tools from a candidate set given a user instruction.
+"""
+
+from typing import Any
+
+from sage.benchmark.benchmark_agent.experiments.base_experiment import (
+    BaseExperiment,
+    ExperimentResult,
+    ToolSelectionConfig,
+)
+
+
+class ToolSelectionQuery:
+    """Query for tool selection."""
+
+    def __init__(
+        self, sample_id: str, instruction: str, context: dict[str, Any], candidate_tools: list[str]
+    ):
+        self.sample_id = sample_id
+        self.instruction = instruction
+        self.context = context
+        self.candidate_tools = candidate_tools
+
+
+class ToolSelectionExperiment(BaseExperiment):
+    """
+    Experiment for tool selection evaluation.
+
+    Workflow:
+    1. Load benchmark samples from DataManager
+    2. For each sample, call selector strategy to get predictions
+    3. Collect predictions and ground truth
+    4. Return ExperimentResult for evaluation
+    """
+
+    def __init__(
+        self, config: ToolSelectionConfig, data_manager: Any = None, adapter_registry: Any = None
+    ):
+        """
+        Initialize tool selection experiment.
+
+        Args:
+            config: Tool selection configuration
+            data_manager: DataManager for data loading
+            adapter_registry: Registry containing selector strategies
+        """
+        super().__init__(config, data_manager, adapter_registry)
+        self.config: ToolSelectionConfig = config
+        self._embedding_client = None
+
+    def _create_embedding_client(self):
+        """Create embedding client for selector."""
+        # Try to create an embedding client wrapper
+        try:
+            import os
+
+            from sage.common.components.sage_embedding import EmbeddingService
+
+            # Determine embedding method and optional model from environment
+            method = os.environ.get("SAGE_EMBEDDING_METHOD") or "hf"
+            model = os.environ.get("SAGE_EMBEDDING_MODEL") or None
+            api_key = os.environ.get("SAGE_EMBEDDING_API_KEY") or None
+            base_url = os.environ.get("SAGE_EMBEDDING_BASE_URL") or None
+
+            # Create a simple embedding client wrapper that uses configured service
+            class EmbeddingClientWrapper:
+                """Wrapper to adapt EmbeddingService to selector interface."""
+
+                def __init__(self):
+                    cfg = {
+                        "method": method,
+                        "model": model,
+                        "api_key": api_key,
+                        "base_url": base_url,
+                        "normalize": True,
+                    }
+                    try:
+                        self.service = EmbeddingService(cfg)
+                        self.service.setup()
+                    except Exception:
+                        # Fallback to mockembedder if configured method not available
+                        cfg["method"] = "mockembedder"
+                        self.service = EmbeddingService(cfg)
+                        self.service.setup()
+
+                def embed(self, texts, model=None, batch_size=32):
+                    """Embed texts and return numpy array."""
+                    import numpy as np
+
+                    result = self.service.embed(texts, batch_size=batch_size)
+                    return np.array(result["vectors"])
+
+                def cleanup(self):
+                    try:
+                        self.service.cleanup()
+                    except Exception:
+                        pass
+
+            return EmbeddingClientWrapper()
+        except ImportError:
+            return None
+        except Exception as e:
+            print(f"Warning: Could not create embedding client: {e}")
+            return None
+
+    def prepare(self):
+        """Prepare experiment: load data and initialize selector."""
+        super().prepare()
+
+        verbose = getattr(self.config, "verbose", False)
+        if verbose:
+            print(f"\n{'=' * 60}")
+            print(f"Tool Selection Experiment: {self.experiment_id}")
+            print(f"{'=' * 60}")
+            print(f"Profile: {self.config.profile}")
+            print(f"Split: {self.config.split}")
+            print(f"Selector: {self.config.selector}")
+            print(f"Top-k: {self.config.top_k}")
+
+        # Load data through DataManager
+        try:
+            agent_eval = self.dm.get_by_usage("agent_eval")
+            profile_data = agent_eval.load_profile(self.config.profile)
+
+            self.benchmark_loader = profile_data.get("benchmark")
+            self.tools_loader = profile_data.get("tools")
+
+            if verbose:
+                print("✓ Loaded benchmark data")
+
+        except Exception as e:
+            print(f"Warning: Could not load data: {e}")
+
+        # Initialize selector strategy with real tools data
+        if self.adapter_registry is not None:
+            try:
+                # Create resources with real tools loader
+                from sage.libs.agentic.agents.action.tool_selection import (
+                    SelectorResources,
+                )
+
+                # Create embedding client if selector needs it
+                embedding_client = None
+                selector_name = self.config.selector.lower()
+                if "embedding" in selector_name or "hybrid" in selector_name:
+                    embedding_client = self._create_embedding_client()
+                    self._embedding_client = embedding_client
+                    if verbose and embedding_client:
+                        print("✓ Initialized embedding client")
+
+                resources = SelectorResources(
+                    tools_loader=self.tools_loader,
+                    embedding_client=embedding_client,
+                )
+
+                self.strategy = self.adapter_registry.get(self.config.selector, resources=resources)
+                if verbose:
+                    print(f"✓ Initialized selector: {self.config.selector}")
+            except Exception as e:
+                print(f"Warning: Could not load selector: {e}")
+                self.strategy = None
+        else:
+            self.strategy = None
+
+    def run(self) -> ExperimentResult:
+        """
+        Run tool selection experiment.
+
+        Returns:
+            ExperimentResult with predictions and references
+        """
+        verbose = getattr(self.config, "verbose", False)
+        if verbose:
+            print("\nRunning experiment...")
+
+        predictions = []
+        references = []
+        metadata = {"total_samples": 0, "failed_samples": 0}
+
+        try:
+            samples = self.benchmark_loader.iter_split(
+                task_type="tool_selection", split=self.config.split
+            )
+
+            for idx, sample in enumerate(samples):
+                if self.config.max_samples and idx >= self.config.max_samples:
+                    break
+
+                metadata["total_samples"] += 1
+
+                try:
+                    # Handle context - may be string or dict
+                    context = sample.context if hasattr(sample, "context") else {}
+                    if isinstance(context, str):
+                        context = {"description": context}
+                    elif context is None:
+                        context = {}
+
+                    query = ToolSelectionQuery(
+                        sample_id=sample.sample_id,
+                        instruction=sample.instruction,
+                        context=context,
+                        candidate_tools=sample.candidate_tools,
+                    )
+
+                    if self.strategy is not None:
+                        pred_tools = self.strategy.predict(query, top_k=self.config.top_k)
+                        pred_dict = {
+                            "sample_id": sample.sample_id,
+                            "predicted_tools": [
+                                {"tool_id": p.tool_id, "score": p.score} for p in pred_tools
+                            ],
+                        }
+                    else:
+                        pred_dict = {"sample_id": sample.sample_id, "predicted_tools": []}
+
+                    predictions.append(pred_dict)
+
+                    gt = sample.get_typed_ground_truth()
+                    ref_dict = {
+                        "sample_id": sample.sample_id,
+                        "ground_truth_tools": gt.top_k,
+                        "explanation": gt.explanation if hasattr(gt, "explanation") else None,
+                    }
+                    references.append(ref_dict)
+
+                    if verbose and (idx + 1) % 10 == 0:
+                        print(f"  Processed {idx + 1} samples...")
+
+                except Exception as e:
+                    metadata["failed_samples"] += 1
+                    if verbose:
+                        print(f"  Error processing sample {idx}: {e}")
+                    continue
+
+        except Exception as e:
+            print(f"Error iterating samples: {e}")
+
+        if verbose:
+            print("\nCompleted:")
+            print(f"  Total samples: {metadata['total_samples']}")
+            print(f"  Failed samples: {metadata['failed_samples']}")
+
+        return self._create_result(
+            predictions=predictions, references=references, metadata=metadata
+        )

sage/benchmark/benchmark_agent/scripts/__init__.py

@@ -0,0 +1,26 @@
+"""
+SAGE-Bench Scripts
+
+统一的 Benchmark 实验脚本入口。
+
+所有实验功能位于 experiments/ 子包:
+- run_paper1_experiments.py: 论文 1 实验统一入口
+- exp_main_*.py: Section 5.2 主实验
+- exp_analysis_*.py: Section 5.3 分析实验
+- exp_cross_dataset.py: Section 5.4 跨数据集泛化
+- exp_training_comparison.py: Section 5.5 训练方法对比
+
+Usage:
+    # CLI 入口
+    sage-bench run --quick
+    sage-bench eval --dataset all
+    sage-bench train --dry-run
+    sage-bench llm status
+
+    # 直接运行
+    python -m sage.benchmark.benchmark_agent.scripts.experiments.run_paper1_experiments --quick
+"""
+
+__all__ = [
+    "experiments",
+]

sage/benchmark/benchmark_agent/scripts/experiments/__init__.py

@@ -0,0 +1,40 @@
+"""
+SAGE-Bench Paper 1 Experiments Package
+
+按论文 Experiment Section 组织的实验脚本集合:
+
+- Section 5.2 (Main Results):
+    - exp_main_timing.py      # RQ1: Timing Detection
+    - exp_main_planning.py    # RQ2: Task Planning
+    - exp_main_selection.py   # RQ3: Tool Selection
+
+- Section 5.3 (Analysis):
+    - exp_analysis_error.py       # 5.3.1 Error Analysis
+    - exp_analysis_scaling.py     # 5.3.2 Scaling Analysis
+    - exp_analysis_robustness.py  # 5.3.3 Robustness Analysis
+    - exp_analysis_ablation.py    # 5.3.4 Ablation Studies
+
+- Section 5.4 (Generalization):
+    - exp_cross_dataset.py    # Cross-dataset evaluation
+
+Usage:
+    sage-bench paper1 run                    # 运行所有实验
+    sage-bench paper1 run --section 5.2      # 仅主实验
+    sage-bench paper1 timing                 # 单个实验
+"""
+
+from .exp_utils import (
+    get_embedding_client,
+    get_llm_client,
+    load_benchmark_data,
+    save_results,
+    setup_experiment_env,
+)
+
+__all__ = [
+    "setup_experiment_env",
+    "load_benchmark_data",
+    "save_results",
+    "get_llm_client",
+    "get_embedding_client",
+]