pi-extension-wandb 0.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Kiran Gadhave
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,88 @@
+# pi-extension-wandb
+
+A [pi coding agent](https://github.com/earendil-works/pi) extension that adds
+[Weights & Biases Inference](https://wandb.ai/site/inference/) as a model
+provider. Model list is discovered dynamically from W&B's OpenAI-compatible
+`/v1/models` endpoint, so new models show up without code changes.
+
+## Install
+
+Get an API key from <https://wandb.ai/authorize>, then:
+
+```bash
+export WANDB_API_KEY=...               # required
+export WANDB_PROJECT=team/project      # optional, recommended
+
+pi install https://github.com/kirangadhave/pi-extension-wandb
+```
+
+Verify:
+
+```bash
+pi --list-models | grep wandb
+```
+
+## Use
+
+```bash
+pi --provider wandb --model deepseek-ai/DeepSeek-V3.1 "hello"
+```
+
+## Configuration
+
+| Env var | Required | Notes |
+| --- | --- | --- |
+| `WANDB_API_KEY` | yes | From <https://wandb.ai/authorize>. If unset, the provider is silently skipped (no `wandb/*` models appear in `--list-models`). |
+| `WANDB_PROJECT` | no\* | Sent as the `OpenAI-Project` header. Format: `team/project`. **If you belong to multiple W&B teams, set this explicitly** &mdash; otherwise W&B picks a default team for attribution, which can land usage on the wrong team. Single-team users can safely omit. |
+| `WANDB_DEBUG` | no | `1` enables stderr logging from the extension (fetch errors, cache hits, etc.). |
+| `WANDB_NO_CACHE` | no | `1` bypasses the model-list cache and fetches fresh on every pi startup. |
+
+## How it works
+
+On every pi startup the extension:
+
+1. Reads `~/.cache/pi-extension-wandb/models.json` if it's fresh (< 1 hour old).
+2. Otherwise fetches `https://api.inference.wandb.ai/v1/models` (5s timeout)
+   and writes the result to the cache.
+3. Registers each returned model under the `wandb` provider with
+   `api: "openai-completions"`.
+
+If the fetch fails and there is no usable cache, the provider is not
+registered &mdash; `pi --list-models` will not show wandb models, and pi falls
+back to its other configured providers cleanly.
+
+## Known limitations
+
+- **Context windows are guesses.** W&B's `/v1/models` doesn't expose limits.
+  The extension hardcodes known windows for a handful of popular models
+  (`KNOWN_CONTEXT_WINDOWS` in `index.ts`); everything else gets a default of
+  128K. PRs to extend the map are welcome.
+- **Reasoning detection is heuristic.** Substring match against the model id
+  (`REASONING_MODEL_PATTERNS`). Edit the list for new families.
+- **No cost tracking.** All models report zero cost. W&B publishes per-token
+  pricing &mdash; PRs to encode it are welcome.
+- **No provider-specific `compat` flags.** Some reasoning models (Qwen3
+  thinking, gpt-oss, DeepSeek-R1) may need provider-specific quirks for
+  thinking mode. If you hit a problem, open an issue with the model id and
+  the exact failure.
+
+## Contributing
+
+PRs welcome. The only file you typically need to touch is `index.ts`:
+
+- `KNOWN_CONTEXT_WINDOWS` &mdash; add accurate context sizes for new models.
+- `REASONING_MODEL_PATTERNS` &mdash; add model family substrings that support
+  reasoning/thinking.
+
+For editor type-checking, after cloning:
+
+```bash
+pnpm install
+pnpm check
+```
+
+pi loads `index.ts` directly at runtime (no build step).
+
+## License
+
+MIT &mdash; see [LICENSE](./LICENSE).

package/index.ts

@@ -0,0 +1,176 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+const BASE_URL = "https://api.inference.wandb.ai/v1";
+const CACHE_DIR = join(homedir(), ".cache", "pi-extension-wandb");
+const CACHE_FILE = join(CACHE_DIR, "models.json");
+const CACHE_TTL_MS = 60 * 60 * 1000;
+const FETCH_TIMEOUT_MS = 5_000;
+
+// Substring matches against model id (case-insensitive). Add new reasoning
+// families here; PRs welcome.
+const REASONING_MODEL_PATTERNS = [
+  "deepseek-r1",
+  "deepseek-v3.1",
+  "deepseek-v3.2",
+  "gpt-oss",
+  "qwen3",
+  "glm-4.5",
+  "glm-5",
+];
+
+// W&B's /v1/models doesn't return context-window metadata. Hardcode the ones
+// we know; unknown ids fall back to DEFAULT_CONTEXT_WINDOW. PRs welcome.
+const KNOWN_CONTEXT_WINDOWS: Record<string, number> = {
+  "deepseek-ai/DeepSeek-V3.1": 128_000,
+  "deepseek-ai/DeepSeek-R1-0528": 161_000,
+  "meta-llama/Llama-3.1-8B-Instruct": 131_072,
+  "meta-llama/Llama-3.3-70B-Instruct": 131_072,
+  "meta-llama/Llama-4-Scout-17B-16E-Instruct": 131_072,
+  "moonshotai/Kimi-K2-Instruct": 131_072,
+  "openai/gpt-oss-20b": 131_072,
+  "openai/gpt-oss-120b": 131_072,
+  "Qwen/Qwen3-235B-A22B-Instruct-2507": 262_144,
+  "Qwen/Qwen3-Coder-480B-A35B-Instruct": 262_144,
+  "zai-org/GLM-4.5": 131_072,
+};
+
+const DEFAULT_CONTEXT_WINDOW = 131_072;
+const DEFAULT_MAX_TOKENS = 16_384;
+
+type ApiModel = {
+  id: string;
+  name?: string;
+};
+
+type CachedModels = {
+  fetchedAt: number;
+  data: ApiModel[];
+};
+
+const isReasoning = (id: string): boolean => {
+  const lower = id.toLowerCase();
+  return REASONING_MODEL_PATTERNS.some((p) => lower.includes(p));
+};
+
+const prettyName = (id: string): string => {
+  if (!id.includes("/")) return id;
+  const [vendor, ...rest] = id.split("/");
+  return `${vendor} · ${rest.join("/")}`;
+};
+
+const log = (msg: string): void => {
+  if (process.env.WANDB_DEBUG === "1") {
+    console.error(`[pi-extension-wandb] ${msg}`);
+  }
+};
+
+async function fetchModels(apiKey: string, project?: string): Promise<ApiModel[]> {
+  const headers: Record<string, string> = {
+    Authorization: `Bearer ${apiKey}`,
+    Accept: "application/json",
+  };
+  if (project) headers["OpenAI-Project"] = project;
+
+  const controller = new AbortController();
+  const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS);
+  try {
+    const res = await fetch(`${BASE_URL}/models`, { headers, signal: controller.signal });
+    if (!res.ok) {
+      throw new Error(`HTTP ${res.status}: ${await res.text()}`);
+    }
+    const payload = (await res.json()) as { data?: ApiModel[] };
+    const data = payload.data ?? [];
+    if (data.length === 0) throw new Error("empty model list");
+    return data;
+  } finally {
+    clearTimeout(timer);
+  }
+}
+
+function readCache(): ApiModel[] | null {
+  try {
+    if (!existsSync(CACHE_FILE)) return null;
+    const cached = JSON.parse(readFileSync(CACHE_FILE, "utf8")) as CachedModels;
+    return cached.data;
+  } catch (err) {
+    log(`cache read failed: ${(err as Error).message}`);
+    return null;
+  }
+}
+
+function isCacheFresh(): boolean {
+  try {
+    if (!existsSync(CACHE_FILE)) return false;
+    const cached = JSON.parse(readFileSync(CACHE_FILE, "utf8")) as CachedModels;
+    return Date.now() - cached.fetchedAt < CACHE_TTL_MS;
+  } catch {
+    return false;
+  }
+}
+
+function writeCache(data: ApiModel[]): void {
+  try {
+    mkdirSync(CACHE_DIR, { recursive: true });
+    const cached: CachedModels = { fetchedAt: Date.now(), data };
+    writeFileSync(CACHE_FILE, JSON.stringify(cached));
+  } catch (err) {
+    log(`cache write failed: ${(err as Error).message}`);
+  }
+}
+
+export default async function (pi: ExtensionAPI): Promise<void> {
+  const apiKey = process.env.WANDB_API_KEY;
+  const project = process.env.WANDB_PROJECT;
+  const skipCache = process.env.WANDB_NO_CACHE === "1";
+
+  if (!apiKey) {
+    log("WANDB_API_KEY not set; provider not registered. See README.");
+    return;
+  }
+
+  let models: ApiModel[] | null = null;
+  if (!skipCache && isCacheFresh()) {
+    models = readCache();
+    if (models) log(`using fresh cached model list (${models.length} models)`);
+  }
+
+  if (!models) {
+    try {
+      models = await fetchModels(apiKey, project);
+      writeCache(models);
+      log(`fetched ${models.length} models from /v1/models`);
+    } catch (err) {
+      log(`fetch failed: ${(err as Error).message}`);
+      models = readCache();
+      if (models) {
+        log(`using stale cached model list (${models.length} models)`);
+      } else {
+        log("no cached models available; provider not registered.");
+        return;
+      }
+    }
+  }
+
+  const headers: Record<string, string> = {};
+  if (project) headers["OpenAI-Project"] = project;
+
+  pi.registerProvider("wandb", {
+    name: "Weights & Biases Inference",
+    baseUrl: BASE_URL,
+    apiKey: "WANDB_API_KEY",
+    api: "openai-completions",
+    headers,
+    models: models.map((m) => ({
+      id: m.id,
+      name: m.name ?? prettyName(m.id),
+      reasoning: isReasoning(m.id),
+      input: ["text"],
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+      contextWindow: KNOWN_CONTEXT_WINDOWS[m.id] ?? DEFAULT_CONTEXT_WINDOW,
+      maxTokens: DEFAULT_MAX_TOKENS,
+    })),
+  });
+}

package/package.json

@@ -0,0 +1,52 @@
+{
+  "name": "pi-extension-wandb",
+  "version": "0.0.0",
+  "description": "pi coding agent extension that adds Weights & Biases Inference as a model provider.",
+  "type": "module",
+  "license": "MIT",
+  "engines": {
+    "node": ">=18"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "author": "Kiran Gadhave",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/kirangadhave/pi-extension-wandb.git"
+  },
+  "homepage": "https://github.com/kirangadhave/pi-extension-wandb",
+  "bugs": {
+    "url": "https://github.com/kirangadhave/pi-extension-wandb/issues"
+  },
+  "keywords": [
+    "pi",
+    "pi-extension",
+    "pi-coding-agent",
+    "wandb",
+    "weights-and-biases",
+    "inference",
+    "llm",
+    "provider"
+  ],
+  "files": [
+    "index.ts",
+    "LICENSE",
+    "README.md"
+  ],
+  "scripts": {
+    "clean": "echo 'nothing to clean'",
+    "build": "echo 'nothing to build'",
+    "check": "tsc --noEmit"
+  },
+  "pi": {
+    "extensions": [
+      "./index.ts"
+    ]
+  },
+  "devDependencies": {
+    "@earendil-works/pi-coding-agent": "^0.75.0",
+    "@types/node": "^22.0.0",
+    "typescript": "^5.5.0"
+  }
+}