sandboxy 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

sandboxy/mlflow/tracing.py

@@ -0,0 +1,126 @@
+"""MLflow tracing support for Sandboxy.
+
+Enables automatic tracing of LLM calls using MLflow's autolog feature.
+When enabled, all OpenAI SDK calls are automatically captured as spans
+within the MLflow run, providing detailed visibility into:
+- Each LLM call (prompt, response, latency, tokens)
+- Tool/function calls made by the LLM
+- The full execution flow
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+_tracing_enabled = False
+
+
+def enable_tracing(
+    tracking_uri: str | None = None,
+    experiment_name: str | None = None,
+) -> bool:
+    """Enable MLflow tracing for OpenAI calls.
+
+    This should be called once before any LLM calls are made.
+    It enables MLflow's autolog feature which automatically
+    captures all OpenAI SDK calls as traces.
+
+    Args:
+        tracking_uri: MLflow tracking server URI (uses env var if not set)
+        experiment_name: Experiment to log traces to
+
+    Returns:
+        True if tracing was enabled successfully, False otherwise
+    """
+    global _tracing_enabled
+
+    if _tracing_enabled:
+        return True
+
+    try:
+        import os
+
+        import mlflow
+
+        # Set tracking URI before enabling autolog
+        uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI")
+        if uri:
+            mlflow.set_tracking_uri(uri)
+
+        # Set experiment before enabling autolog
+        if experiment_name:
+            mlflow.set_experiment(experiment_name)
+
+        # Enable OpenAI autologging - this captures all OpenAI calls as traces
+        mlflow.openai.autolog()
+
+        _tracing_enabled = True
+        logger.debug("MLflow tracing enabled for OpenAI")
+        return True
+
+    except ImportError as e:
+        logger.warning(f"MLflow or OpenAI not installed, tracing disabled: {e}")
+        return False
+    except Exception as e:
+        logger.warning(f"Failed to enable MLflow tracing: {e}")
+        return False
+
+
+def disable_tracing() -> None:
+    """Disable MLflow tracing."""
+    global _tracing_enabled
+
+    if not _tracing_enabled:
+        return
+
+    try:
+        import mlflow
+
+        mlflow.openai.autolog(disable=True)
+        _tracing_enabled = False
+        logger.debug("MLflow tracing disabled")
+
+    except Exception as e:
+        logger.warning(f"Failed to disable MLflow tracing: {e}")
+
+
+@contextmanager
+def trace_span(name: str, span_type: str = "CHAIN") -> Generator[None, None, None]:
+    """Create a manual trace span for non-LLM operations.
+
+    Use this to wrap tool calls, scenario steps, or other operations
+    you want to appear in the trace.
+
+    Args:
+        name: Name of the span (e.g., "tool_call:get_account_activity")
+        span_type: Type of span (CHAIN, TOOL, RETRIEVER, etc.)
+
+    Example:
+        with trace_span("tool_call:search", span_type="TOOL"):
+            result = execute_tool(...)
+    """
+    try:
+        import mlflow
+
+        with mlflow.start_span(name=name, span_type=span_type):
+            yield
+
+    except ImportError:
+        # MLflow not installed, just run without tracing
+        yield
+    except Exception as e:
+        logger.debug(f"Tracing span failed: {e}")
+        yield
+
+
+def is_tracing_enabled() -> bool:
+    """Check if tracing is currently enabled."""
+    return _tracing_enabled

sandboxy/scenarios/loader.py

@@ -28,6 +28,25 @@ class StepSpec(BaseModel):
     params: dict[str, Any] = Field(default_factory=dict)
 
 
+class MLflowYamlConfig(BaseModel):
+    """MLflow configuration from scenario YAML.
+
+    Example:
+        mlflow:
+          enabled: true
+          experiment: "agent-evals"
+          tracking_uri: "http://localhost:5000"
+          tags:
+            team: "support-agents"
+            environment: "staging"
+    """
+
+    enabled: bool = False
+    experiment: str | None = None
+    tracking_uri: str | None = None
+    tags: dict[str, str] = Field(default_factory=dict)
+
+
 class McpServerSpec(BaseModel):
     """Specification for an MCP server connection.
 
@@ -74,9 +93,12 @@ class ScenarioSpec(BaseModel):
 
     # Evaluation
     goals: list[GoalSpec] = Field(default_factory=list)
-    evaluation: list[dict[str, Any]] = Field(default_factory=list)
+    evaluation: list[dict[str, Any]] | dict[str, Any] = Field(default_factory=list)
     scoring: dict[str, Any] = Field(default_factory=dict)
 
+    # MLflow integration (optional)
+    mlflow: MLflowYamlConfig | None = None
+
 
 def load_scenario(path: Path) -> ScenarioSpec:
     """Load a scenario from a YAML file.
@@ -148,7 +170,14 @@ def parse_scenario(raw: dict[str, Any]) -> ScenarioSpec:
 
     # Parse goals
     goals: list[GoalSpec] = []
-    for g in raw.get("goals", []):
+    goals_raw = raw.get("goals", [])
+
+    # Also check for goals nested inside evaluation dict
+    evaluation_raw = raw.get("evaluation", [])
+    if isinstance(evaluation_raw, dict) and "goals" in evaluation_raw:
+        goals_raw = evaluation_raw.get("goals", [])
+
+    for g in goals_raw:
         goals.append(
             GoalSpec(
                 id=g.get("id", f"goal_{len(goals)}"),
@@ -159,6 +188,17 @@ def parse_scenario(raw: dict[str, Any]) -> ScenarioSpec:
             )
         )
 
+    # Parse MLflow config if present
+    mlflow_config: MLflowYamlConfig | None = None
+    if "mlflow" in raw and isinstance(raw["mlflow"], dict):
+        mlflow_raw = raw["mlflow"]
+        mlflow_config = MLflowYamlConfig(
+            enabled=mlflow_raw.get("enabled", False),
+            experiment=mlflow_raw.get("experiment"),
+            tracking_uri=mlflow_raw.get("tracking_uri"),
+            tags=mlflow_raw.get("tags", {}),
+        )
+
     return ScenarioSpec(
         id=raw.get("id", "unnamed"),
         name=raw.get("name", raw.get("id", "Unnamed Scenario")),
@@ -174,6 +214,7 @@ def parse_scenario(raw: dict[str, Any]) -> ScenarioSpec:
         goals=goals,
         evaluation=raw.get("evaluation", []),
         scoring=raw.get("scoring", {}),
+        mlflow=mlflow_config,
     )
 
 
@@ -259,4 +300,5 @@ def apply_scenario_variables(spec: ScenarioSpec, variables: dict[str, Any]) -> S
         goals=spec.goals,
         evaluation=spec.evaluation,
         scoring=spec.scoring,
+        mlflow=spec.mlflow,  # Preserve MLflow config
     )

sandboxy/scenarios/runner.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
+import time
 from pathlib import Path
 from typing import Any
 
@@ -26,6 +27,16 @@ class ScenarioEvent(BaseModel):
     payload: dict[str, Any] = Field(default_factory=dict)
 
 
+class GoalResult(BaseModel):
+    """Result of evaluating a single goal."""
+
+    id: str
+    name: str
+    achieved: bool
+    points: int
+    reason: str = ""
+
+
 class ScenarioResult(BaseModel):
     """Result of running a scenario."""
 
@@ -35,7 +46,12 @@ class ScenarioResult(BaseModel):
     tool_calls: list[dict[str, Any]] = Field(default_factory=list)
     final_state: dict[str, Any] = Field(default_factory=dict)
     goals_achieved: list[str] = Field(default_factory=list)
+    goal_results: list[GoalResult] = Field(default_factory=list)
     score: float = 0.0
+    max_score: float = 0.0
+    latency_ms: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
 
     def to_json(self, indent: int | None = None) -> str:
         """Serialize result to JSON string."""
@@ -70,7 +86,11 @@ class ScenarioResult(BaseModel):
         lines.append("")
         lines.append(f"Tool Calls Made: {len(self.tool_calls)}")
         lines.append(f"Goals Achieved: {len(self.goals_achieved)}")
-        lines.append(f"Score: {self.score}")
+        lines.append(f"Score: {self.score}/{self.max_score}")
+        if self.latency_ms:
+            lines.append(f"Latency: {self.latency_ms}ms")
+        if self.input_tokens or self.output_tokens:
+            lines.append(f"Tokens: {self.input_tokens} in / {self.output_tokens} out")
 
         return "\n".join(lines)
 
@@ -176,6 +196,8 @@ class ScenarioRunner:
         Returns:
             ScenarioResult with events and evaluation
         """
+        start_time = time.perf_counter()
+
         try:
             # Load MCP tools if configured
             await self._load_mcp_tools()
@@ -188,9 +210,21 @@ class ScenarioRunner:
             for step in self.scenario.steps:
                 await self._execute_step(step, max_turns)
 
-            # Evaluate goals
+            # Evaluate goals and build detailed results
             goals_achieved = self._evaluate_goals()
+            goal_results = self._build_goal_results(goals_achieved)
             score = self._compute_score(goals_achieved)
+            max_score = sum(g.points for g in self.scenario.goals)
+
+            # Get token usage from agent if available
+            input_tokens = 0
+            output_tokens = 0
+            if hasattr(self.agent, "get_usage"):
+                usage = self.agent.get_usage()
+                input_tokens = usage.get("input_tokens", 0)
+                output_tokens = usage.get("output_tokens", 0)
+
+            latency_ms = int((time.perf_counter() - start_time) * 1000)
 
             return ScenarioResult(
                 scenario_id=self.scenario.id,
@@ -199,7 +233,12 @@ class ScenarioRunner:
                 tool_calls=self.tool_call_log,
                 final_state=self.env_state.copy(),
                 goals_achieved=goals_achieved,
+                goal_results=goal_results,
                 score=score,
+                max_score=max_score,
+                latency_ms=latency_ms,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
             )
         finally:
             await self._cleanup_mcp()
@@ -440,6 +479,22 @@ class ScenarioRunner:
 
         return list(set(achieved))  # Deduplicate
 
+    def _build_goal_results(self, goals_achieved: list[str]) -> list[GoalResult]:
+        """Build detailed goal results for MLflow logging."""
+        results = []
+        for goal in self.scenario.goals:
+            achieved = goal.id in goals_achieved
+            results.append(
+                GoalResult(
+                    id=goal.id,
+                    name=goal.name or goal.id,
+                    achieved=achieved,
+                    points=goal.points if achieved else 0,
+                    reason="Goal achieved" if achieved else "Goal not achieved",
+                )
+            )
+        return results
+
     def _compute_score(self, goals_achieved: list[str]) -> float:
         """Compute score based on achieved goals."""
         from sandboxy.core.safe_eval import EvaluationError, safe_eval_formula

sandboxy/tools/yaml_tools.py

@@ -288,6 +288,17 @@ class YamlMockTool:
         self.config = config.config
         self.spec = spec
         self._call_log: list[dict[str, Any]] = []
+        self._overrides: dict[str, Any] = {}
+
+    def set_overrides(self, overrides: dict[str, Any]) -> None:
+        """Set response overrides for dataset benchmarking.
+
+        Args:
+            overrides: Dict mapping "tool.action" or "tool" to override response data.
+                      When a matching action is called, returns the override data
+                      instead of the normal mock response.
+        """
+        self._overrides = overrides
 
     @property
     def call_log(self) -> list[dict[str, Any]]:
@@ -347,6 +358,13 @@ class YamlMockTool:
         for effect in action_spec.side_effects:
             effect.apply(env_state, validated_args)
 
+        # Check for override (dataset benchmarking)
+        override_key = f"{self.name}.{action}"
+        if override_key in self._overrides:
+            return ToolResult(success=True, data=self._overrides[override_key])
+        if self.name in self._overrides:
+            return ToolResult(success=True, data=self._overrides[self.name])
+
         # Compute return value
         result_value = self._compute_return(action_spec, validated_args, env_state)