pi-bifrost 0.1.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/CONTEXT.md +56 -0
- package/README.md +170 -0
- package/bifrost-routes.json.example +10 -0
- package/bifrost.json +45 -0
- package/cache.ts +171 -0
- package/classification-pipeline.ts +107 -0
- package/classifier.ts +282 -0
- package/commands.ts +565 -0
- package/config.ts +120 -0
- package/debug.ts +169 -0
- package/docs/adr/0001-dual-classifier.md +14 -0
- package/docs/adr/0002-subprocess-classifier.md +15 -0
- package/index.ts +259 -0
- package/package.json +31 -0
- package/probe.ts +144 -0
- package/routing.ts +173 -0
- package/schema.json +189 -0
- package/tests/cache.test.ts +148 -0
- package/tests/classifier.test.ts +48 -0
- package/tests/commands.test.ts +116 -0
- package/tests/integration.test.mjs +157 -0
- package/tests/pipeline.test.ts +269 -0
- package/tests/routing.test.ts +224 -0
- package/tsconfig.json +19 -0
package/debug.ts
ADDED
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
// ── Structured debug logging with Performance API ──────────────────
|
|
2
|
+
// Buffered async writes. No sync I/O on the hot path — events are
|
|
3
|
+
// queued in memory and flushed via setImmediate. On process exit the
|
|
4
|
+
// remaining buffer is flushed synchronously.
|
|
5
|
+
//
|
|
6
|
+
// Configure: { "debug": { "enabled": true, "path": ".pi/bifrost-debug.jsonl" } }
|
|
7
|
+
|
|
8
|
+
import { existsSync, mkdirSync } from "node:fs";
|
|
9
|
+
import { appendFile, rename, stat } from "node:fs/promises";
|
|
10
|
+
import { dirname } from "node:path";
|
|
11
|
+
import { performance } from "node:perf_hooks";
|
|
12
|
+
import { appendFileSync } from "node:fs";
|
|
13
|
+
|
|
14
|
+
export interface DebugConfig {
|
|
15
|
+
enabled?: boolean;
|
|
16
|
+
path?: string;
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
const DEFAULT_MAX_SIZE_MB = 10;
|
|
20
|
+
|
|
21
|
+
let debugEnabled = false;
|
|
22
|
+
let debugPath: string | null = null;
|
|
23
|
+
const maxSizeBytes: number = DEFAULT_MAX_SIZE_MB * 1024 * 1024;
|
|
24
|
+
let startupDone = false;
|
|
25
|
+
|
|
26
|
+
// ── Async write buffer ────────────────────────────────────────────
|
|
27
|
+
|
|
28
|
+
let buffer: string[] = [];
|
|
29
|
+
let flushScheduled = false;
|
|
30
|
+
let flushing = false;
|
|
31
|
+
|
|
32
|
+
async function rotateIfNeeded(): Promise<void> {
|
|
33
|
+
if (!debugPath) return;
|
|
34
|
+
try {
|
|
35
|
+
const info = await stat(debugPath).catch(() => null);
|
|
36
|
+
if (info && info.size > maxSizeBytes) {
|
|
37
|
+
await rename(debugPath, debugPath.replace(/\.jsonl$/, ".old.jsonl")).catch(() => {});
|
|
38
|
+
}
|
|
39
|
+
} catch { /* ignore */ }
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
async function flushBuffer(): Promise<void> {
|
|
43
|
+
// Write lock: only one flush runs at a time. Prevents interleaved
|
|
44
|
+
// appendFile calls that would corrupt JSONL.
|
|
45
|
+
if (flushing) return;
|
|
46
|
+
flushing = true;
|
|
47
|
+
flushScheduled = false;
|
|
48
|
+
|
|
49
|
+
while (buffer.length > 0 && debugPath) {
|
|
50
|
+
const lines = buffer.join("\n") + "\n";
|
|
51
|
+
buffer = [];
|
|
52
|
+
try {
|
|
53
|
+
await rotateIfNeeded();
|
|
54
|
+
await appendFile(debugPath, lines, "utf-8");
|
|
55
|
+
} catch (err) {
|
|
56
|
+
try {
|
|
57
|
+
process.stderr.write(`[bifrost] debug write failed: ${err}\n`);
|
|
58
|
+
} catch { /* silence */ }
|
|
59
|
+
}
|
|
60
|
+
}
|
|
61
|
+
|
|
62
|
+
flushing = false;
|
|
63
|
+
// New entries may have arrived during the await.
|
|
64
|
+
if (buffer.length > 0) scheduleFlush();
|
|
65
|
+
}
|
|
66
|
+
|
|
67
|
+
function scheduleFlush(): void {
|
|
68
|
+
if (!flushScheduled && !flushing) {
|
|
69
|
+
flushScheduled = true;
|
|
70
|
+
setImmediate(flushBuffer);
|
|
71
|
+
}
|
|
72
|
+
// If already flushing, the running flush will drain the buffer
|
|
73
|
+
// and re-schedule if more entries arrive.
|
|
74
|
+
}
|
|
75
|
+
|
|
76
|
+
/** Synchronous flush for process.exit (must be sync). */
|
|
77
|
+
function flushSync(): void {
|
|
78
|
+
if (buffer.length === 0 || !debugPath) return;
|
|
79
|
+
const lines = buffer.join("\n") + "\n";
|
|
80
|
+
buffer = [];
|
|
81
|
+
try {
|
|
82
|
+
appendFileSync(debugPath, lines, "utf-8");
|
|
83
|
+
} catch { /* silence */ }
|
|
84
|
+
}
|
|
85
|
+
|
|
86
|
+
// ── Setup ─────────────────────────────────────────────────────────
|
|
87
|
+
|
|
88
|
+
export function setupDebug(cfg: DebugConfig, cwd: string) {
|
|
89
|
+
debugEnabled = cfg.enabled ?? false;
|
|
90
|
+
if (cfg.path) {
|
|
91
|
+
debugPath = cfg.path.startsWith("/") || cfg.path.startsWith("~")
|
|
92
|
+
? cfg.path.replace(/^~/, process.env.HOME ?? "/tmp")
|
|
93
|
+
: `${cwd}/${cfg.path}`;
|
|
94
|
+
} else {
|
|
95
|
+
debugPath = `${cwd}/.pi/bifrost-debug.jsonl`;
|
|
96
|
+
}
|
|
97
|
+
|
|
98
|
+
if (debugEnabled && !startupDone) {
|
|
99
|
+
startupDone = true;
|
|
100
|
+
try {
|
|
101
|
+
const dir = dirname(debugPath);
|
|
102
|
+
if (!existsSync(dir)) mkdirSync(dir, { recursive: true });
|
|
103
|
+
} catch (err) {
|
|
104
|
+
process.stderr.write(`[bifrost] debug: cannot create log dir: ${err}\n`);
|
|
105
|
+
debugEnabled = false;
|
|
106
|
+
return;
|
|
107
|
+
}
|
|
108
|
+
|
|
109
|
+
// Flush remaining on exit (sync — exit hook cannot do async I/O).
|
|
110
|
+
process.on("exit", () => {
|
|
111
|
+
buffer.push(JSON.stringify({ ts: new Date().toISOString(), module: "bifrost", event: "shutdown", entryType: "event" }));
|
|
112
|
+
flushSync();
|
|
113
|
+
});
|
|
114
|
+
}
|
|
115
|
+
}
|
|
116
|
+
|
|
117
|
+
// ── Public API ────────────────────────────────────────────────────
|
|
118
|
+
|
|
119
|
+
/** Log a discrete event. Async — queued and flushed via setImmediate. */
|
|
120
|
+
export function debug(
|
|
121
|
+
module: string,
|
|
122
|
+
event: string,
|
|
123
|
+
meta?: Record<string, unknown>,
|
|
124
|
+
) {
|
|
125
|
+
if (!debugEnabled) return;
|
|
126
|
+
try {
|
|
127
|
+
buffer.push(JSON.stringify({ ts: new Date().toISOString(), module, event, entryType: "event", ...meta }));
|
|
128
|
+
} catch {
|
|
129
|
+
buffer.push(JSON.stringify({ ts: new Date().toISOString(), module, event, entryType: "event", _meta_error: "unserializable" }));
|
|
130
|
+
}
|
|
131
|
+
scheduleFlush();
|
|
132
|
+
}
|
|
133
|
+
|
|
134
|
+
/**
|
|
135
|
+
* Start a performance measure. Returns a stop function that records
|
|
136
|
+
* a measure entry via the Performance API. Async — queued like debug().
|
|
137
|
+
*/
|
|
138
|
+
export function debugMeasure(module: string, event: string) {
|
|
139
|
+
if (!debugEnabled) return (_meta?: Record<string, unknown>) => {};
|
|
140
|
+
const name = `bifrost:${module}:${event}`;
|
|
141
|
+
performance.mark(`${name}:start`);
|
|
142
|
+
return (meta?: Record<string, unknown>) => {
|
|
143
|
+
try {
|
|
144
|
+
performance.mark(`${name}:end`);
|
|
145
|
+
performance.measure(name, `${name}:start`, `${name}:end`);
|
|
146
|
+
const entry = performance.getEntriesByName(name, "measure")[0];
|
|
147
|
+
if (entry) {
|
|
148
|
+
const data: Record<string, unknown> = {
|
|
149
|
+
ts: new Date().toISOString(),
|
|
150
|
+
module,
|
|
151
|
+
event,
|
|
152
|
+
entryType: "measure",
|
|
153
|
+
duration_ms: +entry.duration.toFixed(3),
|
|
154
|
+
...meta,
|
|
155
|
+
};
|
|
156
|
+
try {
|
|
157
|
+
buffer.push(JSON.stringify(data));
|
|
158
|
+
} catch {
|
|
159
|
+
buffer.push(JSON.stringify({ ...data, _meta_error: "unserializable" }));
|
|
160
|
+
}
|
|
161
|
+
}
|
|
162
|
+
} finally {
|
|
163
|
+
performance.clearMarks(`${name}:start`);
|
|
164
|
+
performance.clearMarks(`${name}:end`);
|
|
165
|
+
performance.clearMeasures(name);
|
|
166
|
+
}
|
|
167
|
+
scheduleFlush();
|
|
168
|
+
};
|
|
169
|
+
}
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
# Dual classifier: LLM-first with regex fallback
|
|
2
|
+
|
|
3
|
+
pi-bifrost uses two classification systems: an LLM classifier (tried first) and regex rules (fallback). The LLM reads the prompt and returns a tier name; regex rules are simple pattern→tier mappings.
|
|
4
|
+
|
|
5
|
+
**Why not LLM-only?** The classifier model can be unavailable (local server down, API outage, rate limit). A silent model router that breaks when a sidecar is down is worse than no router at all. The regex fallback ensures routing always works.
|
|
6
|
+
|
|
7
|
+
**Why not regex-only?** Regex rules are rigid. "Help me design a rate limiter" and "what's a rate limiter?" both contain "rate limiter" but need different tiers. An LLM can distinguish intent — regex cannot.
|
|
8
|
+
|
|
9
|
+
**Why LLM first, not regex first?** The LLM is more accurate. If it's available, use it. Regex is the safety net, not the primary path. Running regex first and only falling back to LLM would mean paying for LLM tokens only when regex fails — but regex never truly "fails," it just matches or doesn't. The current order means: try the accurate thing, fall back to the reliable thing.
|
|
10
|
+
|
|
11
|
+
**Consequences:**
|
|
12
|
+
- Every cache miss costs LLM tokens (mitigated by using cheap classifier models and a fuzzy cache)
|
|
13
|
+
- The classifier model must understand the tier names — custom tiers need a custom system prompt
|
|
14
|
+
- Two classification paths to test and maintain
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
# Subprocess classifier invocation alongside direct HTTP
|
|
2
|
+
|
|
3
|
+
The LLM classifier supports two invocation methods: `direct` (OpenAI-compatible HTTP POST to the model's base URL) and `subprocess` (spawns a fresh `pi --no-extensions` process).
|
|
4
|
+
|
|
5
|
+
**Why two methods?** Direct HTTP is fast but only works for OpenAI-compatible APIs. Many pi providers use non-OpenAI protocols (Anthropic Messages, Codex Responses) or require custom auth flows that pi handles internally. The subprocess method delegates to pi itself — any model pi can use, the classifier can use.
|
|
6
|
+
|
|
7
|
+
**Why not subprocess-only?** Subprocess spawn is slow — cold start, config loading, model resolution. Direct HTTP is a single fetch. For local LM Studio/Ollama servers (the common classifier case), direct HTTP is near-instant.
|
|
8
|
+
|
|
9
|
+
**Why not direct-only?** Would lock the classifier out of many providers. A user whose only cheap model is an Anthropic Haiku key couldn't use the classifier.
|
|
10
|
+
|
|
11
|
+
**Consequences:**
|
|
12
|
+
- Two code paths to maintain in `classifier.ts` (direct HTTP ~50 lines, subprocess ~50 lines)
|
|
13
|
+
- Subprocess method is a recursive pi invocation — pi-bifrost loaded in the subprocess must not interfere (handled via `--no-extensions` in the subprocess)
|
|
14
|
+
- `auto` method tries direct first, then subprocess — doubles latency on failure but maximizes availability
|
|
15
|
+
- The `endpoint` config field forces `direct` mode since there's no pi process to delegate to for arbitrary endpoints
|
package/index.ts
ADDED
|
@@ -0,0 +1,259 @@
|
|
|
1
|
+
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
2
|
+
import { fileURLToPath } from "node:url";
|
|
3
|
+
import { classifyWithLLM as invokeClassifier, type ClassifierModel } from "./classifier.js";
|
|
4
|
+
import {
|
|
5
|
+
createPipeline,
|
|
6
|
+
type ClassificationPipeline,
|
|
7
|
+
} from "./classification-pipeline.js";
|
|
8
|
+
import {
|
|
9
|
+
cachePath,
|
|
10
|
+
lookupCache,
|
|
11
|
+
touchCacheEntry,
|
|
12
|
+
loadCache,
|
|
13
|
+
saveCache,
|
|
14
|
+
updateCache,
|
|
15
|
+
DEFAULT_MAX_ENTRIES,
|
|
16
|
+
DEFAULT_THRESHOLD,
|
|
17
|
+
type CacheEntry,
|
|
18
|
+
} from "./cache.js";
|
|
19
|
+
import {
|
|
20
|
+
loadConfig,
|
|
21
|
+
loadRules,
|
|
22
|
+
type BifrostConfig,
|
|
23
|
+
} from "./config.js";
|
|
24
|
+
import {
|
|
25
|
+
findCandidates,
|
|
26
|
+
getStrategy,
|
|
27
|
+
modelKey,
|
|
28
|
+
resolveModel,
|
|
29
|
+
} from "./routing.js";
|
|
30
|
+
import { createCommandRouter, log, uiBusy, uiDone, type BifrostState } from "./commands.js";
|
|
31
|
+
import { setupDebug, debug, debugMeasure } from "./debug.js";
|
|
32
|
+
|
|
33
|
+
// ── Pipeline builder (composition root) ────────────────────────
|
|
34
|
+
|
|
35
|
+
function resolveClassifierModels(
|
|
36
|
+
ctx: ExtensionContext,
|
|
37
|
+
config: BifrostConfig,
|
|
38
|
+
): ClassifierModel[] {
|
|
39
|
+
const pattern = config.classifier?.model;
|
|
40
|
+
if (!pattern) return [];
|
|
41
|
+
return findCandidates(ctx, pattern)
|
|
42
|
+
.slice(0, 3)
|
|
43
|
+
.map((model) => ({ kind: "registry" as const, model }));
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
function endpointClassifier(id: string, endpoint: string): ClassifierModel {
|
|
47
|
+
return { kind: "endpoint", id, baseUrl: endpoint };
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
function buildPipeline(
|
|
51
|
+
ctx: ExtensionContext,
|
|
52
|
+
config: BifrostConfig,
|
|
53
|
+
cacheEntries: CacheEntry[],
|
|
54
|
+
classifierEnabled: boolean,
|
|
55
|
+
): ClassificationPipeline {
|
|
56
|
+
const tiers = Object.keys(config.models ?? {});
|
|
57
|
+
const cacheCfg = config.cache;
|
|
58
|
+
const cacheEnabled = cacheCfg?.enabled ?? true;
|
|
59
|
+
const threshold = cacheCfg?.threshold ?? DEFAULT_THRESHOLD;
|
|
60
|
+
|
|
61
|
+
// Resolve classifier models once at pipeline construction.
|
|
62
|
+
// If classifier is disabled, pass empty array — pipeline skips LLM stage.
|
|
63
|
+
let classifierModels: ClassifierModel[] = [];
|
|
64
|
+
if (classifierEnabled && tiers.length > 0) {
|
|
65
|
+
const classifierEndpoint = config.classifier?.endpoint;
|
|
66
|
+
if (classifierEndpoint) {
|
|
67
|
+
const rawModel = config.classifier?.model;
|
|
68
|
+
const modelId = Array.isArray(rawModel)
|
|
69
|
+
? rawModel[0]
|
|
70
|
+
: (rawModel ?? "classifier");
|
|
71
|
+
classifierModels = [endpointClassifier(modelId, classifierEndpoint)];
|
|
72
|
+
} else {
|
|
73
|
+
classifierModels = resolveClassifierModels(ctx, config);
|
|
74
|
+
}
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
return createPipeline({
|
|
78
|
+
cacheLookup: (text) => {
|
|
79
|
+
if (!cacheEnabled) return undefined;
|
|
80
|
+
const entry = lookupCache(cacheEntries, text, threshold);
|
|
81
|
+
if (entry) {
|
|
82
|
+
touchCacheEntry(entry);
|
|
83
|
+
return entry.category;
|
|
84
|
+
}
|
|
85
|
+
return undefined;
|
|
86
|
+
},
|
|
87
|
+
classifierModels,
|
|
88
|
+
classifyWithLLM: (model, text, tiers) =>
|
|
89
|
+
invokeClassifier(ctx, model, tiers, text, {
|
|
90
|
+
systemPrompt: config.classifier?.systemPrompt,
|
|
91
|
+
maxTokens: config.classifier?.maxTokens,
|
|
92
|
+
temperature: config.classifier?.temperature,
|
|
93
|
+
method: config.classifier?.method,
|
|
94
|
+
}),
|
|
95
|
+
regexRules: loadRules(process.cwd(), config),
|
|
96
|
+
defaultTier: config.default,
|
|
97
|
+
tiers,
|
|
98
|
+
});
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
export default function bifrostExtension(pi: ExtensionAPI) {
|
|
102
|
+
const extensionDir = fileURLToPath(new URL(".", import.meta.url));
|
|
103
|
+
|
|
104
|
+
// Setup debug logging first — so startup errors are captured.
|
|
105
|
+
const bootConfig = loadConfig(process.cwd(), extensionDir);
|
|
106
|
+
if (bootConfig.debug?.enabled) {
|
|
107
|
+
setupDebug(bootConfig.debug, process.cwd());
|
|
108
|
+
debug("bifrost", "startup", { extensionDir });
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
const config = bootConfig;
|
|
112
|
+
const cacheEntries = loadCache(cachePath(process.cwd(), config.cache?.path));
|
|
113
|
+
let selfSelecting = false;
|
|
114
|
+
let pipeline: ClassificationPipeline | undefined;
|
|
115
|
+
|
|
116
|
+
function getPipeline(ctx: ExtensionContext): ClassificationPipeline {
|
|
117
|
+
if (!pipeline) {
|
|
118
|
+
pipeline = buildPipeline(ctx, state.config, state.cacheEntries, state.classifierEnabled);
|
|
119
|
+
}
|
|
120
|
+
return pipeline;
|
|
121
|
+
}
|
|
122
|
+
|
|
123
|
+
function invalidatePipeline() {
|
|
124
|
+
debug("bifrost", "pipeline.invalidate");
|
|
125
|
+
pipeline = undefined;
|
|
126
|
+
}
|
|
127
|
+
|
|
128
|
+
// Mutable state shared with command handlers.
|
|
129
|
+
const state: BifrostState = {
|
|
130
|
+
config,
|
|
131
|
+
enabled: config.enabled ?? true,
|
|
132
|
+
classifierEnabled: config.classifier?.enabled ?? true,
|
|
133
|
+
pinned: false,
|
|
134
|
+
cacheEntries,
|
|
135
|
+
extensionDir,
|
|
136
|
+
getPipeline,
|
|
137
|
+
invalidatePipeline,
|
|
138
|
+
};
|
|
139
|
+
|
|
140
|
+
const handleCommand = createCommandRouter(state);
|
|
141
|
+
|
|
142
|
+
pi.registerCommand("bifrost", {
|
|
143
|
+
description: "Bifrost model router control",
|
|
144
|
+
handler: async (args, ctx) => {
|
|
145
|
+
await handleCommand(args, ctx);
|
|
146
|
+
},
|
|
147
|
+
});
|
|
148
|
+
|
|
149
|
+
pi.on("model_select", async (_event, ctx) => {
|
|
150
|
+
if (selfSelecting) {
|
|
151
|
+
selfSelecting = false;
|
|
152
|
+
return;
|
|
153
|
+
}
|
|
154
|
+
if (!state.enabled) return;
|
|
155
|
+
|
|
156
|
+
state.pinned = true;
|
|
157
|
+
debug("bifrost", "model_select", { model: modelKey(ctx.model) });
|
|
158
|
+
log(
|
|
159
|
+
ctx,
|
|
160
|
+
`Model manually changed to ${modelKey(ctx.model)}; Bifrost pinned.`,
|
|
161
|
+
);
|
|
162
|
+
});
|
|
163
|
+
|
|
164
|
+
pi.on("input", async (event, ctx) => {
|
|
165
|
+
if (event.source === "extension") return { action: "continue" };
|
|
166
|
+
if (!state.enabled || state.pinned) {
|
|
167
|
+
debug("input", "bypass", { enabled: state.enabled, pinned: state.pinned });
|
|
168
|
+
return { action: "continue" };
|
|
169
|
+
}
|
|
170
|
+
|
|
171
|
+
const text = event.text.trim();
|
|
172
|
+
if (text.startsWith("/")) return { action: "continue" };
|
|
173
|
+
|
|
174
|
+
const endInput = debugMeasure("input", "total");
|
|
175
|
+
debug("input", "prompt", { length: text.length });
|
|
176
|
+
|
|
177
|
+
if (state.classifierEnabled) {
|
|
178
|
+
const endRefresh = debugMeasure("input", "registry.refresh");
|
|
179
|
+
try {
|
|
180
|
+
await ctx.modelRegistry.refresh();
|
|
181
|
+
invalidatePipeline();
|
|
182
|
+
endRefresh();
|
|
183
|
+
} catch (err) {
|
|
184
|
+
debug("input", "registry.refresh.error", { error: String(err) });
|
|
185
|
+
console.error(`[bifrost] model registry refresh failed: ${err}`);
|
|
186
|
+
}
|
|
187
|
+
}
|
|
188
|
+
|
|
189
|
+
uiBusy(ctx, "Bifrost classifying...");
|
|
190
|
+
const endClassify = debugMeasure("input", "classify");
|
|
191
|
+
const classification = await getPipeline(ctx).classify(text);
|
|
192
|
+
endClassify({ kind: classification.kind, tier: classification.kind !== "unclassified" ? classification.tier : undefined });
|
|
193
|
+
uiDone(ctx);
|
|
194
|
+
|
|
195
|
+
if (classification.kind === "unclassified") {
|
|
196
|
+
log(ctx, "Bifrost: no tier matched — using default model", "warning");
|
|
197
|
+
debug("input", "unclassified");
|
|
198
|
+
endInput();
|
|
199
|
+
return { action: "continue" };
|
|
200
|
+
}
|
|
201
|
+
|
|
202
|
+
const tier = classification.tier;
|
|
203
|
+
const source = classification.kind === "classified"
|
|
204
|
+
? classification.source
|
|
205
|
+
: "fallback";
|
|
206
|
+
const pattern = state.config.models?.[tier] ?? tier;
|
|
207
|
+
const strategy = getStrategy(state.config.categoryStrategies, state.config.strategy, tier);
|
|
208
|
+
const model = resolveModel(ctx, pattern, strategy);
|
|
209
|
+
if (!model) {
|
|
210
|
+
debug("input", "no_model", { tier });
|
|
211
|
+
log(ctx, `Bifrost: tier "${tier}" matched but no model available`, "warning");
|
|
212
|
+
endInput();
|
|
213
|
+
return { action: "continue" };
|
|
214
|
+
}
|
|
215
|
+
|
|
216
|
+
// Only cache LLM classifier results (regex is deterministic).
|
|
217
|
+
if (
|
|
218
|
+
classification.kind === "classified" &&
|
|
219
|
+
classification.source === "classifier"
|
|
220
|
+
) {
|
|
221
|
+
const maxEntries = state.config.cache?.maxEntries ?? DEFAULT_MAX_ENTRIES;
|
|
222
|
+
if (state.config.cache?.enabled ?? true) {
|
|
223
|
+
const endCacheSave = debugMeasure("input", "cacheSave");
|
|
224
|
+
state.cacheEntries = updateCache(state.cacheEntries, text, tier, maxEntries);
|
|
225
|
+
saveCache(cachePath(process.cwd(), state.config.cache?.path), state.cacheEntries);
|
|
226
|
+
invalidatePipeline();
|
|
227
|
+
endCacheSave({ entries: state.cacheEntries.length });
|
|
228
|
+
}
|
|
229
|
+
}
|
|
230
|
+
|
|
231
|
+
if (modelKey(model) === modelKey(ctx.model)) {
|
|
232
|
+
uiDone(ctx);
|
|
233
|
+
log(ctx, `Bifrost: ${tier} → ${modelKey(model)} (already active, ${source})`);
|
|
234
|
+
debug("input", "model_unchanged", { model: modelKey(model) });
|
|
235
|
+
endInput({ model: modelKey(model), tier, strategy, source });
|
|
236
|
+
return { action: "continue" };
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
uiBusy(ctx, `Bifrost routing to ${modelKey(model)}...`);
|
|
240
|
+
selfSelecting = true;
|
|
241
|
+
const endSwitch = debugMeasure("input", "setModel");
|
|
242
|
+
const ok = await pi.setModel(model);
|
|
243
|
+
endSwitch({ model: modelKey(model), ok });
|
|
244
|
+
uiDone(ctx);
|
|
245
|
+
if (!ok) {
|
|
246
|
+
selfSelecting = false;
|
|
247
|
+
log(ctx, `Bifrost: no API key for ${modelKey(model)}`, "error");
|
|
248
|
+
endInput({ model: modelKey(model), ok: false });
|
|
249
|
+
return { action: "continue" };
|
|
250
|
+
}
|
|
251
|
+
|
|
252
|
+
const doneMsg = classification.kind === "classified"
|
|
253
|
+
? `Bifrost: ${tier} → ${modelKey(model)} (${classification.source})`
|
|
254
|
+
: `Bifrost: ${tier} → ${modelKey(model)} (fallback)`;
|
|
255
|
+
log(ctx, doneMsg);
|
|
256
|
+
endInput({ model: modelKey(model), tier, strategy, source });
|
|
257
|
+
return { action: "continue" };
|
|
258
|
+
});
|
|
259
|
+
}
|
package/package.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
{
|
|
2
|
+
"name": "pi-bifrost",
|
|
3
|
+
"version": "0.1.0",
|
|
4
|
+
"description": "query-aware model router for the pi",
|
|
5
|
+
"type": "module",
|
|
6
|
+
"keywords": [
|
|
7
|
+
"pi-package",
|
|
8
|
+
"pi-extension",
|
|
9
|
+
"model-router"
|
|
10
|
+
],
|
|
11
|
+
"license": "MIT",
|
|
12
|
+
"pi": {
|
|
13
|
+
"extensions": [
|
|
14
|
+
"./index.ts"
|
|
15
|
+
],
|
|
16
|
+
"minPiVersion": "0.80.8"
|
|
17
|
+
},
|
|
18
|
+
"scripts": {
|
|
19
|
+
"test": "node --test --experimental-strip-types tests/cache.test.ts tests/classifier.test.ts tests/routing.test.ts tests/pipeline.test.ts tests/commands.test.ts",
|
|
20
|
+
"test:integration": "cd .. && node --test pi-bifrost/tests/integration.test.mjs",
|
|
21
|
+
"typecheck": "tsc --noEmit",
|
|
22
|
+
"lint": "tsc --noEmit"
|
|
23
|
+
},
|
|
24
|
+
"peerDependencies": {
|
|
25
|
+
"@earendil-works/pi-ai": "*",
|
|
26
|
+
"@earendil-works/pi-coding-agent": "*"
|
|
27
|
+
},
|
|
28
|
+
"devDependencies": {
|
|
29
|
+
"@types/node": "^26.1.1"
|
|
30
|
+
}
|
|
31
|
+
}
|
package/probe.ts
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
1
|
+
// ── Model probe: availability testing ────────────────────────────
|
|
2
|
+
// Sends a tiny prompt to every available model and records results.
|
|
3
|
+
// Used by /bifrost probe to surface real-world model health.
|
|
4
|
+
|
|
5
|
+
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
6
|
+
import type { Api, Model } from "@earendil-works/pi-ai";
|
|
7
|
+
import { existsSync, mkdirSync, writeFileSync } from "node:fs";
|
|
8
|
+
import { dirname, join } from "node:path";
|
|
9
|
+
|
|
10
|
+
export interface ProbeResult {
|
|
11
|
+
provider: string;
|
|
12
|
+
model: string;
|
|
13
|
+
cost_input: number;
|
|
14
|
+
cost_output: number;
|
|
15
|
+
status: "ok" | "error" | "timeout" | "skipped";
|
|
16
|
+
duration_ms: number;
|
|
17
|
+
error?: string;
|
|
18
|
+
tokens?: number;
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
const PROBE_PROMPT = "1+1=";
|
|
22
|
+
export const PROBE_PROMPT_TEXT = PROBE_PROMPT;
|
|
23
|
+
const PROBE_TIMEOUT_MS = 10_000;
|
|
24
|
+
const PROBE_MAX_TOKENS = 5;
|
|
25
|
+
|
|
26
|
+
export async function runProbe(
|
|
27
|
+
ctx: ExtensionContext,
|
|
28
|
+
onProgress?: (done: number, total: number, last: ProbeResult) => void,
|
|
29
|
+
): Promise<{ results: ProbeResult[]; path: string }> {
|
|
30
|
+
const available = ctx.modelRegistry.getAvailable();
|
|
31
|
+
const total = available.length;
|
|
32
|
+
const CONCURRENCY = 8;
|
|
33
|
+
const results: ProbeResult[] = new Array(total);
|
|
34
|
+
let cursor = 0;
|
|
35
|
+
let completed = 0;
|
|
36
|
+
|
|
37
|
+
async function worker() {
|
|
38
|
+
while (cursor < total) {
|
|
39
|
+
const i = cursor++;
|
|
40
|
+
results[i] = await probeOne(ctx, available[i]);
|
|
41
|
+
completed++;
|
|
42
|
+
onProgress?.(completed, total, results[i]);
|
|
43
|
+
}
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
const workers = Array.from({ length: Math.min(CONCURRENCY, total) }, () => worker());
|
|
47
|
+
await Promise.all(workers);
|
|
48
|
+
|
|
49
|
+
const outputPath = join(process.cwd(), ".pi", "bifrost-probe.json");
|
|
50
|
+
const dir = dirname(outputPath);
|
|
51
|
+
if (!existsSync(dir)) mkdirSync(dir, { recursive: true });
|
|
52
|
+
writeFileSync(outputPath, JSON.stringify(results, null, 2), "utf-8");
|
|
53
|
+
return { results, path: outputPath };
|
|
54
|
+
}
|
|
55
|
+
|
|
56
|
+
async function probeOne(
|
|
57
|
+
ctx: ExtensionContext,
|
|
58
|
+
model: Model<Api>,
|
|
59
|
+
): Promise<ProbeResult> {
|
|
60
|
+
const base: ProbeResult = {
|
|
61
|
+
provider: model.provider,
|
|
62
|
+
model: model.id,
|
|
63
|
+
cost_input: model.cost?.input ?? 0,
|
|
64
|
+
cost_output: model.cost?.output ?? 0,
|
|
65
|
+
status: "skipped",
|
|
66
|
+
duration_ms: 0,
|
|
67
|
+
};
|
|
68
|
+
|
|
69
|
+
// Only probe OpenAI-compatible HTTP models.
|
|
70
|
+
const api = model.api as string | undefined;
|
|
71
|
+
if (
|
|
72
|
+
!api ||
|
|
73
|
+
(!api.startsWith("openai-") && api !== "mistral-conversations")
|
|
74
|
+
) {
|
|
75
|
+
base.error = `unsupported api: ${api}`;
|
|
76
|
+
return base;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
const start = performance.now();
|
|
80
|
+
try {
|
|
81
|
+
const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
|
|
82
|
+
const apiKey = auth.ok ? auth.apiKey : undefined;
|
|
83
|
+
const headers: Record<string, string> = {
|
|
84
|
+
"Content-Type": "application/json",
|
|
85
|
+
...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
|
|
86
|
+
...(auth.ok && auth.headers ? auth.headers : {}),
|
|
87
|
+
};
|
|
88
|
+
|
|
89
|
+
const controller = new AbortController();
|
|
90
|
+
const timer = setTimeout(() => controller.abort(), PROBE_TIMEOUT_MS);
|
|
91
|
+
|
|
92
|
+
try {
|
|
93
|
+
let baseUrl = model.baseUrl.endsWith("/")
|
|
94
|
+
? model.baseUrl
|
|
95
|
+
: `${model.baseUrl}/`;
|
|
96
|
+
// Guard against providers that already include the endpoint path.
|
|
97
|
+
if (baseUrl.endsWith("/chat/completions/") || baseUrl.endsWith("/v1/chat/completions/")) {
|
|
98
|
+
// URL is already the full chat/completions endpoint.
|
|
99
|
+
} else {
|
|
100
|
+
baseUrl = new URL("chat/completions", baseUrl).toString();
|
|
101
|
+
}
|
|
102
|
+
const url = baseUrl;
|
|
103
|
+
|
|
104
|
+
const response = await fetch(url, {
|
|
105
|
+
method: "POST",
|
|
106
|
+
headers,
|
|
107
|
+
body: JSON.stringify({
|
|
108
|
+
model: model.id,
|
|
109
|
+
messages: [{ role: "user", content: PROBE_PROMPT }],
|
|
110
|
+
max_tokens: PROBE_MAX_TOKENS,
|
|
111
|
+
temperature: 0,
|
|
112
|
+
stream: false,
|
|
113
|
+
}),
|
|
114
|
+
signal: controller.signal,
|
|
115
|
+
});
|
|
116
|
+
|
|
117
|
+
base.duration_ms = +(performance.now() - start).toFixed(1);
|
|
118
|
+
|
|
119
|
+
if (!response.ok) {
|
|
120
|
+
const body = await response.text().catch(() => "");
|
|
121
|
+
base.status = "error";
|
|
122
|
+
base.error = `HTTP ${response.status}: ${body.slice(0, 200)}`;
|
|
123
|
+
return base;
|
|
124
|
+
}
|
|
125
|
+
|
|
126
|
+
const data = (await response.json()) as {
|
|
127
|
+
choices?: Array<{ message?: { content?: string } }>;
|
|
128
|
+
usage?: { total_tokens?: number };
|
|
129
|
+
};
|
|
130
|
+
base.status = "ok";
|
|
131
|
+
base.tokens = data.usage?.total_tokens;
|
|
132
|
+
} finally {
|
|
133
|
+
clearTimeout(timer);
|
|
134
|
+
}
|
|
135
|
+
} catch (err) {
|
|
136
|
+
base.duration_ms = +(performance.now() - start).toFixed(1);
|
|
137
|
+
base.status = err instanceof DOMException && err.name === "AbortError"
|
|
138
|
+
? "timeout"
|
|
139
|
+
: "error";
|
|
140
|
+
base.error = String(err).slice(0, 200);
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
return base;
|
|
144
|
+
}
|