wafer-core 0.1.21__tar.gz → 0.1.23__tar.gz

{wafer_core-0.1.21 → wafer_core-0.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-core
-Version: 0.1.21
+Version: 0.1.23
 Summary: Core utilities and environments for Wafer GPU kernel optimization
 Requires-Python: >=3.10
 Requires-Dist: aiohttp>=3.9.0

{wafer_core-0.1.21 → wafer_core-0.1.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-core"
-version = "0.1.21"
+version = "0.1.23"
 description = "Core utilities and environments for Wafer GPU kernel optimization"
 requires-python = ">=3.10"
 dependencies = [

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/auth.py

@@ -41,6 +41,18 @@ PROVIDERS = {
         "display_name": "Modal",
         "key_url": "https://modal.com/settings",
     },
+    "anthropic": {
+        "env_var": "ANTHROPIC_API_KEY",
+        "alt_env_var": "WAFER_ANTHROPIC_API_KEY",  # Check this first
+        "display_name": "Anthropic",
+        "key_url": "https://console.anthropic.com/settings/keys",
+    },
+    "openai": {
+        "env_var": "OPENAI_API_KEY",
+        "alt_env_var": "WAFER_OPENAI_KEY",  # Check this first
+        "display_name": "OpenAI",
+        "key_url": "https://platform.openai.com/api-keys",
+    },
 }
 
 
@@ -78,11 +90,12 @@ def get_api_key(provider: str) -> str | None:
     """Get API key for a provider.
 
     Checks in order:
-    1. Environment variable (e.g., WAFER_RUNPOD_API_KEY)
-    2. ~/.wafer/auth.json
+    1. Alt environment variable if defined (e.g., WAFER_ANTHROPIC_API_KEY)
+    2. Primary environment variable (e.g., ANTHROPIC_API_KEY)
+    3. ~/.wafer/auth.json
 
     Args:
-        provider: Provider name (runpod, digitalocean, modal)
+        provider: Provider name (runpod, digitalocean, modal, anthropic, openai)
 
     Returns:
         API key string or None if not found
@@ -92,7 +105,13 @@ def get_api_key(provider: str) -> str | None:
 
     config = PROVIDERS[provider]
 
-    # Check environment variable first
+    # Check alt environment variable first (e.g., WAFER_ANTHROPIC_API_KEY)
+    if "alt_env_var" in config:
+        alt_key = os.environ.get(config["alt_env_var"], "").strip()
+        if alt_key:
+            return alt_key
+
+    # Check primary environment variable
     env_key = os.environ.get(config["env_var"], "").strip()
     if env_key:
         return env_key
@@ -154,7 +173,7 @@ def get_auth_status(provider: str) -> AuthStatus:
     """Get authentication status for a provider.
 
     Args:
-        provider: Provider name (runpod, digitalocean, modal)
+        provider: Provider name (runpod, digitalocean, modal, anthropic, openai)
 
     Returns:
         AuthStatus with details about the auth state
@@ -164,7 +183,20 @@ def get_auth_status(provider: str) -> AuthStatus:
 
     config = PROVIDERS[provider]
 
-    # Check environment variable first
+    # Check alt environment variable first (e.g., WAFER_ANTHROPIC_API_KEY)
+    if "alt_env_var" in config:
+        alt_key = os.environ.get(config["alt_env_var"], "").strip()
+        if alt_key:
+            return AuthStatus(
+                provider=provider,
+                display_name=config["display_name"],
+                is_authenticated=True,
+                source="env",
+                key_preview=_format_key_preview(alt_key),
+                key_url=config["key_url"],
+            )
+
+    # Check primary environment variable
     env_key = os.environ.get(config["env_var"], "").strip()
     if env_key:
         return AuthStatus(

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/environments/coding.py

@@ -34,6 +34,8 @@ from wafer_core.tools import (
     GLOB_TOOL,
     GREP_TOOL,
     READ_TOOL,
+    SEARCH_DOCS_TOOL,
+    SKILL_TOOL,
     WRITE_TOOL,
     ApprovalCallback,
     exec_bash,
@@ -41,6 +43,8 @@ from wafer_core.tools import (
     exec_glob,
     exec_grep,
     exec_read,
+    exec_search_docs,
+    exec_skill,
     exec_write,
 )
 
@@ -61,6 +65,8 @@ ALL_TOOLS = {
     "glob": GLOB_TOOL,
     "grep": GREP_TOOL,
     "bash": BASH_TOOL,
+    "search_docs": SEARCH_DOCS_TOOL,
+    "skill": SKILL_TOOL,
     # TODO(wafer-tool): "wafer": WAFER_TOOL,
 }
 
@@ -208,6 +214,8 @@ class CodingEnvironment:
                 self.bash_approval_callback,
                 self._sandbox_policy,
             ),
+            "search_docs": lambda tc: exec_search_docs(tc),
+            "skill": lambda tc: exec_skill(tc),
             # TODO(wafer-tool): "wafer": lambda tc: exec_wafer(
             #     tc, self.working_dir, self.enabled_tools, self.allow_spawn, cancel_scope
             # ),

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/rollouts/dtypes.py

@@ -1562,6 +1562,10 @@ class EvalConfig:
     resume_dir: Path | None = None
     report_batch_size: int = 1  # Write report after each sample for best recovery
 
+    # Custom metadata (flows to report.json for dashboard filtering)
+    # e.g., {"waferbench_category": "gemm", "github_runner": "elliot"}
+    metadata: dict[str, Any] | None = None
+
 
 # ── Session Types ──────────────────────────────────────────────────────────────
 # Types for persisting agent sessions (trajectories, config, environment state).

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/rollouts/environments/localfs.py

@@ -331,9 +331,9 @@ def generate_diff(old_content: str, new_content: str, context_lines: int = 3) ->
 
 # Tool preset configurations
 TOOL_PRESETS = {
-    "full": ["read", "write", "edit", "bash", "web_fetch"],
+    "full": ["read", "write", "edit", "bash", "web_fetch", "skill"],
     "readonly": ["read"],
-    "no-write": ["read", "edit", "bash", "web_fetch"],
+    "no-write": ["read", "edit", "bash", "web_fetch", "skill"],
 }
 
 
@@ -630,6 +630,24 @@ class LocalFilesystemEnvironment:
                     required=["url", "prompt"],
                 ),
             ),
+            # skill tool
+            Tool(
+                type="function",
+                function=ToolFunction(
+                    name="skill",
+                    description="Load a skill's full instructions. Skills provide domain-specific knowledge and workflows. Use this when you need detailed guidance for a task mentioned in your available skills.",
+                    parameters=ToolFunctionParameter(
+                        type="object",
+                        properties={
+                            "name": {
+                                "type": "string",
+                                "description": "Name of the skill to load (e.g., 'wafer-guide')",
+                            },
+                        },
+                    ),
+                    required=["name"],
+                ),
+            ),
         ]
 
     async def on_assistant_message(self, message: Message, state: AgentState) -> AgentState:
@@ -655,6 +673,8 @@ class LocalFilesystemEnvironment:
                 return await self._exec_bash(tool_call, current_state.session_id, cancel_scope)
             elif tool_call.name == "web_fetch":
                 return await self._exec_web_fetch(tool_call, current_state.session_id)
+            elif tool_call.name == "skill":
+                return await self._exec_skill(tool_call)
             else:
                 return ToolResult(
                     tool_call_id=tool_call.id,
@@ -1155,3 +1175,31 @@ class LocalFilesystemEnvironment:
             content=header + final_content,
             details={"output_file": output_file_path} if output_file_path else None,
         )
+
+    async def _exec_skill(self, tool_call: ToolCall) -> ToolResult:
+        """Load a skill's full instructions."""
+        from ..skills import load_skill
+
+        skill_name = tool_call.args["name"]
+        skill = load_skill(skill_name)
+
+        if skill is None:
+            # List available skills in error message
+            from ..skills import discover_skills
+
+            available = discover_skills()
+            available_names = [s.name for s in available]
+            return ToolResult(
+                tool_call_id=tool_call.id,
+                is_error=True,
+                content="",
+                error=f"Skill not found: {skill_name}. Available skills: {', '.join(available_names) or 'none'}",
+            )
+
+        # Return the full skill content
+        header = f"# Skill: {skill.name}\n\n"
+        return ToolResult(
+            tool_call_id=tool_call.id,
+            is_error=False,
+            content=header + skill.content,
+        )

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/rollouts/evaluation.py

@@ -642,6 +642,7 @@ class EvalReport:
     timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
     git_info: dict[str, Any] = field(default_factory=_get_git_info)
     config_path: str | None = None  # Path to config file relative to repo root
+    metadata: dict[str, Any] | None = None  # Custom metadata (waferbench_category, github_runner, etc.)
 
     async def save(self, output_dir: Path) -> None:
         """Save evaluation results to directory."""
@@ -675,6 +676,8 @@ class EvalReport:
             "config_path": self.config_path,
             "sample_ids": [s.id for s in self.sample_results],
         }
+        if self.metadata:
+            summary["metadata"] = self.metadata
         # Sanitize API keys in the summary before saving
         summary = sanitize_api_keys(summary)
         report_file = output_dir / "report.json"
@@ -761,6 +764,9 @@ def _write_partial_report(
         "config_path": config.config_path,
     }
 
+    if config.metadata:
+        report["metadata"] = config.metadata
+
     if resume_from:
         report["resume_from"] = resume_from
 
@@ -1279,6 +1285,7 @@ async def evaluate(
             summary_metrics={},
             sample_results=[],
             config={"resumed_from": str(config.resume_dir)},
+            metadata=config.metadata,
         )
 
     if config.verbose:
@@ -1489,6 +1496,7 @@ async def evaluate(
             "evaluation_timestamp": datetime.now().isoformat(),
         },
         config_path=config.config_path,
+        metadata=config.metadata,
     )
 
     # Save if output directory specified
@@ -1546,7 +1554,7 @@ def compute_summary_metrics(results: list[Sample]) -> dict[str, float]:
             for m in r.score.metrics:
                 all_metric_names.add(m.name)
 
-    # Compute mean, min, max, std for each metric
+    # Compute mean, median, min, max, std for each metric
     for metric_name in all_metric_names:
         values = []
         for r in results:
@@ -1557,7 +1565,15 @@ def compute_summary_metrics(results: list[Sample]) -> dict[str, float]:
                         break
         if values:
             mean_val = sum(values) / len(values)
+            sorted_values = sorted(values)
+            n = len(sorted_values)
+            if n % 2 == 0:
+                median_val = (sorted_values[n // 2 - 1] + sorted_values[n // 2]) / 2
+            else:
+                median_val = sorted_values[n // 2]
+            
             summary[f"mean_{metric_name}"] = mean_val
+            summary[f"median_{metric_name}"] = median_val
             summary[f"min_{metric_name}"] = min(values)
             summary[f"max_{metric_name}"] = max(values)
             summary[f"std_{metric_name}"] = (

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/rollouts/prompt.py

@@ -232,7 +232,7 @@ Detailed docs: {docs}
 If asked about your capabilities, read these files."""
 
 
-def build_system_prompt(
+def build_system_prompt(  # noqa: PLR0913
     env_name: str,
     tools: list[Tool],
     cwd: Path | None = None,
@@ -240,6 +240,7 @@ def build_system_prompt(
     env_system_prompt: str | None = None,
     include_self_docs: bool = True,
     include_project_context: bool = True,
+    include_skills: bool = False,
 ) -> str:
     """Build complete system prompt with dynamic tool info.
 
@@ -251,6 +252,7 @@ def build_system_prompt(
         env_system_prompt: Environment-provided system prompt (from env.get_system_prompt())
         include_self_docs: Whether to include rollouts documentation paths
         include_project_context: Whether to load AGENTS.md/ROLLOUTS.md files
+        include_skills: Whether to discover and list available skills
     """
     # Assertions (Tiger Style: 2+ per function, split compound)
     assert env_name, "env_name required"
@@ -277,11 +279,19 @@ def build_system_prompt(
     if guidelines:
         sections.append("Guidelines:\n" + "\n".join(f"- {g}" for g in guidelines))
 
-    # 5. Self-documentation
+    # 5. Available skills (metadata only - agent loads full content via skill tool)
+    if include_skills:
+        from .skills import discover_skills, format_skill_metadata_for_prompt
+
+        skill_metadata = discover_skills()
+        if skill_metadata:
+            sections.append(format_skill_metadata_for_prompt(skill_metadata))
+
+    # 6. Self-documentation
     if include_self_docs:
         sections.append(build_self_doc_section())
 
-    # 6. Project context files (AGENTS.md, ROLLOUTS.md, etc.)
+    # 7. Project context files (AGENTS.md, ROLLOUTS.md, etc.)
     if include_project_context:
         context_files = load_project_context(working_dir)
         if context_files:
@@ -290,7 +300,7 @@ def build_system_prompt(
                 ctx_section += f"\n## {path}\n\n{content}\n"
             sections.append(ctx_section)
 
-    # 7. Runtime context
+    # 8. Runtime context
     now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     sections.append(f"Current time: {now}\nWorking directory: {working_dir}")
 

wafer_core-0.1.23/wafer_core/rollouts/skills.py

@@ -0,0 +1,176 @@
+"""Skill discovery and loading.
+
+Skills are documentation files that agents can load on demand.
+Format follows agentskills.io spec: SKILL.md with YAML frontmatter.
+
+Discovery order:
+1. ~/.wafer/skills/{name}/SKILL.md (user-installed)
+2. Bundled skills (wafer-cli package)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from .paths import get_config_dir
+
+
+@dataclass(frozen=True)
+class SkillMetadata:
+    """Lightweight skill metadata for system prompt injection."""
+
+    name: str
+    description: str
+    path: Path  # Path to SKILL.md file
+
+
+@dataclass(frozen=True)
+class Skill:
+    """Full skill with content."""
+
+    name: str
+    description: str
+    content: str  # Full markdown content (without frontmatter)
+    path: Path
+
+
+def _parse_skill_file(path: Path) -> tuple[dict[str, str], str] | None:
+    """Parse SKILL.md file into (frontmatter, content).
+
+    Returns None if file doesn't exist or is malformed.
+    """
+    if not path.exists():
+        return None
+
+    try:
+        text = path.read_text()
+    except (OSError, PermissionError):
+        return None
+
+    # Parse YAML frontmatter (between --- markers)
+    if not text.startswith("---"):
+        return None
+
+    # Find closing ---
+    end_idx = text.find("---", 3)
+    if end_idx == -1:
+        return None
+
+    frontmatter_text = text[3:end_idx].strip()
+    content = text[end_idx + 3 :].strip()
+
+    # Parse YAML (simple key: value format, no dependencies)
+    frontmatter: dict[str, str] = {}
+    for raw_line in frontmatter_text.split("\n"):
+        stripped = raw_line.strip()
+        if not stripped or ":" not in stripped:
+            continue
+        key, _, value = stripped.partition(":")
+        frontmatter[key.strip()] = value.strip()
+
+    # Validate required fields
+    if "name" not in frontmatter or "description" not in frontmatter:
+        return None
+
+    return frontmatter, content
+
+
+def _get_bundled_skills_dir() -> Path | None:
+    """Get path to bundled skills in wafer-cli package."""
+    # Try to find wafer-cli's skills directory
+    try:
+        import wafer
+
+        wafer_cli_path = Path(wafer.__file__).parent
+        skills_dir = wafer_cli_path / "skills"
+        if skills_dir.exists():
+            return skills_dir
+    except ImportError:
+        pass
+
+    return None
+
+
+def discover_skills() -> list[SkillMetadata]:
+    """Discover all available skills.
+
+    Returns list of SkillMetadata (name + description only).
+    """
+    skills: dict[str, SkillMetadata] = {}
+
+    # 1. User-installed skills (~/.wafer/skills/)
+    user_skills_dir = get_config_dir() / "skills"
+    if user_skills_dir.exists():
+        for skill_dir in user_skills_dir.iterdir():
+            if not skill_dir.is_dir():
+                continue
+            skill_file = skill_dir / "SKILL.md"
+            parsed = _parse_skill_file(skill_file)
+            if parsed:
+                frontmatter, _ = parsed
+                skills[frontmatter["name"]] = SkillMetadata(
+                    name=frontmatter["name"],
+                    description=frontmatter["description"],
+                    path=skill_file,
+                )
+
+    # 2. Bundled skills (wafer-cli package)
+    bundled_dir = _get_bundled_skills_dir()
+    if bundled_dir:
+        for skill_dir in bundled_dir.iterdir():
+            if not skill_dir.is_dir():
+                continue
+            skill_file = skill_dir / "SKILL.md"
+            parsed = _parse_skill_file(skill_file)
+            if parsed:
+                frontmatter, _ = parsed
+                # User skills take precedence
+                if frontmatter["name"] not in skills:
+                    skills[frontmatter["name"]] = SkillMetadata(
+                        name=frontmatter["name"],
+                        description=frontmatter["description"],
+                        path=skill_file,
+                    )
+
+    return list(skills.values())
+
+
+def load_skill(name: str) -> Skill | None:
+    """Load a skill by name.
+
+    Returns full Skill with content, or None if not found.
+    """
+    # Find the skill
+    for metadata in discover_skills():
+        if metadata.name == name:
+            parsed = _parse_skill_file(metadata.path)
+            if parsed:
+                frontmatter, content = parsed
+                return Skill(
+                    name=frontmatter["name"],
+                    description=frontmatter["description"],
+                    content=content,
+                    path=metadata.path,
+                )
+    return None
+
+
+def format_skill_metadata_for_prompt(skills: list[SkillMetadata]) -> str:
+    """Format skill metadata for system prompt injection.
+
+    Returns a compact section listing available skills.
+    """
+    if not skills:
+        return ""
+
+    lines = ["## Available Skills", ""]
+    lines.append(
+        "You have access to the following skills. Use the `skill` tool to load full instructions when needed."
+    )
+    lines.append("")
+
+    for skill in skills:
+        lines.append(f"- **{skill.name}**: {skill.description}")
+
+    return "\n".join(lines)

{wafer_core-0.1.21 → wafer_core-0.1.23}/wafer_core/rollouts/templates/base.py

@@ -222,6 +222,9 @@ class TemplateConfig:
     # Example: {"corpus": "./docs/", "format": "markdown"}
     defaults: dict[str, str] = field(default_factory=dict)
 
+    # Skill discovery - if True, discovers skills and adds skill tool
+    include_skills: bool = False
+
     def interpolate_prompt(self, args: dict[str, str] | None = None) -> str:
         """Interpolate template variables into the system prompt.