@fastino-ai/pioneer-cli 0.1.0 → 0.2.0

package/src/tools/wandb.ts

@@ -0,0 +1,348 @@
+/**
+ * Weights & Biases integration for experiment tracking
+ */
+
+import { spawn } from "child_process";
+import type { Tool, ToolResult } from "../agent/types.js";
+
+export interface WandbConfig {
+  apiKey?: string;
+  entity?: string;
+  project?: string;
+  timeout?: number;
+}
+
+// Check if wandb CLI is available
+export async function isWandbAvailable(): Promise<boolean> {
+  return new Promise((resolve) => {
+    const proc = spawn("wandb", ["--version"], { stdio: "pipe" });
+    proc.on("close", (code) => resolve(code === 0));
+    proc.on("error", () => resolve(false));
+  });
+}
+
+// Run wandb command
+async function runWandb(
+  args: string[],
+  options: WandbConfig = {}
+): Promise<{ stdout: string; stderr: string; exitCode: number }> {
+  const timeout = options.timeout || 60000;
+
+  return new Promise((resolve) => {
+    let stdout = "";
+    let stderr = "";
+    let timedOut = false;
+
+    const env: NodeJS.ProcessEnv = { ...process.env };
+    if (options.apiKey) env.WANDB_API_KEY = options.apiKey;
+    if (options.entity) env.WANDB_ENTITY = options.entity;
+    if (options.project) env.WANDB_PROJECT = options.project;
+
+    const proc = spawn("wandb", args, {
+      stdio: ["pipe", "pipe", "pipe"],
+      env,
+    });
+
+    const timer = setTimeout(() => {
+      timedOut = true;
+      proc.kill("SIGTERM");
+    }, timeout);
+
+    proc.stdout.on("data", (data: Buffer) => {
+      stdout += data.toString();
+    });
+
+    proc.stderr.on("data", (data: Buffer) => {
+      stderr += data.toString();
+    });
+
+    proc.on("close", (code) => {
+      clearTimeout(timer);
+      resolve({
+        stdout,
+        stderr,
+        exitCode: timedOut ? -1 : (code ?? -1),
+      });
+    });
+
+    proc.on("error", (err) => {
+      clearTimeout(timer);
+      resolve({
+        stdout: "",
+        stderr: err.message,
+        exitCode: -1,
+      });
+    });
+  });
+}
+
+export function createWandbLogTool(config: WandbConfig = {}): Tool {
+  return {
+    name: "wandb_log",
+    description: "Log metrics, artifacts, or data to Weights & Biases for experiment tracking.",
+    parameters: [
+      {
+        name: "run_name",
+        type: "string",
+        description: "Name for the W&B run",
+        required: true,
+      },
+      {
+        name: "metrics",
+        type: "object",
+        description: "Metrics to log as JSON object (e.g., {\"loss\": 0.5, \"accuracy\": 0.9})",
+        required: true,
+      },
+      {
+        name: "project",
+        type: "string",
+        description: "W&B project name (optional, uses config default)",
+        required: false,
+      },
+    ],
+
+    async execute(args: Record<string, unknown>): Promise<ToolResult> {
+      const runName = args.run_name as string;
+      const metrics = args.metrics as Record<string, unknown>;
+      const project = (args.project as string) || config.project;
+
+      if (!runName || !metrics) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "run_name and metrics are required",
+        };
+      }
+
+      const available = await isWandbAvailable();
+      if (!available) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "wandb CLI not installed. Run: pip install wandb",
+        };
+      }
+
+      // Create a Python script to log to wandb
+      const fs = await import("fs");
+      const path = await import("path");
+      const os = await import("os");
+
+      const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "wandb-"));
+      const tempFile = path.join(tempDir, "log_metrics.py");
+
+      const pythonCode = `
+import wandb
+import json
+
+metrics = json.loads('''${JSON.stringify(metrics)}''')
+run = wandb.init(
+    project="${project || "pioneer-agent"}",
+    name="${runName}",
+    reinit=True
+)
+wandb.log(metrics)
+wandb.finish()
+print(f"Logged metrics to W&B run: {run.url}")
+`;
+
+      fs.writeFileSync(tempFile, pythonCode);
+
+      try {
+        const env: NodeJS.ProcessEnv = { ...process.env };
+        if (config.apiKey) env.WANDB_API_KEY = config.apiKey;
+        if (config.entity) env.WANDB_ENTITY = config.entity;
+
+        return new Promise((resolve) => {
+          let stdout = "";
+          let stderr = "";
+
+          const proc = spawn("python3", [tempFile], {
+            stdio: ["pipe", "pipe", "pipe"],
+            env,
+          });
+
+          proc.stdout.on("data", (data: Buffer) => {
+            stdout += data.toString();
+          });
+
+          proc.stderr.on("data", (data: Buffer) => {
+            stderr += data.toString();
+          });
+
+          proc.on("close", (code) => {
+            let output = stdout;
+            if (stderr && !stderr.includes("wandb:")) {
+              output += (output ? "\n" : "") + stderr;
+            }
+
+            resolve({
+              toolCallId: "",
+              output: output || "Metrics logged successfully",
+              success: code === 0,
+              error: code !== 0 ? `Exit code: ${code}` : undefined,
+            });
+          });
+
+          proc.on("error", (err) => {
+            resolve({
+              toolCallId: "",
+              output: "",
+              success: false,
+              error: err.message,
+            });
+          });
+        });
+      } finally {
+        try {
+          fs.rmSync(tempDir, { recursive: true, force: true });
+        } catch {
+          // Ignore
+        }
+      }
+    },
+  };
+}
+
+export function createWandbQueryTool(config: WandbConfig = {}): Tool {
+  return {
+    name: "wandb_query",
+    description: "Query runs and metrics from Weights & Biases.",
+    parameters: [
+      {
+        name: "project",
+        type: "string",
+        description: "W&B project name to query",
+        required: true,
+      },
+      {
+        name: "filters",
+        type: "object",
+        description: "Filters for the query as JSON (e.g., {\"state\": \"finished\"})",
+        required: false,
+      },
+      {
+        name: "limit",
+        type: "number",
+        description: "Maximum number of runs to return (default: 10)",
+        required: false,
+      },
+    ],
+
+    async execute(args: Record<string, unknown>): Promise<ToolResult> {
+      const project = args.project as string;
+      const filters = (args.filters as Record<string, unknown>) || {};
+      const limit = (args.limit as number) || 10;
+
+      if (!project) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "project is required",
+        };
+      }
+
+      const available = await isWandbAvailable();
+      if (!available) {
+        return {
+          toolCallId: "",
+          output: "",
+          success: false,
+          error: "wandb CLI not installed. Run: pip install wandb",
+        };
+      }
+
+      const fs = await import("fs");
+      const path = await import("path");
+      const os = await import("os");
+
+      const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "wandb-query-"));
+      const tempFile = path.join(tempDir, "query_runs.py");
+
+      const entity = config.entity ? `"${config.entity}"` : "None";
+      const filtersJson = JSON.stringify(filters);
+
+      const pythonCode = `
+import wandb
+import json
+
+api = wandb.Api()
+entity = ${entity}
+project = "${project}"
+
+path = f"{entity}/{project}" if entity else project
+runs = api.runs(path, filters=json.loads('${filtersJson}'))
+
+results = []
+for run in list(runs)[:${limit}]:
+    results.append({
+        "id": run.id,
+        "name": run.name,
+        "state": run.state,
+        "url": run.url,
+        "summary": dict(run.summary),
+        "config": dict(run.config),
+    })
+
+print(json.dumps(results, indent=2, default=str))
+`;
+
+      fs.writeFileSync(tempFile, pythonCode);
+
+      try {
+        const env: NodeJS.ProcessEnv = { ...process.env };
+        if (config.apiKey) env.WANDB_API_KEY = config.apiKey;
+
+        return new Promise((resolve) => {
+          let stdout = "";
+          let stderr = "";
+
+          const proc = spawn("python3", [tempFile], {
+            stdio: ["pipe", "pipe", "pipe"],
+            env,
+          });
+
+          proc.stdout.on("data", (data: Buffer) => {
+            stdout += data.toString();
+          });
+
+          proc.stderr.on("data", (data: Buffer) => {
+            stderr += data.toString();
+          });
+
+          proc.on("close", (code) => {
+            resolve({
+              toolCallId: "",
+              output: stdout || stderr || "(no results)",
+              success: code === 0,
+              error: code !== 0 ? `Exit code: ${code}` : undefined,
+            });
+          });
+
+          proc.on("error", (err) => {
+            resolve({
+              toolCallId: "",
+              output: "",
+              success: false,
+              error: err.message,
+            });
+          });
+        });
+      } finally {
+        try {
+          fs.rmSync(tempDir, { recursive: true, force: true });
+        } catch {
+          // Ignore
+        }
+      }
+    },
+  };
+}
+
+export function createWandbTools(config: WandbConfig = {}): Tool[] {
+  return [createWandbLogTool(config), createWandbQueryTool(config)];
+}
+