synth-ai 0.1.9__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/examples/verilog/engine.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Dict, Any, Tuple, Optional
+from dataclasses import dataclass
+
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.environments.tasks.core import TaskInstance
+from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent
+
+
+@dataclass
+class VerilogEngineSnapshot(StatefulEngineSnapshot):
+    task_instance_dict: Dict
+    engine_snapshot: Dict
+
+    def model_dump(self) -> Dict:
+        """Convert dataclass to dictionary for compatibility with Pydantic models."""
+        return {
+            "task_instance_dict": self.task_instance_dict,
+            "engine_snapshot": self.engine_snapshot,
+        }
+
+
+@dataclass
+class VerilogPublicState:
+    files: Dict[str, str]
+    build_dir: str
+    task_completed: bool = False
+    last_compile_output: Optional[str] = None
+    last_simulate_output: Optional[str] = None
+
+
+@dataclass
+class VerilogPrivateState:
+    reward_last: float
+    total_reward: float
+    terminated: bool
+    truncated: bool
+
+
+class VerilogCompileSuccessComponent(RewardComponent):
+    async def score(self, state: VerilogPublicState, action: Any) -> float:
+        if hasattr(action, "get") and action.get("type") == "compile":
+            # Check if compilation was successful (returncode 0)
+            if action.get("returncode") == 0:
+                return 0.1
+        return 0.0
+
+
+class VerilogSimulationPassComponent(RewardComponent):
+    async def score(self, state: VerilogPublicState, action: Any) -> float:
+        if hasattr(action, "get") and action.get("type") == "simulate":
+            # Check if simulation passed
+            if action.get("passed", False):
+                return 1.0
+        return 0.0
+
+
+class VerilogStepPenaltyComponent(RewardComponent):
+    def __init__(self, penalty: float = -0.01):
+        self.penalty = penalty
+
+    async def score(self, state: Any, action: Any) -> float:
+        return self.penalty
+
+
+class VerilogEngine(StatefulEngine):
+    """
+    Stateful Verilog evaluation engine with persistent artifact snapshots.
+    """
+
+    def __init__(self, task_instance: TaskInstance):
+        self.task_instance = task_instance
+        self._total_reward = 0.0
+        self._current_action_for_reward: Optional[Dict[str, Any]] = None
+
+        self.reward_stack = RewardStack(
+            components=[
+                VerilogCompileSuccessComponent(),
+                VerilogSimulationPassComponent(),
+                VerilogStepPenaltyComponent(penalty=-0.01),
+            ]
+        )
+
+        # Initialize paths - will be set properly in _reset_engine
+        self.snapshot_dir: Optional[Path] = None
+        self.build_dir: Optional[Path] = None
+
+        # Track last compile/simulate outputs
+        self._last_compile_output: Optional[str] = None
+        self._last_simulate_output: Optional[str] = None
+
+    async def _reset_engine(
+        self, *, seed: Optional[int] = None
+    ) -> Tuple[VerilogPrivateState, VerilogPublicState]:
+        """Initialize the Verilog environment with task files."""
+        self._total_reward = 0.0
+        self._current_action_for_reward = None
+        self._last_compile_output = None
+        self._last_simulate_output = None
+
+        # Initialize snapshot from task instance
+        self._init_snapshot()
+
+        priv = VerilogPrivateState(
+            reward_last=0.0, total_reward=0.0, terminated=False, truncated=False
+        )
+
+        pub = VerilogPublicState(
+            files=self._get_file_contents(),
+            build_dir=str(self.build_dir),
+            task_completed=False,
+        )
+
+        return priv, pub
+
+    async def _step_engine(
+        self, action_result: Dict[str, Any]
+    ) -> Tuple[VerilogPrivateState, VerilogPublicState]:
+        """Process an action result and update engine state."""
+        self._current_action_for_reward = action_result
+
+        # Update last outputs if this is a compile or simulate action
+        if action_result.get("type") == "compile":
+            stdout = action_result.get("stdout", "")
+            stderr = action_result.get("stderr", "")
+            # Combine stdout and stderr for compile output, stderr has the error info
+            self._last_compile_output = stderr if stderr else stdout
+        elif action_result.get("type") == "simulate":
+            self._last_simulate_output = action_result.get("stdout")
+
+        # Calculate reward using RewardStack
+        current_pub_state = VerilogPublicState(
+            files=self._get_file_contents(),
+            build_dir=str(self.build_dir),
+            task_completed=action_result.get("passed", False),
+        )
+
+        reward_from_stack = await self.reward_stack.step_reward(
+            state=current_pub_state, action=self._current_action_for_reward
+        )
+        self._current_action_for_reward = None
+
+        self._total_reward += reward_from_stack
+
+        # Check termination conditions
+        terminated = action_result.get("passed", False) or action_result.get("submitted", False)
+
+        priv = VerilogPrivateState(
+            reward_last=reward_from_stack,
+            total_reward=self._total_reward,
+            terminated=terminated,
+            truncated=False,
+        )
+
+        pub = VerilogPublicState(
+            files=self._get_file_contents(),
+            build_dir=str(self.build_dir),
+            task_completed=action_result.get("passed", False),
+            last_compile_output=self._last_compile_output,
+            last_simulate_output=self._last_simulate_output,
+        )
+
+        return priv, pub
+
+    def _init_snapshot(self) -> None:
+        """Initialize snapshot directory from task instance data."""
+        if not hasattr(self.task_instance, "snapshot_dir"):
+            raise ValueError("Task instance must have a snapshot_dir attribute")
+
+        self.snapshot_dir = Path(self.task_instance.snapshot_dir)
+
+        if self.snapshot_dir.exists() and any(self.snapshot_dir.iterdir()):
+            # Already initialized
+            self.build_dir = self.snapshot_dir / "build"
+            self.build_dir.mkdir(exist_ok=True)
+            return
+
+        # Copy pristine files from task data
+        pristine_dir = getattr(self.task_instance, "pristine_dir", None)
+        if pristine_dir and Path(pristine_dir).exists():
+            shutil.copytree(pristine_dir, self.snapshot_dir, dirs_exist_ok=True)
+        else:
+            # Create basic structure if no pristine dir
+            self.snapshot_dir.mkdir(parents=True, exist_ok=True)
+
+        self.build_dir = self.snapshot_dir / "build"
+        self.build_dir.mkdir(exist_ok=True)
+
+    def _get_file_contents(self) -> Dict[str, str]:
+        """Get contents of all Verilog files in the snapshot directory."""
+        if not self.snapshot_dir:
+            return {}
+
+        files = {}
+        for p in self.snapshot_dir.rglob("*.v"):
+            try:
+                relative_path = p.relative_to(self.snapshot_dir)
+                files[str(relative_path)] = p.read_text()
+            except Exception:
+                continue
+        return files
+
+    async def write_file(self, path: str, content: str) -> Dict[str, Any]:
+        """Write content to a file in the snapshot directory."""
+        if not self.snapshot_dir:
+            return {"ok": False, "error": "Snapshot directory not initialized"}
+
+        file_path = self.snapshot_dir / path
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(content)
+        return {"ok": True, "type": "write_file"}
+
+    async def compile(
+        self, sources: Optional[list] = None, testbench: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Compile Verilog sources with iverilog."""
+        if not self.snapshot_dir or not self.build_dir:
+            return {"ok": False, "error": "Directories not initialized"}
+
+        # Default to all .v files if no sources specified
+        if sources is None:
+            sources = [str(p.relative_to(self.snapshot_dir)) for p in self.snapshot_dir.glob("*.v")]
+
+        src_paths = [self.snapshot_dir / src for src in sources]
+
+        # Add testbench if specified
+        if testbench:
+            tb_path = self.snapshot_dir / testbench
+            if tb_path.exists():
+                src_paths.append(tb_path)
+
+        binary = self.build_dir / "a.out"
+        cmd = ["iverilog", "-g2012", "-o", str(binary)] + [str(p) for p in src_paths]
+
+        try:
+            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
+            return {
+                "ok": proc.returncode == 0,
+                "type": "compile",
+                "stdout": proc.stdout,
+                "stderr": proc.stderr,
+                "returncode": proc.returncode,
+                "binary": str(binary) if proc.returncode == 0 else None,
+            }
+        except subprocess.TimeoutExpired:
+            return {"ok": False, "error": "Compilation timeout", "type": "compile"}
+        except Exception as e:
+            return {"ok": False, "error": str(e), "type": "compile"}
+
+    async def simulate(self, binary: Optional[str] = None) -> Dict[str, Any]:
+        """Run vvp on compiled binary."""
+        if not self.build_dir:
+            return {"ok": False, "error": "Build directory not initialized"}
+
+        bin_path = binary if binary else str(self.build_dir / "a.out")
+
+        try:
+            proc = subprocess.run(["vvp", bin_path], capture_output=True, text=True, timeout=30)
+
+            # Check for various success indicators
+            stdout = proc.stdout
+            passed = (
+                "ALL_TESTS_PASSED" in stdout
+                or ("Mismatches: 0 " in stdout and "samples" in stdout)
+                or ("no mismatches" in stdout.lower() and "errors" not in stdout.lower())
+            )
+
+            return {
+                "ok": True,
+                "type": "simulate",
+                "stdout": proc.stdout,
+                "stderr": proc.stderr,
+                "returncode": proc.returncode,
+                "passed": passed,
+            }
+        except subprocess.TimeoutExpired:
+            return {"ok": False, "error": "Simulation timeout", "type": "simulate"}
+        except Exception as e:
+            return {"ok": False, "error": str(e), "type": "simulate"}
+
+    async def submit(self) -> Dict[str, Any]:
+        """Submit solution for grading."""
+        # For now, simple check based on last simulation
+        # In a full implementation, this would call the task's verify method
+        return {
+            "ok": True,
+            "type": "submit",
+            "passed": True,  # Placeholder
+            "detail": "Submission processed",
+            "submitted": True,
+        }
+
+    async def _serialize_engine(self) -> VerilogEngineSnapshot:
+        """Serialize engine state to a snapshot."""
+        engine_data = {
+            "total_reward": self._total_reward,
+            "snapshot_dir": str(self.snapshot_dir) if self.snapshot_dir else None,
+            "build_dir": str(self.build_dir) if self.build_dir else None,
+        }
+
+        task_instance_dict = await self.task_instance.serialize()
+
+        return VerilogEngineSnapshot(
+            task_instance_dict=task_instance_dict, engine_snapshot=engine_data
+        )
+
+    @classmethod
+    async def _deserialize_engine(cls, snapshot: VerilogEngineSnapshot) -> "VerilogEngine":
+        """Deserialize engine from snapshot."""
+        # This would need proper task instance deserialization
+        # For now, create a minimal implementation
+        from synth_ai.environments.examples.verilog.taskset import VerilogTaskInstance
+
+        task_instance = await VerilogTaskInstance.deserialize(snapshot.task_instance_dict)
+        engine = cls(task_instance)
+
+        engine_data = snapshot.engine_snapshot
+        engine._total_reward = engine_data.get("total_reward", 0.0)
+
+        if engine_data.get("snapshot_dir"):
+            engine.snapshot_dir = Path(engine_data["snapshot_dir"])
+        if engine_data.get("build_dir"):
+            engine.build_dir = Path(engine_data["build_dir"])
+
+        return engine

synth_ai/environments/examples/verilog/environment.py

@@ -0,0 +1,349 @@
+from typing import List, Optional, Any, Dict, Union
+from pydantic import BaseModel
+
+from synth_ai.environments.examples.verilog.engine import (
+    VerilogEngine,
+    VerilogPrivateState,
+    VerilogPublicState,
+    VerilogEngineSnapshot,
+)
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.tasks.core import TaskInstance
+from synth_ai.environments.environment.tools import (
+    AbstractTool,
+    EnvToolCall,
+    ToolResult,
+    TOOL_REGISTRY,
+    register_tool,
+)
+
+
+# Tool Input Schemas
+class WriteFileInput(BaseModel):
+    path: str
+    content: str
+
+
+class CompileInput(BaseModel):
+    sources: Optional[List[str]] = None
+    testbench: Optional[str] = None
+
+
+class SimulateInput(BaseModel):
+    binary: Optional[str] = None
+
+
+class SubmitInput(BaseModel):
+    pass  # No arguments needed for submit
+
+
+# Tool Implementations
+class VerilogWriteFileTool(AbstractTool):
+    name = "write_file"
+    description = "Write content to a Verilog file"
+    call_schema = WriteFileInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: VerilogEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            validated_args = self.call_schema(**call.args)
+            result = await self.engine.write_file(validated_args.path, validated_args.content)
+            return ToolResult(ok=result["ok"], payload=result)
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e))
+
+
+class VerilogCompileTool(AbstractTool):
+    name = "compile"
+    description = "Compile Verilog sources with iverilog"
+    call_schema = CompileInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: VerilogEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            validated_args = self.call_schema(**call.args)
+            result = await self.engine.compile(validated_args.sources, validated_args.testbench)
+            return ToolResult(ok=result["ok"], payload=result)
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e))
+
+
+class VerilogSimulateTool(AbstractTool):
+    name = "simulate"
+    description = "Run vvp on compiled binary"
+    call_schema = SimulateInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: VerilogEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            validated_args = self.call_schema(**call.args)
+            result = await self.engine.simulate(validated_args.binary)
+            return ToolResult(ok=result["ok"], payload=result)
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e))
+
+
+class VerilogSubmitTool(AbstractTool):
+    name = "submit"
+    description = "Submit solution for grading"
+    call_schema = SubmitInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: VerilogEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            result = await self.engine.submit()
+            return ToolResult(ok=result["ok"], payload=result)
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e))
+
+
+class VerilogObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: VerilogPublicState, priv: VerilogPrivateState
+    ) -> InternalObservation:
+        files_summary = f"{len(pub.files)} Verilog files available"
+        if pub.files:
+            files_summary += f": {', '.join(pub.files.keys())}"
+
+        compile_status = ""
+        if pub.last_compile_output is not None:
+            # Check for common error indicators in compile output
+            output_lower = pub.last_compile_output.lower()
+            is_success = not any(
+                indicator in output_lower for indicator in ["error", "failed", "syntax"]
+            )
+            if is_success:
+                compile_status = "Last compile: Success"
+            else:
+                # Include the actual error message to help the agent debug
+                compile_status = f"Last compile: Failed\n{pub.last_compile_output}"
+
+        simulate_status = ""
+        if pub.last_simulate_output:
+            # Use same success detection logic as in engine
+            stdout = pub.last_simulate_output
+            passed = (
+                "ALL_TESTS_PASSED" in stdout
+                or ("Mismatches: 0 " in stdout and "samples" in stdout)
+                or ("no mismatches" in stdout.lower() and "errors" not in stdout.lower())
+            )
+            simulate_status = f"Last simulation: {'Passed' if passed else 'Failed'}"
+
+        observation: Dict[str, Any] = {
+            "files": pub.files,
+            "build_dir": pub.build_dir,
+            "files_summary": files_summary,
+            "task_completed": pub.task_completed,
+            "reward_last": priv.reward_last,
+            "total_reward": priv.total_reward,
+            "terminated": priv.terminated,
+            "compile_status": compile_status,
+            "simulate_status": simulate_status,
+        }
+        return observation  # type: ignore[return-value]
+
+
+class VerilogEnvironment(StatefulEnvironment):
+    def __init__(
+        self,
+        task_instance: TaskInstance,
+        custom_obs: Optional[GetObservationCallable] = None,
+    ):
+        self.name = "VerilogEval"
+        self.task_instance = task_instance
+        self.custom_observation_callable = custom_obs or VerilogObservationCallable()
+        self.engine: VerilogEngine = VerilogEngine(task_instance)
+
+        # Initialize tools
+        self._tools_instances = {
+            "write_file": VerilogWriteFileTool(self.engine),
+            "compile": VerilogCompileTool(self.engine),
+            "simulate": VerilogSimulateTool(self.engine),
+            "submit": VerilogSubmitTool(self.engine),
+        }
+
+        # Register tools
+        for tool_name, tool_instance in self._tools_instances.items():
+            if tool_name not in TOOL_REGISTRY:
+                register_tool(tool_instance)
+
+    async def initialize(self) -> InternalObservation:
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub)
+
+    async def terminate(self) -> InternalObservation:
+        # Get current state and mark as terminated
+        try:
+            # Try to get current state from engine
+            current_files = self.engine._get_file_contents()
+            build_dir = str(self.engine.build_dir) if self.engine.build_dir else ""
+
+            priv = VerilogPrivateState(
+                reward_last=0.0,
+                total_reward=self.engine._total_reward,
+                terminated=True,
+                truncated=False,
+            )
+
+            pub = VerilogPublicState(files=current_files, build_dir=build_dir, task_completed=False)
+        except Exception:
+            # Fallback if engine state is not accessible
+            priv = VerilogPrivateState(
+                reward_last=0.0, total_reward=0.0, terminated=True, truncated=False
+            )
+
+            pub = VerilogPublicState(files={}, build_dir="", task_completed=False)
+
+        obs = await self._to_observation(priv, pub)
+        if isinstance(obs, dict):
+            obs["terminated"] = True
+            obs["message"] = "Environment terminated."
+        return obs
+
+    def validate_tool_calls(
+        self,
+        tool_calls: Union[
+            EnvToolCall,
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+            Dict[str, Any],
+        ],
+    ) -> EnvToolCall:
+        """Normalize and validate tool calls to a single EnvToolCall."""
+        raw_call_data: Dict[str, Any]
+
+        if isinstance(tool_calls, list):
+            if not tool_calls:
+                raise ValueError("Received empty list of tool calls.")
+            first_item = tool_calls[0]
+            if isinstance(first_item, list):
+                if not first_item:
+                    raise ValueError("Received empty inner list of tool calls.")
+                raw_call_data = first_item[0]
+            elif isinstance(first_item, dict):
+                raw_call_data = first_item
+            elif isinstance(first_item, EnvToolCall):
+                return first_item
+            else:
+                raise TypeError(f"Unexpected type in tool_calls list: {type(first_item)}")
+        elif isinstance(tool_calls, dict):
+            raw_call_data = tool_calls
+        elif isinstance(tool_calls, EnvToolCall):
+            return tool_calls
+        else:
+            raise TypeError(f"Unexpected type for tool_calls: {type(tool_calls)}")
+
+        if not isinstance(raw_call_data, dict):
+            raise TypeError(f"Processed call data is not a dict: {type(raw_call_data)}")
+
+        # Convert dict to EnvToolCall instance
+        tool_name = raw_call_data.get("tool")
+        tool_args = raw_call_data.get("args", {})
+
+        valid_tools = {"write_file", "compile", "simulate", "submit"}
+        if tool_name not in valid_tools:
+            raise ValueError(f"Unknown tool: {tool_name}. Expected one of: {valid_tools}")
+
+        return EnvToolCall(tool=tool_name, args=tool_args)
+
+    async def step(
+        self,
+        tool_calls: Union[
+            EnvToolCall,
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+            Dict[str, Any],
+        ],
+    ) -> InternalObservation:
+        agent_call = self.validate_tool_calls(tool_calls)
+
+        # Get the appropriate tool
+        tool_instance = self._tools_instances.get(agent_call.tool)
+        if not tool_instance:
+            tool_instance = TOOL_REGISTRY.get(agent_call.tool)
+            if not tool_instance:
+                raise ValueError(f"Tool '{agent_call.tool}' not found.")
+
+        # Execute the tool
+        tool_result: ToolResult = await tool_instance(agent_call)
+
+        # Update engine state with tool result
+        if tool_result.payload:
+            action_result = tool_result.payload
+        elif not tool_result.ok:
+            action_result = {
+                "ok": False,
+                "error": tool_result.error,
+                "type": agent_call.tool,
+            }
+        else:
+            action_result = {}
+
+        priv_state, pub_state = await self.engine._step_engine(action_result)
+
+        return await self._to_observation(priv_state, pub_state)
+
+    async def checkpoint(self) -> InternalObservation:
+        engine_snapshot: VerilogEngineSnapshot = await self.engine._serialize_engine()
+
+        # Get current state for observation
+        try:
+            current_files = self.engine._get_file_contents()
+            build_dir = str(self.engine.build_dir) if self.engine.build_dir else ""
+
+            priv = VerilogPrivateState(
+                reward_last=0.0,
+                total_reward=self.engine._total_reward,
+                terminated=False,
+                truncated=False,
+            )
+
+            pub = VerilogPublicState(files=current_files, build_dir=build_dir, task_completed=False)
+
+            obs_data = await self._to_observation(priv, pub)
+        except Exception:
+            obs_data = {"message": "Checkpoint created"}
+
+        if isinstance(obs_data, dict):
+            obs_data["engine_snapshot_data"] = engine_snapshot.model_dump()
+
+        return obs_data
+
+    async def _to_observation(
+        self,
+        priv: VerilogPrivateState,
+        pub: VerilogPublicState,
+        extra_obs: Optional[Dict[str, Any]] = None,
+    ) -> InternalObservation:
+        observation = await self.custom_observation_callable.get_observation(pub, priv)
+        if extra_obs and isinstance(observation, dict):
+            observation.update(extra_obs)
+        return observation
+
+    async def _serialize_engine(self) -> VerilogEngineSnapshot:
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snapshot: VerilogEngineSnapshot, task_instance: TaskInstance
+    ) -> "VerilogEnvironment":
+        eng = await VerilogEngine._deserialize_engine(snapshot)
+        env = cls(task_instance)
+        env.engine = eng
+        return env