wafer-core 0.1.24__py3-none-any.whl → 0.1.25__py3-none-any.whl

wafer_core/environments/coding.py

@@ -34,7 +34,6 @@ from wafer_core.tools import (
     GLOB_TOOL,
     GREP_TOOL,
     READ_TOOL,
-    SEARCH_DOCS_TOOL,
     SKILL_TOOL,
     WRITE_TOOL,
     ApprovalCallback,
@@ -43,7 +42,6 @@ from wafer_core.tools import (
     exec_glob,
     exec_grep,
     exec_read,
-    exec_search_docs,
     exec_skill,
     exec_write,
 )
@@ -65,7 +63,6 @@ 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,
 }
@@ -214,7 +211,6 @@ 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/rollouts/upload.py

@@ -7,10 +7,13 @@ import logging
 import os
 from pathlib import Path
 
+import httpx
+
 logger = logging.getLogger(__name__)
 
 SUPABASE_URL = "https://hvlpthcnxlywlquiciqe.supabase.co"
 BUCKET_NAME = "traces"
+API_BASE = os.environ.get("WAFER_API_URL", "https://api.wafer.ai")
 
 
 def upload_results_to_supabase(output_dir: Path, log: logging.Logger | None = None) -> bool:
@@ -95,6 +98,12 @@ def upload_results_to_supabase(output_dir: Path, log: logging.Logger | None = No
             )
 
         log.info(f"Uploaded {len(uploaded)} files to Supabase: {run_name}")
+
+        # Auto-index in database for trace viewer
+        # Fail if indexing fails - user can re-run (everything is idempotent)
+        if not _index_run_in_database(run_name, report_path, log):
+            return False
+
         return True
 
     except ImportError:
@@ -103,3 +112,39 @@ def upload_results_to_supabase(output_dir: Path, log: logging.Logger | None = No
     except Exception as e:
         log.error(f"Failed to upload to Supabase: {e}")
         return False
+
+
+def _index_run_in_database(run_name: str, report_path: Path, log: logging.Logger) -> bool:
+    """Index a run in the trace_runs database table for fast querying.
+
+    Calls POST /v1/eval-traces/runs to upsert the run metadata.
+    This enables the trace viewer to show the run immediately without manual sync.
+
+    Args:
+        run_name: Name of the run (folder name)
+        report_path: Path to the report.json file
+        log: Logger instance
+
+    Returns:
+        True if indexing succeeded, False otherwise
+    """
+    try:
+        with open(report_path) as f:
+            report = json.load(f)
+
+        response = httpx.post(
+            f"{API_BASE}/v1/eval-traces/runs",
+            json={"name": run_name, "report": report},
+            timeout=30.0,
+        )
+
+        if response.status_code == 200:
+            log.info(f"Indexed run in database: {run_name}")
+            return True
+        else:
+            log.error(f"Failed to index run {run_name}: {response.status_code} {response.text}")
+            return False
+
+    except Exception as e:
+        log.error(f"Failed to index run {run_name} in database: {e}")
+        return False

wafer_core/tools/__init__.py

@@ -72,10 +72,6 @@ from wafer_core.tools.write_kernel_tool import (
     KernelSubmission,
     exec_write_kernel,
 )
-from wafer_core.tools.search_docs_tool import (
-    SEARCH_DOCS_TOOL,
-    exec_search_docs,
-)
 
 __all__ = [
     # File tools
@@ -137,7 +133,4 @@ __all__ = [
     "exec_tracelens_report",
     "exec_tracelens_compare",
     "exec_tracelens_collective",
-    # Search docs tool
-    "SEARCH_DOCS_TOOL",
-    "exec_search_docs",
 ]

{wafer_core-0.1.24.dist-info → wafer_core-0.1.25.dist-info}/METADATA

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

{wafer_core-0.1.24.dist-info → wafer_core-0.1.25.dist-info}/RECORD

@@ -12,7 +12,7 @@ wafer_core/config/__init__.py,sha256=hKywfjA4YXd4lBeBFEcBoMwFoflPHJTiBnkTq7_JYOQ
 wafer_core/config/loader.py,sha256=k7JnILmO13TWUzIv9Lm8fvmj3UfYHZDgaFurjQ-GXpY,6623
 wafer_core/config/schema.py,sha256=2WhFlnG0VYYX4T-70BLeJK8Janvi4KEa8KKGZA7331w,3898
 wafer_core/environments/__init__.py,sha256=SIsResVtm22tr_d-oHPeeSxrkhFdmPOFico3DqtRqK8,238
-wafer_core/environments/coding.py,sha256=N-ELZwJu5vKLCVtwO25c6JSty6fmqf85VR2d3WJ4RXw,8559
+wafer_core/environments/coding.py,sha256=T-_JFU-n5OxPR8xAWp8qar4Y5xyC-TWTIBjRy4PDel8,8418
 wafer_core/environments/gpumode.py,sha256=8Da08nltvN_YloNyYI6-omN2D4n5C7aptKDCtUgT2bQ,17191
 wafer_core/lib/__init__.py,sha256=4-4p3mhwlquejWGglYXU8_nHdA0LoPaa_jGzcm13USA,1325
 wafer_core/lib/kernel_scope/__init__.py,sha256=WW2vu8jUlqOu-MCpgO40lIYacCA9N2u-uuECIs_JO2w,2817
@@ -359,7 +359,7 @@ wafer_core/rollouts/skills.py,sha256=ATYoG02Cc6_VrtE415TnseBFJrKOMq27z-5YgBgPpZQ
 wafer_core/rollouts/slice.py,sha256=darOZO53BuSPfvv_KjOSzulGVSWbL4OuoE3k6xXpBFg,20195
 wafer_core/rollouts/store.py,sha256=UDP9idDOEVs_0Pslx0K_Y8E1i-BeoqVSaxdQiaqtz1E,18051
 wafer_core/rollouts/transform_messages.py,sha256=yldzdLgugNYb5Zxju7myFBel1tmrHXx9M399ImqPLGI,20891
-wafer_core/rollouts/upload.py,sha256=dR0W6QREQWIn7F0j95K7ut6uQ5phS95ijnYJK0J24bo,3432
+wafer_core/rollouts/upload.py,sha256=hEqZfwgb0b4GYbrwRSA3fuqF70pqo6hyaQU59j3vM7E,4890
 wafer_core/rollouts/_logging/__init__.py,sha256=rCXeAssQ3gIrduuMzvKPD-ikt6rXejVL9h5XtDRyIQg,498
 wafer_core/rollouts/_logging/color_formatter.py,sha256=x3qRKwHsUCFkgcIl8x_Ajjw82X2EedbTe14sCxMU4Kc,2267
 wafer_core/rollouts/_logging/json_formatter.py,sha256=jJIa2IZCsu2C_Y1HXQi7hbI33x6L6shN_dqu-hmhxp4,2380
@@ -586,11 +586,10 @@ wafer_core/sessions/hooks.py,sha256=A-txm6ufnRGQCdtP3vwh7oEOdlLN9Tv0XsjORMihuAI,
 wafer_core/targets/__init__.py,sha256=sHndC7AAOaHXlrmDXFLB53a5Y8DBjuyqS6nwsO2nj-Y,1728
 wafer_core/targets/digitalocean.py,sha256=cvoYpYjtSyy5t2lQAPi7ERruuuibronah_ivOiduAHQ,16550
 wafer_core/targets/runpod.py,sha256=LrVmNvA6qjzL5nbGSWvtw7CHrK6bDu7_o3vKIek00Tc,20286
-wafer_core/tools/__init__.py,sha256=deGQQlcdSD6zQx8JHizfSXgF5-EntdBOF_ngtob1-VU,3506
+wafer_core/tools/__init__.py,sha256=wBQD45GdSfkxcT6NHzIv0IMeXCc0enwwkpm3T_9j1X8,3341
 wafer_core/tools/bash_tool.py,sha256=daoKOVGSgL0x9X_3l8Apd6-wFH4VMXMGJwVemw2FIfc,16828
 wafer_core/tools/glob_tool.py,sha256=9X5PdOjQJj7kiVNqqCZC0-1LmnE6wHx3Zc9zfMjtXdc,3533
 wafer_core/tools/grep_tool.py,sha256=cStyDz-J47oDLLZCL83yOvYo8Ijv4qu3D372JKT_ptM,4580
-wafer_core/tools/search_docs_tool.py,sha256=WY4hY83sseX8Fpxvw6DZxiG-F95F2t3-4PyfMD1Lpkg,6809
 wafer_core/tools/skill_tool.py,sha256=JXsT5hBTUH5U4tmzHEywU7eHHt5xCEF79tL2tsuk4-c,2067
 wafer_core/tools/wafer_tool.py,sha256=-dgPTHbWXq3I3wFj0mP7-lj5iZqGRoFvFf9IEEo3plQ,6345
 wafer_core/tools/write_kernel_tool.py,sha256=dJjhr-WBhVNe06hcJQVmBZTbS8mid64KF1MwlE2s2R4,21547
@@ -674,6 +673,6 @@ wafer_core/utils/modal_execution/modal_app.py,sha256=VfS2cX8gHtnlPXemmMcEwDPeQdh
 wafer_core/utils/modal_execution/modal_config.py,sha256=7cGX9TGqilQ3qxI3OFGXV5orjtyRU-PEDOJ4vP2oxno,4421
 wafer_core/utils/modal_execution/modal_execution.py,sha256=gChjnV6jqA3A7IRP3DfvV5cSfm_MN0X4f7JZufXgdZE,24594
 wafer_core/utils/modal_execution/test_modal.py,sha256=_jqou_hrLs1Daf1590Pnb0a_lXMMa2rczAPpW9HpoNQ,8153
-wafer_core-0.1.24.dist-info/METADATA,sha256=h2zO5zgoRFyd1aZbWSugm8JWl8RzYYd9w5h0CDQ2pa4,1420
-wafer_core-0.1.24.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-wafer_core-0.1.24.dist-info/RECORD,,
+wafer_core-0.1.25.dist-info/METADATA,sha256=hNv2xTrSbhdaA1G8mBmLxe1HbYRE0NLX3C6w89sht3k,1420
+wafer_core-0.1.25.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+wafer_core-0.1.25.dist-info/RECORD,,

wafer_core/tools/search_docs_tool.py

@@ -1,196 +0,0 @@
-"""Search documentation tool for GPU programming corpora.
-
-Provides semantic and keyword search over documentation for CuTeDSL, CUDA, etc.
-
-Corpora are downloaded via `wafer corpus download <name>` and stored in ~/.cache/wafer/corpora/.
-"""
-
-import re
-from pathlib import Path
-
-from wafer_core.rollouts.dtypes import Tool, ToolCall, ToolFunction, ToolFunctionParameter, ToolResult
-
-# Cache directory where wafer corpus download stores files
-CACHE_DIR = Path.home() / ".cache" / "wafer" / "corpora"
-
-# Available corpora (names match wafer corpus download)
-AVAILABLE_CORPORA = ["cutlass", "cutedsl", "cuda", "hip", "amd"]
-
-SEARCH_DOCS_TOOL = Tool(
-    type="function",
-    function=ToolFunction(
-        name="search_docs",
-        description="""Search GPU programming documentation for relevant information.
-
-Use this tool to find documentation about:
-- CUTLASS C++ (cute:: namespace, gemm tutorials, tensor cores, TMA, Blackwell)
-- CuTeDSL Python API (@cute.kernel, @cute.jit, cute.arch functions)
-- CUDA programming concepts
-- GPU kernel optimization techniques
-- Code examples and patterns
-
-Available corpora:
-- 'cutlass' - NVIDIA CUTLASS C++ docs + GitHub examples (gemm, hopper, blackwell)
-- 'cutedsl' - CuTeDSL Python documentation
-- 'cuda' - General CUDA programming docs
-- 'hip' - AMD HIP programming docs
-- 'amd' - AMD GPU kernel development (rocWMMA, CK, etc.)
-
-Note: Corpora must be downloaded first with `wafer corpus download <name>`.
-Returns relevant documentation snippets with file paths.""",
-        parameters=ToolFunctionParameter(
-            type="object",
-            properties={
-                "query": {
-                    "type": "string",
-                    "description": "Search query - describe what you're looking for",
-                },
-                "corpus": {
-                    "type": "string",
-                    "description": "Which docs to search: 'cutlass', 'cutedsl', 'cuda', 'hip', 'amd' (default: cutlass)",
-                },
-                "max_results": {
-                    "type": "integer",
-                    "description": "Maximum number of results to return (default: 5)",
-                },
-            },
-        ),
-        required=["query"],
-    )
-)
-
-
-def _get_corpus_path(corpus_name: str) -> Path | None:
-    """Get the path to a corpus in the cache directory.
-    
-    Corpora are stored at ~/.cache/wafer/corpora/<corpus_name>/
-    """
-    if corpus_name not in AVAILABLE_CORPORA:
-        return None
-    
-    corpus_path = CACHE_DIR / corpus_name
-    if corpus_path.exists():
-        return corpus_path
-    
-    return None
-
-
-def _search_files(corpus_path: Path, query: str, max_results: int = 5) -> list[dict]:
-    """Simple keyword search through documentation files."""
-    results = []
-    query_terms = query.lower().split()
-    
-    # Search .md, .py, .cu, .hpp, and .h files (for CUTLASS examples)
-    for pattern in ["**/*.md", "**/*.py", "**/*.cu", "**/*.hpp", "**/*.h", "**/*.cuh"]:
-        for file_path in corpus_path.glob(pattern):
-            if file_path.is_file():
-                try:
-                    content = file_path.read_text(encoding="utf-8", errors="ignore")
-                    content_lower = content.lower()
-                    
-                    # Score based on term matches
-                    score = sum(content_lower.count(term) for term in query_terms)
-                    
-                    if score > 0:
-                        # Extract relevant snippets
-                        snippets = _extract_snippets(content, query_terms)
-                        results.append({
-                            "file": str(file_path),  # Return absolute path so read tool can access it
-                            "score": score,
-                            "snippets": snippets[:3],  # Top 3 snippets
-                        })
-                except Exception:
-                    continue
-    
-    # Sort by score and return top results
-    results.sort(key=lambda x: x["score"], reverse=True)
-    return results[:max_results]
-
-
-def _extract_snippets(content: str, terms: list[str], context_lines: int = 5) -> list[str]:
-    """Extract snippets containing search terms."""
-    snippets = []
-    lines = content.split("\n")
-    
-    for i, line in enumerate(lines):
-        line_lower = line.lower()
-        if any(term in line_lower for term in terms):
-            # Get context around the match
-            start = max(0, i - context_lines)
-            end = min(len(lines), i + context_lines + 1)
-            snippet = "\n".join(lines[start:end])
-            
-            # Skip very short snippets
-            if len(snippet.strip()) > 50:
-                snippets.append(snippet)
-    
-    return snippets
-
-
-async def exec_search_docs(
-    tool_call: ToolCall,
-    corpus_override: str | None = None,
-) -> ToolResult:
-    """Execute search_docs tool.
-    
-    Args:
-        tool_call: The tool call with query and optional corpus
-        corpus_override: Override corpus path (for testing)
-    """
-    query = tool_call.args.get("query", "")
-    corpus_name = tool_call.args.get("corpus", "cutlass")
-    max_results = tool_call.args.get("max_results", 5)
-    
-    if not query:
-        return ToolResult(
-            tool_call_id=tool_call.id,
-            content="",
-            error="query parameter is required",
-        )
-    
-    # Find corpus path
-    if corpus_override:
-        corpus_path = Path(corpus_override)
-    else:
-        corpus_path = _get_corpus_path(corpus_name)
-        if corpus_path is None:
-            return ToolResult(
-                tool_call_id=tool_call.id,
-                content="",
-                error=f"Unknown corpus: {corpus_name}. Available: {AVAILABLE_CORPORA}",
-            )
-    
-    if not corpus_path.exists():
-        return ToolResult(
-            tool_call_id=tool_call.id,
-            content="",
-            error=f"Corpus '{corpus_name}' not downloaded. Run: wafer corpus download {corpus_name}",
-        )
-    
-    # Search
-    results = _search_files(corpus_path, query, max_results)
-    
-    if not results:
-        return ToolResult(
-            tool_call_id=tool_call.id,
-            content=f"No results found for query: {query}",
-            error=None,
-        )
-    
-    # Format output
-    output_parts = [f"Found {len(results)} results for: {query}\n"]
-    
-    for i, result in enumerate(results, 1):
-        output_parts.append(f"\n{'='*60}")
-        output_parts.append(f"[{i}] {result['file']} (score: {result['score']})")
-        output_parts.append("=" * 60)
-        
-        for snippet in result["snippets"]:
-            output_parts.append(snippet)
-            output_parts.append("-" * 40)
-    
-    return ToolResult(
-        tool_call_id=tool_call.id,
-        content="\n".join(output_parts),
-        error=None,
-    )