|
/** |
|
* Merge Gateway provider for pi. |
|
* |
|
* Merge Gateway (https://gateway.merge.dev) routes every major LLM through one |
|
* OpenAI-compatible endpoint. This extension discovers the live model catalog from |
|
* GET /v1/models and registers it as a pi provider, so the picker always reflects |
|
* what Gateway currently serves (no hardcoded model list to go stale). |
|
* |
|
* Usage: |
|
* export MERGE_GATEWAY_API_KEY=mg_... |
|
* Then pick a model via /model (filter on "merge-gateway"). The special |
|
* "default_routing" entry hands requests to your Gateway routing policy. |
|
*/ |
|
|
|
import { homedir } from "node:os"; |
|
import { join } from "node:path"; |
|
import { readFileSync } from "node:fs"; |
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; |
|
|
|
const BASE_URL = "https://api-gateway.merge.dev/v1/openai"; |
|
const CATALOG_URL = "https://api-gateway.merge.dev/v1/models"; |
|
const API_KEY_ENV = "MERGE_GATEWAY_API_KEY"; |
|
|
|
// Resolve the key from the environment, falling back to pi's credential store |
|
// (~/.pi/agent/auth.json, populated by /login). The provider's apiKey in |
|
// registerProvider stays "$MERGE_GATEWAY_API_KEY" so env-set keys work too; this |
|
// only covers the catalog fetch, which needs a concrete key at load time. |
|
function resolveApiKey(): string | undefined { |
|
const env = process.env[API_KEY_ENV]; |
|
if (env) return env; |
|
try { |
|
const auth = JSON.parse( |
|
readFileSync(join(homedir(), ".pi", "agent", "auth.json"), "utf8"), |
|
) as Record<string, { key?: string }>; |
|
return auth["merge-gateway"]?.key; |
|
} catch { |
|
return undefined; |
|
} |
|
} |
|
|
|
type Vendor = { |
|
context_window?: number; |
|
max_output_tokens?: number; |
|
capabilities?: { |
|
input?: string[]; |
|
output?: string[]; |
|
supports_reasoning?: boolean; |
|
}; |
|
pricing?: { |
|
input_per_million?: number; |
|
output_per_million?: number; |
|
cache_read_per_million?: number; |
|
cache_write_per_million?: number; |
|
}; |
|
}; |
|
|
|
type CatalogModel = { |
|
model: string; |
|
display_name?: string; |
|
vendors?: Record<string, Vendor>; |
|
}; |
|
|
|
type CatalogResponse = { |
|
data?: CatalogModel[]; |
|
has_more?: boolean; |
|
next_cursor?: string; |
|
}; |
|
|
|
type PiModel = { |
|
id: string; |
|
name: string; |
|
reasoning: boolean; |
|
input: ("text" | "image")[]; |
|
contextWindow: number; |
|
maxTokens: number; |
|
cost: { input: number; output: number; cacheRead: number; cacheWrite: number }; |
|
}; |
|
|
|
const DEFAULT_ROUTING: PiModel = { |
|
id: "default_routing", |
|
name: "Merge Gateway (policy routing)", |
|
reasoning: true, |
|
input: ["text", "image"], |
|
contextWindow: 128000, |
|
maxTokens: 8192, |
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, |
|
}; |
|
|
|
function toPiModel(m: CatalogModel): PiModel | null { |
|
const v = m.vendors && Object.values(m.vendors)[0]; |
|
if (!v) return null; |
|
const caps = v.capabilities ?? {}; |
|
const outModes = caps.output ?? []; |
|
// Keep only chat/completion models (skip image/audio/video gen, embeddings, tts). |
|
if (!outModes.includes("text") && !outModes.includes("tool_use")) return null; |
|
const input = (caps.input ?? []).filter( |
|
(x): x is "text" | "image" => x === "text" || x === "image", |
|
); |
|
const pr = v.pricing ?? {}; |
|
return { |
|
id: m.model, |
|
name: m.display_name || m.model, |
|
reasoning: Boolean(caps.supports_reasoning), |
|
input: input.length ? input : ["text"], |
|
contextWindow: v.context_window ?? 128000, |
|
maxTokens: v.max_output_tokens ?? 8192, |
|
cost: { |
|
input: pr.input_per_million ?? 0, |
|
output: pr.output_per_million ?? 0, |
|
cacheRead: pr.cache_read_per_million ?? 0, |
|
cacheWrite: pr.cache_write_per_million ?? 0, |
|
}, |
|
}; |
|
} |
|
|
|
async function fetchCatalog(apiKey: string, signal?: AbortSignal): Promise<PiModel[]> { |
|
const models: PiModel[] = []; |
|
let cursor: string | undefined; |
|
do { |
|
const url = new URL(CATALOG_URL); |
|
if (cursor) url.searchParams.set("cursor", cursor); |
|
const res = await fetch(url, { |
|
headers: { Authorization: `Bearer ${apiKey}` }, |
|
signal: signal ?? AbortSignal.timeout(8000), |
|
}); |
|
if (!res.ok) throw new Error(`catalog status ${res.status}`); |
|
const payload = (await res.json()) as CatalogResponse; |
|
for (const m of payload.data ?? []) { |
|
const pm = toPiModel(m); |
|
if (pm) models.push(pm); |
|
} |
|
cursor = payload.has_more ? payload.next_cursor : undefined; |
|
} while (cursor); |
|
return models; |
|
} |
|
|
|
export default function (pi: ExtensionAPI) { |
|
const apiKey = resolveApiKey(); |
|
|
|
pi.registerProvider("merge-gateway", { |
|
name: "Merge Gateway", |
|
baseUrl: BASE_URL, |
|
// omp (oh-my-pi) does not interpolate $ENV syntax in registerProvider; pass |
|
// the concrete key so both pi and omp work. Falls back to pi's auth store. |
|
apiKey: apiKey ?? `$${API_KEY_ENV}`, |
|
api: "openai-completions", |
|
models: [DEFAULT_ROUTING], |
|
// Live catalog discovery happens on model refresh, never during the factory — |
|
// an awaited fetch here blocked pi startup by ~1.6s. |
|
async refreshModels({ signal }: { signal: AbortSignal }) { |
|
if (!apiKey) return [DEFAULT_ROUTING]; |
|
const discovered = await fetchCatalog(apiKey, signal); |
|
return discovered.length ? [...discovered, DEFAULT_ROUTING] : [DEFAULT_ROUTING]; |
|
}, |
|
}); |
|
|
|
// Budget display: Gateway echoes credit balance / key limit headers on every |
|
// response. Show remaining balance in the footer status after each turn. |
|
pi.on("after_provider_response", (event, ctx) => { |
|
const h = event.headers as Record<string, string | string[] | undefined>; |
|
if (!h) return; |
|
const pick = (...names: string[]) => { |
|
for (const n of names) { |
|
const v = h[n] ?? h[n.toLowerCase()]; |
|
if (v != null) return Array.isArray(v) ? v[0] : v; |
|
} |
|
return undefined; |
|
}; |
|
const num = (v?: string) => |
|
v == null ? undefined : Number(v).toFixed(2).replace(/\.?0+$/, ""); |
|
const balance = num(pick("x-credit-balance-usd")); |
|
const keyLeft = num(pick("x-key-limit-remaining-usd")); |
|
const projLeft = num(pick("x-project-budget-remaining-usd")); |
|
const warn = pick("x-budget-warning"); |
|
const parts: string[] = []; |
|
if (balance != null) parts.push(`budget: $${balance}`); |
|
if (keyLeft != null) parts.push(`api key: $${keyLeft}`); |
|
if (projLeft != null) parts.push(`proj left: $${projLeft}`); |
|
if (warn) parts.push(`⚠ ${warn}`); |
|
if (parts.length) ctx.ui.setStatus("merge-gateway", parts.join(" · ")); |
|
}); |
|
} |
