@aliou/pi-neuralwatt 0.16.3 → 0.17.0

package/extensions/provider/api/anthropic-messages.ts

@@ -1,12 +1,12 @@
 import type { StreamOptions } from "@earendil-works/pi-ai";
 import { stream, streamSimple } from "@earendil-works/pi-ai/compat";
+import { configuredApiBaseUrl } from "../../../src/config/loader";
 import {
-  NEURALWATT_BASE_URL,
   NEURALWATT_PROVIDER_ID,
   NEURALWATT_REQUEST_HEADERS,
 } from "../constants";
 import { buildAnthropicThinkingLevelMap } from "../models/build";
-import type { NeuralwattModel } from "../models/catalog";
+import type { NeuralwattChatModel } from "../models/catalog";
 import type { AnyStreamSimple } from "../stream-simple";
 import type { NeuralwattApiHandler } from "./types";
 
@@ -61,7 +61,7 @@ function toMessagesBaseUrl(baseUrl: string): string {
   return baseUrl.replace(/\/v1\/?$/, "");
 }
 
-function stampAnthropicModels(models: NeuralwattModel[]) {
+function stampAnthropicModels(models: NeuralwattChatModel[]) {
   return models.map((model) => {
     const { reasoningContract, ...compiled } = model;
     // No retained contract: the identity map is the alias-free special case.
@@ -77,7 +77,7 @@ function stampAnthropicModels(models: NeuralwattModel[]) {
       ...compiled,
       api: "anthropic-messages" as const,
       provider: NEURALWATT_PROVIDER_ID,
-      baseUrl: toMessagesBaseUrl(model.baseUrl ?? NEURALWATT_BASE_URL),
+      baseUrl: toMessagesBaseUrl(model.baseUrl ?? configuredApiBaseUrl()),
       headers: NEURALWATT_REQUEST_HEADERS,
       compat: {
         forceAdaptiveThinking: true,
@@ -105,10 +105,10 @@ export function createAnthropicMessagesApi(options?: {
     stream: (model, context, streamOptions) =>
       stream(model, context, withReasoning(streamOptions) as never),
     streamSimple: (model, context, simpleOptions) =>
-      (options?.streamSimple ?? streamSimple)(
+      ((options?.streamSimple ?? streamSimple) as AnyStreamSimple)(
         model,
         context,
-        withReasoning(simpleOptions) as never,
+        withReasoning(simpleOptions),
       ),
   };
 }

package/extensions/provider/api/openai-completions.ts

@@ -1,11 +1,11 @@
 import type { StreamOptions } from "@earendil-works/pi-ai";
 import { stream, streamSimple } from "@earendil-works/pi-ai/compat";
+import { configuredApiBaseUrl } from "../../../src/config/loader";
 import {
-  NEURALWATT_BASE_URL,
   NEURALWATT_PROVIDER_ID,
   NEURALWATT_REQUEST_HEADERS,
 } from "../constants";
-import type { NeuralwattModel } from "../models/catalog";
+import type { NeuralwattChatModel } from "../models/catalog";
 import type { AnyStreamSimple } from "../stream-simple";
 import type { NeuralwattApiHandler } from "./types";
 
@@ -60,24 +60,24 @@ export function createOpenAiCompletionsApi(options?: {
   });
 
   return {
-    stampModels: (models: NeuralwattModel[]) =>
+    stampModels: (models: NeuralwattChatModel[]) =>
       models.map((model) => {
         const { reasoningContract: _reasoningContract, ...compiled } = model;
         return {
           ...compiled,
           api: "openai-completions" as const,
           provider: NEURALWATT_PROVIDER_ID,
-          baseUrl: model.baseUrl ?? NEURALWATT_BASE_URL,
+          baseUrl: model.baseUrl ?? configuredApiBaseUrl(),
           headers: NEURALWATT_REQUEST_HEADERS,
         };
       }),
     stream: (model, context, streamOptions) =>
       stream(model, context, withReasoningReplay(streamOptions) as never),
     streamSimple: (model, context, simpleOptions) =>
-      (options?.streamSimple ?? streamSimple)(
+      ((options?.streamSimple ?? streamSimple) as AnyStreamSimple)(
         model,
         context,
-        withReasoningReplay(simpleOptions) as never,
+        withReasoningReplay(simpleOptions),
       ),
   };
 }

package/extensions/provider/api/system-one.ts

@@ -0,0 +1,242 @@
+import type {
+  ClassifierAnswer,
+  ClassifierApi,
+  ClassifierContext,
+  ClassifierFunction,
+  ClassifierModel,
+  ClassifierOptions,
+  ClassifierResult,
+} from "@earendil-works/pi-ai";
+import { calculateCost } from "@earendil-works/pi-ai";
+import type { Static } from "typebox";
+import { Type } from "typebox";
+import { Check, Parse, ParseError } from "typebox/value";
+import { configuredApiBaseUrl } from "../../../src/config/loader";
+import {
+  NEURALWATT_PROVIDER_ID,
+  NEURALWATT_REQUEST_HEADERS,
+} from "../constants";
+import {
+  NEURALWATT_SYSTEM_ONE_API,
+  type NeuralwattClassifierModel,
+} from "../models/build";
+
+const UsageSchema = Type.Object({
+  input_tokens: Type.Integer({ minimum: 0 }),
+  output_tokens: Type.Integer({ minimum: 0 }),
+});
+const AnswerSchemas = {
+  choice: Type.Object({
+    type: Type.Literal("choice"),
+    choice: Type.String(),
+    probabilities: Type.Record(Type.String(), Type.Number()),
+    confidence: Type.Number(),
+  }),
+  score: Type.Object({
+    type: Type.Literal("score"),
+    score: Type.Number(),
+    confidence: Type.Number(),
+  }),
+  bool: Type.Object({ type: Type.Literal("noul"), noul: Type.Number() }),
+};
+const ErrorSchema = Type.Union([
+  Type.Object({
+    error: Type.Object({
+      message: Type.String(),
+      type: Type.String(),
+      code: Type.String(),
+    }),
+  }),
+  Type.Object({ detail: Type.String() }),
+]);
+
+function responseSchema(context: ClassifierContext) {
+  return Type.Object({
+    answers: Type.Object(
+      Object.fromEntries(
+        Object.entries(context.questions).map(([id, question]) => [
+          id,
+          AnswerSchemas[question.type],
+        ]),
+      ),
+    ),
+    usage: Type.Optional(UsageSchema),
+  });
+}
+
+type SystemOneResponse = Static<ReturnType<typeof responseSchema>>;
+
+function mapAnswer(
+  value: SystemOneResponse["answers"][string],
+): ClassifierAnswer {
+  if (value.type === "noul") return { type: "bool", probability: value.noul };
+  if (value.type === "score")
+    return { type: "score", score: value.score, confidence: value.confidence };
+  return {
+    type: "choice",
+    choice: value.choice,
+    probabilities: value.probabilities,
+    confidence: value.confidence,
+  };
+}
+
+function mapUsage(
+  model: ClassifierModel<ClassifierApi>,
+  value: SystemOneResponse["usage"],
+): ClassifierResult["usage"] {
+  if (!value) return undefined;
+  const { input_tokens: input, output_tokens: output } = value;
+  const usage = {
+    input,
+    output,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: input + output,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+  calculateCost(model, usage);
+  return usage;
+}
+
+function billedUsageOnInvalidResponse(
+  model: ClassifierModel<ClassifierApi>,
+  error: ParseError,
+) {
+  const usage = (error.cause.value as { usage?: unknown } | null)?.usage;
+  return Check(UsageSchema, usage) ? mapUsage(model, usage) : undefined;
+}
+
+function formatRequest(
+  model: ClassifierModel<ClassifierApi>,
+  context: ClassifierContext,
+) {
+  return {
+    model: model.id,
+    state: context.state,
+    questions: Object.fromEntries(
+      Object.entries(context.questions).map(([id, question]) => [
+        id,
+        question.type === "bool" ? { ...question, type: "noul" } : question,
+      ]),
+    ),
+  };
+}
+
+function requestHeaders(
+  model: ClassifierModel<ClassifierApi>,
+  options: ClassifierOptions,
+) {
+  const headers = new Headers({
+    authorization: `Bearer ${options.apiKey || "-"}`,
+    "content-type": "application/json",
+  });
+  const overrides = [model.headers, options.headers].flatMap((source) =>
+    Object.entries(source ?? {}),
+  );
+  for (const [name, value] of overrides) {
+    if (value === null) headers.delete(name);
+    else headers.set(name, value);
+  }
+  return Object.fromEntries(headers);
+}
+
+function requestSignal({ signal, timeoutMs }: ClassifierOptions) {
+  if (timeoutMs === undefined) return signal;
+  const timeout = AbortSignal.timeout(timeoutMs);
+  return signal ? AbortSignal.any([signal, timeout]) : timeout;
+}
+
+async function request(
+  model: ClassifierModel<ClassifierApi>,
+  context: ClassifierContext,
+  options: ClassifierOptions,
+) {
+  if (model.api !== NEURALWATT_SYSTEM_ONE_API) {
+    throw new Error(`Unsupported classifier API: ${model.api}`);
+  }
+
+  if (context.images?.length) {
+    throw new Error("Neuralwatt classifier does not support image input");
+  }
+
+  options.signal?.throwIfAborted();
+
+  const rawPayload = formatRequest(model, context);
+  const payload = (await options.onPayload?.(rawPayload, model)) ?? rawPayload;
+
+  const url = new URL("systemone", `${model.baseUrl.replace(/\/+$/u, "")}/`);
+
+  const requestFetch = options.fetch ?? globalThis.fetch;
+
+  const response = await requestFetch(url, {
+    method: "POST",
+    headers: requestHeaders(model, options),
+    body: JSON.stringify(payload),
+    signal: requestSignal(options),
+  });
+
+  await options.onResponse?.(
+    { status: response.status, headers: Object.fromEntries(response.headers) },
+    model,
+  );
+
+  const body: unknown = await response.json();
+  if (!response.ok) {
+    const failure = Parse(ErrorSchema, body);
+    const message = "error" in failure ? failure.error.message : failure.detail;
+    throw new Error(
+      `Neuralwatt System One API returned ${response.status}: ${message}`,
+    );
+  }
+
+  return Parse(responseSchema(context), body);
+}
+
+export const classify: ClassifierFunction = async (
+  model,
+  context,
+  options = {},
+) => {
+  const result: ClassifierResult = {
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    answers: {},
+    stopReason: "stop",
+    timestamp: Date.now(),
+  };
+  try {
+    const body = await request(model, context, options);
+    result.usage = mapUsage(model, body.usage);
+    result.answers = Object.fromEntries(
+      Object.keys(context.questions).map((id) => [
+        id,
+        mapAnswer(body.answers[id]),
+      ]),
+    );
+  } catch (error) {
+    result.stopReason = options.signal?.aborted ? "aborted" : "error";
+    if (error instanceof ParseError) {
+      result.usage = billedUsageOnInvalidResponse(model, error);
+      result.errorMessage = error.cause.errors
+        .map(({ instancePath, message }) => `${instancePath}: ${message}`)
+        .join("; ");
+      return result;
+    }
+    result.errorMessage =
+      error instanceof Error ? error.message : String(error);
+  }
+  return result;
+};
+
+export function stampClassifierModels(
+  models: NeuralwattClassifierModel[],
+): ClassifierModel<ClassifierApi>[] {
+  return models.map((model) => ({
+    ...model,
+    api: NEURALWATT_SYSTEM_ONE_API,
+    provider: NEURALWATT_PROVIDER_ID,
+    baseUrl: model.baseUrl ?? configuredApiBaseUrl(),
+    headers: NEURALWATT_REQUEST_HEADERS,
+  }));
+}

package/extensions/provider/api/types.ts

@@ -1,24 +1,24 @@
 import type {
   Api,
   AssistantMessageEventStream,
-  Context,
   Model,
   SimpleStreamOptions,
   StreamOptions,
+  TranscriptContext,
 } from "@earendil-works/pi-ai";
-import type { NeuralwattModel } from "../models/catalog";
+import type { NeuralwattChatModel } from "../models/catalog";
 
-/** One Neuralwatt API surface: model stamping plus submission plumbing. */
+/** One Neuralwatt API surface: model stamping plus chat submission plumbing. */
 export interface NeuralwattApiHandler {
-  stampModels(models: NeuralwattModel[]): Model<Api>[];
+  stampModels(models: NeuralwattChatModel[]): Model<Api>[];
   stream(
     model: Model<Api>,
-    context: Context,
+    context: TranscriptContext,
     options?: StreamOptions,
   ): AssistantMessageEventStream;
   streamSimple(
     model: Model<Api>,
-    context: Context,
+    context: TranscriptContext,
     options?: SimpleStreamOptions,
   ): AssistantMessageEventStream;
 }

package/extensions/provider/commands/settings/api-base-url-editor.ts

@@ -0,0 +1,129 @@
+import type { Component, SettingsListTheme } from "@earendil-works/pi-tui";
+import {
+  Input,
+  Key,
+  matchesKey,
+  truncateToWidth,
+} from "@earendil-works/pi-tui";
+
+export interface ApiBaseUrlEditorOptions {
+  initial: string;
+  defaultBaseUrl: string;
+  theme: SettingsListTheme;
+  requestRender: () => void;
+  onSubmit: (value: string) => void;
+  onCancel: () => void;
+}
+
+type Status =
+  | { kind: "idle" }
+  | { kind: "validating"; url: string }
+  | { kind: "error"; message: string };
+
+export function normalizeApiBaseUrlInput(value: string): string {
+  return value.trim().replace(/\/+$/u, "");
+}
+
+export async function checkApiBase(
+  base: string,
+  timeoutMs = 8_000,
+): Promise<string | null> {
+  try {
+    const response = await fetch(`${base}/models`, {
+      headers: { Accept: "application/json" },
+      signal: AbortSignal.timeout(timeoutMs),
+    });
+    return response.ok ? null : `HTTP ${response.status} from ${base}/models`;
+  } catch (error) {
+    return error instanceof Error && error.name === "TimeoutError"
+      ? `Timed out reaching ${base}/models`
+      : `Could not reach ${base}/models`;
+  }
+}
+
+export class ApiBaseUrlEditor implements Component {
+  private readonly input = new Input();
+  private readonly defaultBaseUrl: string;
+  private readonly theme: SettingsListTheme;
+  private readonly requestRender: () => void;
+  private readonly onSubmit: (value: string) => void;
+  private readonly onCancel: () => void;
+  private status: Status = { kind: "idle" };
+
+  constructor(options: ApiBaseUrlEditorOptions) {
+    this.defaultBaseUrl = options.defaultBaseUrl;
+    this.theme = options.theme;
+    this.requestRender = options.requestRender;
+    this.onSubmit = options.onSubmit;
+    this.onCancel = options.onCancel;
+    this.input.setValue(options.initial);
+    this.input.onSubmit = (value) => void this.submit(value);
+    this.input.onEscape = () => this.onCancel();
+  }
+
+  private async submit(raw: string): Promise<void> {
+    const base = normalizeApiBaseUrlInput(raw);
+    if (base) {
+      this.status = { kind: "validating", url: `${base}/models` };
+      this.requestRender();
+      const error = await checkApiBase(base);
+      if (error) {
+        this.status = { kind: "error", message: error };
+        this.requestRender();
+        return;
+      }
+    }
+    this.onSubmit(base);
+  }
+
+  invalidate(): void {}
+
+  handleInput(data: string): void {
+    if (matchesKey(data, Key.escape)) {
+      this.onCancel();
+      return;
+    }
+    if (this.status.kind === "validating") return;
+    this.status = { kind: "idle" };
+    this.input.handleInput(data);
+  }
+
+  render(width: number): string[] {
+    const lines: string[] = [];
+    lines.push(
+      truncateToWidth(this.theme.label("API base URL", true), width, "", true),
+    );
+    lines.push(...this.input.render(width));
+    if (this.status.kind === "validating") {
+      lines.push(
+        truncateToWidth(
+          this.theme.description(`Validating ${this.status.url} …`),
+          width,
+          "",
+          true,
+        ),
+      );
+    } else if (this.status.kind === "error") {
+      lines.push(
+        truncateToWidth(
+          this.theme.value(this.status.message, true),
+          width,
+          "",
+          true,
+        ),
+      );
+    } else {
+      lines.push(
+        truncateToWidth(
+          this.theme.description(
+            `Enter to save · empty = direct upstream (${this.defaultBaseUrl})`,
+          ),
+          width,
+          "",
+          true,
+        ),
+      );
+    }
+    return lines;
+  }
+}

package/extensions/provider/commands/settings/index.ts

@@ -14,6 +14,7 @@ import {
   NEURALWATT_CONFIG_UPDATED_EVENT,
   type NeuralwattFeatureId,
 } from "../../../../src/events";
+import { ApiBaseUrlEditor } from "./api-base-url-editor";
 
 export interface RegisterNeuralwattSettingsOptions {
   getLoadedFeatures: () => Set<NeuralwattFeatureId>;
@@ -62,15 +63,16 @@ export function registerNeuralwattSettings(
   options: RegisterNeuralwattSettingsOptions,
 ): void {
   const { getLoadedFeatures } = options;
-  // The provider stamps `provider.api` at extension load; a saved change only
-  // reaches it after `/reload`.
+  // The provider stamps `provider.api` and assembles against the configured
+  // base at extension load; a saved change only reaches it after `/reload`.
   let pendingApi: NeuralwattApi | undefined;
+  let pendingApiBaseUrl: string | undefined;
 
   registerSettingsCommand<NeuralwattConfig, ResolvedNeuralwattConfig>(pi, {
     commandName: "neuralwatt:settings",
     title: "Neuralwatt Settings",
     configStore: configLoader,
-    buildSections: (tabConfig, resolved): SettingsSection[] => {
+    buildSections: (tabConfig, resolved, buildCtx): SettingsSection[] => {
       const loaded = getLoadedFeatures();
       return [
         {
@@ -80,10 +82,41 @@ export function registerNeuralwattSettings(
               id: "api",
               label: "API",
               description:
-                "Serve models via the OpenAI-compatible chat/completions endpoint or the Anthropic-compatible /v1/messages endpoint",
+                "Serve models via the OpenAI completions endpoint or the Anthropic endpoint",
               currentValue: tabConfig?.provider?.api ?? resolved.provider.api,
               values: ["openai-completions", "anthropic-messages"],
             },
+            {
+              id: "apiBaseUrl",
+              label: "API base URL",
+              description:
+                "Base URL for API requests; empty uses the default upstream",
+              currentValue: tabConfig?.provider?.apiBaseUrl || "(not set)",
+              submenu: (current, submenuDone, submenuCtx) => {
+                const display = (value: string) => value || "(not set)";
+                return new ApiBaseUrlEditor({
+                  initial: current === "(not set)" ? "" : current,
+                  defaultBaseUrl: resolved.provider.apiBaseUrl,
+                  theme: buildCtx.theme,
+                  requestRender: submenuCtx.requestRender,
+                  onSubmit: (value) => {
+                    const updated = structuredClone(
+                      tabConfig ?? {},
+                    ) as NeuralwattConfig;
+                    updated.provider = { ...updated.provider };
+                    if (value) {
+                      updated.provider.apiBaseUrl = value;
+                    } else {
+                      delete updated.provider.apiBaseUrl;
+                    }
+                    buildCtx.setDraft(updated);
+                    pendingApiBaseUrl = value;
+                    submenuDone(display(value));
+                  },
+                  onCancel: () => submenuDone(display(current)),
+                });
+              },
+            },
           ],
         },
         {
@@ -165,9 +198,12 @@ export function registerNeuralwattSettings(
     },
     onSave: async (ctx) => {
       emitConfigUpdated(pi);
-      if (pendingApi === undefined) return;
+      if (pendingApi === undefined && pendingApiBaseUrl === undefined) {
+        return;
+      }
       pendingApi = undefined;
-      ctx.ui.notify("Run /reload to apply the new API", "info");
+      pendingApiBaseUrl = undefined;
+      ctx.ui.notify("Run /reload to apply the new provider settings", "info");
     },
   });
 }

package/extensions/provider/index.ts

@@ -30,7 +30,10 @@ import {
   parseRateLimitHeaders,
 } from "./rate-limit-error";
 import { updateQuotasFromSseComment } from "./sse-quotas";
-import { wrapNeuralwattStreamSimple } from "./stream-simple";
+import {
+  type NeuralwattStreamCallbacks,
+  wrapNeuralwattStreamSimple,
+} from "./stream-simple";
 
 const HEADER_EMIT_THROTTLE_MS = 5_000;
 
@@ -42,7 +45,7 @@ function emitConfigUpdated(pi: ExtensionAPI): void {
 
 function registerNeuralwattProvider(
   pi: ExtensionAPI,
-  onSseQuota: (line: string) => void,
+  streamCallbacks: NeuralwattStreamCallbacks,
 ): void {
   const staticModels = buildNeuralwattProviderModels();
 
@@ -51,7 +54,7 @@ function registerNeuralwattProvider(
   const streamSimple = baseStreamSimple
     ? (wrapNeuralwattStreamSimple(
         baseStreamSimple as never,
-        onSseQuota,
+        streamCallbacks,
       ) as never)
     : undefined;
 
@@ -60,7 +63,7 @@ function registerNeuralwattProvider(
   const messagesStreamSimple = messagesBaseStreamSimple
     ? (wrapNeuralwattStreamSimple(
         messagesBaseStreamSimple as never,
-        onSseQuota,
+        streamCallbacks,
       ) as never)
     : undefined;
 
@@ -101,7 +104,17 @@ export default async function (pi: ExtensionAPI) {
     emitQuotas(quotas, "sse");
   };
 
-  registerNeuralwattProvider(pi, handleSseQuota);
+  // Rate-limit info captured by the per-request fetch wrapper from the most
+  // recent 429 response. Used in message_end to rewrite the generic error
+  // text with actionable details from Neuralwatt's response headers.
+  let pendingRateLimitInfo: NeuralwattRateLimitInfo | undefined;
+
+  registerNeuralwattProvider(pi, {
+    onSseQuota: handleSseQuota,
+    onRateLimit: (info) => {
+      pendingRateLimitInfo = info;
+    },
+  });
 
   const loadedFeatures = new Set<NeuralwattFeatureId>();
 
@@ -125,10 +138,6 @@ export default async function (pi: ExtensionAPI) {
     pi.events.emit(NEURALWATT_QUOTAS_UPDATED_EVENT, { quotas, source });
   }
 
-  // Stored rate-limit info from the most recent 429 response.
-  // Used in message_end to rewrite the generic error text with
-  // actionable details from Neuralwatt's response headers.
-  let pendingRateLimitInfo: NeuralwattRateLimitInfo | undefined;
   let currentModelRegistry: ModelRegistry | undefined;
 
   pi.on("message_end", (event, ctx) => {
@@ -148,14 +157,14 @@ export default async function (pi: ExtensionAPI) {
       return { message };
     }
 
-    // Fallback for 429s where no layer-specific headers were captured. The
-    // streamSimple wrap (wrapNeuralwattStreamSimple) already formats a
-    // detailed message via formatRateLimitError when it captures headers;
-    // detect that case by the `"429 rate limit:"` prefix it emits and leave
-    // it untouched. This branch only fires for genuinely headerless 429s
-    // (e.g. anonymous playground limits, or a 429 from infra in front of
-    // Neuralwatt), since after_provider_response cannot observe 429s — the
-    // OpenAI SDK throws before Pi's onResponse hook runs.
+    // Fallback for 429s where no layer-specific headers were captured (the
+    // fetch wrapper reported `undefined` via onRateLimit, or the 429 came
+    // from infra in front of Neuralwatt). The `"429 rate limit:"` prefix
+    // guard skips messages the first branch already rewrote. On the
+    // openai-completions surface after_provider_response cannot observe
+    // 429s — the OpenAI SDK throws before Pi's onResponse hook runs — so
+    // this is the only repair path there; on anthropic-messages
+    // after_provider_response also feeds pendingRateLimitInfo directly.
     if (
       event.message.role === "assistant" &&
       event.message.stopReason === "error" &&

package/extensions/provider/models/build.ts

@@ -4,8 +4,24 @@ import type {
   NeuralwattReasoningEffort,
 } from "../../../src/types/models-api";
 
+export type ProviderChatModelConfig = Extract<
+  ProviderModelConfig,
+  { type?: "chat" }
+>;
+
+export type ProviderClassifierModelConfig = Extract<
+  ProviderModelConfig,
+  { type: "classifier" }
+>;
+
+export const NEURALWATT_SYSTEM_ONE_API = "typesafe-system-one" as const;
+
+export type NeuralwattClassifierModel = ProviderClassifierModelConfig & {
+  api: typeof NEURALWATT_SYSTEM_ONE_API;
+};
+
 export type ThinkingLevelMap = NonNullable<
-  ProviderModelConfig["thinkingLevelMap"]
+  ProviderChatModelConfig["thinkingLevelMap"]
 >;
 
 /**
@@ -13,7 +29,7 @@ export type ThinkingLevelMap = NonNullable<
  * retained for anthropic-messages map derivation. Rides the models store
  * (JSON passthrough); stripped from stamped runtime models.
  */
-export type NeuralwattCompiledModel = ProviderModelConfig & {
+export type NeuralwattCompiledModel = ProviderChatModelConfig & {
   reasoningContract?: NeuralwattReasoningMapSource;
 };
 
@@ -157,7 +173,7 @@ export function buildNeuralwattModel(
 ): NeuralwattCompiledModel {
   const vision = variant.vision ?? family.vision;
 
-  const compat: NonNullable<ProviderModelConfig["compat"]> = {
+  const compat: NonNullable<ProviderChatModelConfig["compat"]> = {
     supportsDeveloperRole: false,
     maxTokensField: "max_tokens",
   };

package/extensions/provider/models/catalog.ts

@@ -1,19 +1,45 @@
-import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
 import type { NeuralwattApiModel } from "../../../src/types/models-api";
 import {
   buildThinkingLevelMap,
+  NEURALWATT_SYSTEM_ONE_API,
+  type NeuralwattClassifierModel,
   type NeuralwattCompiledModel,
+  type ProviderChatModelConfig,
   resolveMaxTokens,
   type ThinkingLevelMap,
 } from "./build";
 import { NEURALWATT_MODELS } from "./public-models";
 
-export type NeuralwattModel = NeuralwattCompiledModel;
+export type NeuralwattChatModel = NeuralwattCompiledModel;
+
+export type NeuralwattModel = NeuralwattChatModel | NeuralwattClassifierModel;
+
+export function isNeuralwattClassifierModel(
+  model: NeuralwattModel,
+): model is NeuralwattClassifierModel {
+  return model.type === "classifier";
+}
+
+export function partitionNeuralwattModels(models: readonly NeuralwattModel[]): {
+  chat: NeuralwattChatModel[];
+  classifiers: NeuralwattClassifierModel[];
+} {
+  const chat: NeuralwattChatModel[] = [];
+  const classifiers: NeuralwattClassifierModel[] = [];
+  for (const model of models) {
+    if (isNeuralwattClassifierModel(model)) {
+      classifiers.push(model);
+    } else {
+      chat.push(model);
+    }
+  }
+  return { chat, classifiers };
+}
 
 // Chat-template thinking: the API exposes a `reasoning` block, but the
 // underlying mechanism is chat_template_kwargs, so Pi needs the mapping.
 const COMPAT_OVERRIDES: Partial<
-  Record<string, Partial<NonNullable<ProviderModelConfig["compat"]>>>
+  Record<string, Partial<NonNullable<ProviderChatModelConfig["compat"]>>>
 > = {
   "Qwen/Qwen3.8-27B-FP8": {
     thinkingFormat: "chat-template",
@@ -23,7 +49,36 @@ const COMPAT_OVERRIDES: Partial<
   },
 };
 
-function apiModelToProviderModel(model: NeuralwattApiModel): NeuralwattModel {
+function apiModelToClassifierModel(
+  model: NeuralwattApiModel,
+): NeuralwattClassifierModel {
+  const meta = model.metadata;
+  if (!meta)
+    throw new Error(
+      `Neuralwatt API returned model "${model.id}" without metadata`,
+    );
+
+  return {
+    type: "classifier",
+    api: NEURALWATT_SYSTEM_ONE_API,
+    id: model.id,
+    name: meta.display_name ?? model.id,
+    input: meta.capabilities.vision
+      ? (["text", "image"] as const)
+      : (["text"] as const),
+    cost: {
+      input: meta.pricing.input_per_million,
+      output: meta.pricing.output_per_million,
+      cacheRead: meta.pricing.cached_input_per_million ?? 0,
+      cacheWrite: meta.pricing.cached_output_per_million ?? 0,
+    },
+    contextWindow: model.max_model_len,
+  };
+}
+
+function apiModelToProviderModel(
+  model: NeuralwattApiModel,
+): NeuralwattChatModel {
   const meta = model.metadata;
   if (!meta)
     throw new Error(
@@ -32,7 +87,7 @@ function apiModelToProviderModel(model: NeuralwattApiModel): NeuralwattModel {
 
   const reasoning = meta.capabilities.reasoning;
 
-  const compat: NonNullable<ProviderModelConfig["compat"]> = {
+  const compat: NonNullable<ProviderChatModelConfig["compat"]> = {
     supportsDeveloperRole: meta.capabilities.developer_role,
     maxTokensField: "max_tokens",
     ...COMPAT_OVERRIDES[model.id],
@@ -83,13 +138,18 @@ export function buildNeuralwattProviderModelsFromApi(
         m.metadata &&
         !m.metadata.deprecated &&
         !m.metadata.pricing.pricing_tbd &&
-        // Exclude non-chat models (e.g. embeddings) by task
         !(
           m.metadata.capabilities.task &&
-          !["chat", "completions"].includes(m.metadata.capabilities.task)
+          !["chat", "completions", "decision"].includes(
+            m.metadata.capabilities.task,
+          )
         ),
     )
-    .map(apiModelToProviderModel);
+    .map((m) =>
+      m.metadata?.capabilities.task === "decision"
+        ? apiModelToClassifierModel(m)
+        : apiModelToProviderModel(m),
+    );
   return models;
 }
 

package/extensions/provider/models/public-models.ts

@@ -1,8 +1,8 @@
-import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
 import {
   buildNeuralwattFamily,
   type NeuralwattModelFamily,
   type NeuralwattVariantSpec,
+  type ProviderChatModelConfig,
 } from "./build";
 
 // Public models returned by https://api.neuralwatt.com/v1/models.
@@ -375,6 +375,14 @@ const FAMILIES: [NeuralwattModelFamily, NeuralwattVariantSpec[]][] = [
         maxOutputTokens: 393216,
         reasoning: true,
       },
+      {
+        id: "mimo-v2.6-pro-flex",
+        name: "MiMo-V2.6-Pro (flex)",
+        contextWindow: 1048560,
+        maxOutputTokens: 393216,
+        reasoning: true,
+        costMultiplier: 0.65,
+      },
     ],
   ],
   [
@@ -492,6 +500,6 @@ const FAMILIES: [NeuralwattModelFamily, NeuralwattVariantSpec[]][] = [
 // `costMultiplier: 0.65` per variant.
 // https://docs.neuralwatt.com/guides/flex-tier.md
 
-export const NEURALWATT_MODELS: ProviderModelConfig[] = FAMILIES.flatMap(
+export const NEURALWATT_MODELS: ProviderChatModelConfig[] = FAMILIES.flatMap(
   ([family, variants]) => buildNeuralwattFamily(family, variants),
 );

package/extensions/provider/models/refresh.ts

@@ -1,7 +1,9 @@
 import type {
+  AnyModel,
   ModelsStoreEntry,
   RefreshModelsContext,
 } from "@earendil-works/pi-ai";
+import { configuredApiBaseUrl } from "../../../src/config/loader";
 import type { NeuralwattApiModel } from "../../../src/types/models-api";
 import type {
   buildNeuralwattProviderModels,
@@ -17,7 +19,8 @@ export const MODEL_STORE_TTL_MS = 4 * 60 * 60 * 1000;
  * key-scoped catalog (preview, grant-gated, private models), so an entry
  * stamped "public" must not shadow a keyed refresh — and a key-scoped entry
  * must not be replayed for an anonymous user. Matches the anonymous-key
- * convention in src/lib/neuralwatt-api.ts (authHeaders).
+ * convention in src/lib/neuralwatt-api.ts (authHeaders). The base URL is part
+ * of the key because a gateway can serve a different catalog than upstream.
  */
 const CATALOG_SCOPE_VERSION = "v2";
 type CatalogScope = "public" | "key";
@@ -30,18 +33,15 @@ function catalogScope(apiKey: string | undefined): CatalogScope {
     : "public";
 }
 
-function storedCatalogKey(entry: ScopedModelsStoreEntry): string | undefined {
-  return entry.catalogKey;
+function catalogKey(scope: CatalogScope): string {
+  return `${scope} ${CATALOG_SCOPE_VERSION} ${configuredApiBaseUrl()}`;
 }
 
 function catalogKeyMatches(
   entry: ScopedModelsStoreEntry | undefined,
   scope: CatalogScope,
 ): boolean {
-  return (
-    entry !== undefined &&
-    storedCatalogKey(entry) === `${scope} ${CATALOG_SCOPE_VERSION}`
-  );
+  return entry !== undefined && entry.catalogKey === catalogKey(scope);
 }
 
 export type FetchNeuralwattApiModels = (
@@ -49,6 +49,10 @@ export type FetchNeuralwattApiModels = (
   signal?: AbortSignal,
 ) => Promise<readonly NeuralwattApiModel[]>;
 
+function asStoredChatModels(models: readonly AnyModel[]): NeuralwattModel[] {
+  return models as unknown as NeuralwattModel[];
+}
+
 function isFreshStoreEntry(
   entry: Readonly<ModelsStoreEntry> | undefined,
 ): entry is ModelsStoreEntry {
@@ -82,11 +86,11 @@ export function createNeuralwattRefreshModels(
       const stored = context.stored as ScopedModelsStoreEntry | undefined;
       if (!context.allowNetwork) {
         return stored !== undefined && catalogKeyMatches(stored, scope)
-          ? buildFromStore(stored.models)
+          ? buildFromStore(asStoredChatModels(stored.models))
           : fallback;
       }
       if (!context.force && isUsableStoreEntry(stored, scope)) {
-        return buildFromStore(stored.models);
+        return buildFromStore(asStoredChatModels(stored.models));
       }
       const apiModels = await fetchApiModels(apiKey, context.signal);
       context.signal.throwIfAborted();
@@ -94,7 +98,7 @@ export function createNeuralwattRefreshModels(
       const entry: ScopedModelsStoreEntry = {
         models: models as unknown as ModelsStoreEntry["models"],
         checkedAt: Date.now(),
-        catalogKey: `${scope} ${CATALOG_SCOPE_VERSION}`,
+        catalogKey: catalogKey(scope),
       };
       await context.publish({ persist: entry }).catch(() => undefined);
       context.signal.throwIfAborted();

package/extensions/provider/provider.ts

@@ -2,6 +2,7 @@ import type { Provider } from "@earendil-works/pi-ai";
 import type { NeuralwattApi } from "../../src/config";
 import { createAnthropicMessagesApi } from "./api/anthropic-messages";
 import { createOpenAiCompletionsApi } from "./api/openai-completions";
+import { classify, stampClassifierModels } from "./api/system-one";
 import type { NeuralwattApiHandler } from "./api/types";
 import {
   NEURALWATT_API_KEY_ENV,
@@ -13,6 +14,7 @@ import type { NeuralwattModel } from "./models/catalog";
 import {
   buildNeuralwattProviderModelsFromApi,
   buildNeuralwattProviderModelsFromStore,
+  partitionNeuralwattModels,
 } from "./models/catalog";
 import {
   createNeuralwattRefreshModels,
@@ -98,11 +100,22 @@ export function createNeuralwattProvider(
           if (envKey) {
             return { auth: { apiKey: envKey }, source: NEURALWATT_API_KEY_ENV };
           }
-          return { auth: { apiKey: "" }, source: "anonymous" };
+          // pi-ai rejects empty keys before dispatch; "-" is the neutral
+          // placeholder that gateways strip or replace server-side.
+          return { auth: { apiKey: "-" }, source: "anonymous" };
         },
       },
     },
-    getModels: () => handler.stampModels(canonicalModels),
+    getModels: () =>
+      handler.stampModels(partitionNeuralwattModels(canonicalModels).chat),
+    getAllModels: () => {
+      const { chat, classifiers } = partitionNeuralwattModels(canonicalModels);
+      return [
+        ...handler.stampModels(chat),
+        ...stampClassifierModels(classifiers),
+      ];
+    },
+    classify: (model, context, options) => classify(model, context, options),
     refreshModels: async (context) => {
       const models = await refreshCatalog(context);
       await context.publish({

package/extensions/provider/stream-simple.ts

@@ -1,32 +1,39 @@
 // Neuralwatt is OpenAI-compatible, but the OpenAI SDK throws on non-2xx
 // responses before Pi's after_provider_response hook can see the raw headers.
-// We wrap the built-in openai-completions streamSimple so 429 rate-limit
-// headers can be captured before the SDK turns them into a generic error, while
-// still delegating normal streaming behavior to Pi's provider implementation.
+// We inject a per-request `options.fetch` wrapper so 429
+// rate-limit headers can be captured before the SDK turns them into a generic
+// error, and tee successful SSE bodies for live quota comments, while still
+// delegating normal streaming behavior to Pi's provider implementation.
+// `globalThis.fetch` is never touched.
 //
 // The SSE tee used for live quota comments is inspired by:
 // https://github.com/monotykamary/pi-neuralwatt-provider
 
-import {
-  type AssistantMessageEventStream,
-  type Context,
-  createAssistantMessageEventStream,
-  type Model,
-  type SimpleStreamOptions,
+import type {
+  AssistantMessageEventStream,
+  FetchFunction,
+  Model,
+  SimpleStreamOptions,
+  TranscriptContext,
 } from "@earendil-works/pi-ai";
+import { configuredApiBaseUrl } from "../../src/config/loader";
 import {
   type NeuralwattRateLimitInfo,
-  normalizeNeuralwattRateLimitError,
   parseRateLimitHeaders,
 } from "./rate-limit-error";
 import { readQuotaCommentsFromTee } from "./sse-quotas";
 
 export type AnyStreamSimple = (
   model: Model<string>,
-  context: Context,
+  context: TranscriptContext,
   options?: SimpleStreamOptions,
 ) => AssistantMessageEventStream;
 
+export interface NeuralwattStreamCallbacks {
+  onSseQuota: (line: string) => void;
+  onRateLimit: (info: NeuralwattRateLimitInfo | undefined) => void;
+}
+
 function headersToRecord(headers: Headers): Record<string, string> {
   const record: Record<string, string> = {};
   headers.forEach((value, key) => {
@@ -58,56 +65,34 @@ function isProviderStreamUrl(
   }
 }
 
-async function forwardStream(
-  stream: AssistantMessageEventStream,
-  outer: AssistantMessageEventStream,
-  getRateLimitInfo: () => NeuralwattRateLimitInfo | undefined,
-  restoreFetch: () => void,
-): Promise<void> {
-  try {
-    for await (const event of stream) {
-      const rateLimitInfo = getRateLimitInfo();
-      if (event.type === "error" && rateLimitInfo) {
-        outer.push({
-          ...event,
-          error: normalizeNeuralwattRateLimitError(event.error, rateLimitInfo),
-        });
-      } else {
-        outer.push(event);
-      }
-    }
-  } finally {
-    restoreFetch();
-    outer.end();
-  }
-}
-
 export function wrapNeuralwattStreamSimple(
   base: AnyStreamSimple,
-  onSseQuota: (line: string) => void,
+  callbacks: NeuralwattStreamCallbacks,
 ): AnyStreamSimple {
   return (model, context, options = {}) => {
-    let rateLimitInfo: NeuralwattRateLimitInfo | undefined;
-    let sseQuotaTask: Promise<void> | undefined;
-    const outer = createAssistantMessageEventStream();
-    const providerOrigin = new URL(
-      model.baseUrl ?? "https://api.neuralwatt.com/v1",
-    ).origin;
-    const originalFetch = globalThis.fetch;
-    const wrappedFetch: typeof fetch = async (input, init) => {
-      const response = await originalFetch(input, init);
+    const providerOrigin = new URL(model.baseUrl ?? configuredApiBaseUrl())
+      .origin;
+    const callerFetch = options.fetch;
+
+    const neuralwattFetch: FetchFunction = async (input, init) => {
+      // Chain onto a caller-provided fetch when present; otherwise resolve
+      // the global at request time so late host/test overrides still apply.
+      const response = await (callerFetch ?? globalThis.fetch)(input, init);
 
       if (!isProviderStreamUrl(input, providerOrigin)) return response;
 
-      const headers = headersToRecord(response.headers);
       if (response.status === 429) {
-        rateLimitInfo = parseRateLimitHeaders(headers);
+        callbacks.onRateLimit(
+          parseRateLimitHeaders(headersToRecord(response.headers)),
+        );
         return response;
       }
 
       if (response.ok && response.body) {
         const [sdkBody, quotaBody] = response.body.tee();
-        sseQuotaTask = readQuotaCommentsFromTee(quotaBody, onSseQuota);
+        // Best-effort side channel; failures inside are swallowed by the
+        // reader so quota comments can never break the SDK stream.
+        void readQuotaCommentsFromTee(quotaBody, callbacks.onSseQuota);
         return new Response(sdkBody, {
           headers: response.headers,
           status: response.status,
@@ -118,22 +103,6 @@ export function wrapNeuralwattStreamSimple(
       return response;
     };
 
-    globalThis.fetch = wrappedFetch;
-
-    const restoreFetch = () => {
-      if (globalThis.fetch === wrappedFetch) globalThis.fetch = originalFetch;
-      sseQuotaTask?.catch(() => {});
-    };
-
-    const stream = base(model, context, options);
-    const originalOuterEnd = outer.end.bind(outer);
-    outer.end = (result?: Parameters<typeof originalOuterEnd>[0]) => {
-      restoreFetch();
-      originalOuterEnd(result);
-    };
-
-    void forwardStream(stream, outer, () => rateLimitInfo, restoreFetch);
-
-    return outer;
+    return base(model, context, { ...options, fetch: neuralwattFetch });
   };
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aliou/pi-neuralwatt",
-  "version": "0.16.3",
+  "version": "0.17.0",
   "license": "MIT",
   "type": "module",
   "private": false,
@@ -38,21 +38,22 @@
     "@aliou/pi-utils-ui": "^0.5.0"
   },
   "peerDependencies": {
-    "@earendil-works/pi-ai": ">=0.84.0",
-    "@earendil-works/pi-coding-agent": ">=0.84.0",
+    "@earendil-works/pi-ai": ">=1.0.0",
+    "@earendil-works/pi-coding-agent": ">=1.0.0",
     "@earendil-works/pi-tui": "*"
   },
   "devDependencies": {
     "@aliou/biome-plugins": "^0.11.0",
     "@biomejs/biome": "^2.4.15",
     "@changesets/cli": "^2.27.11",
-    "@earendil-works/pi-ai": "0.84.0",
-    "@earendil-works/pi-coding-agent": "0.84.0",
-    "@earendil-works/pi-tui": "0.84.0",
+    "@earendil-works/pi-ai": "1.1.0",
+    "@earendil-works/pi-coding-agent": "1.1.0",
+    "@earendil-works/pi-tui": "1.1.0",
     "@types/node": "^25.0.10",
     "husky": "^9.1.7",
     "ts-json-schema-generator": "^2.4.0",
     "tsx": "4.23.12",
+    "typebox": "1.3.27",
     "typescript": "^5.9.3",
     "vitest": "^4.0.18"
   },

package/schema.json

@@ -76,6 +76,10 @@
         "api": {
           "$ref": "#/definitions/NeuralwattApi",
           "description": "Which API serves model requests."
+        },
+        "apiBaseUrl": {
+          "type": "string",
+          "description": "Fallback base URL for Neuralwatt API calls, used only by models that carry no baseUrl of their own (models.json, custom providers, and proxy routing all take precedence) and by the extension's own catalog and quota fetches. Defaults to the direct upstream."
         }
       },
       "additionalProperties": false

package/src/config/defaults.ts

@@ -1,5 +1,7 @@
 import type { ResolvedNeuralwattConfig } from "./types";
 
+export const NEURALWATT_API_BASE_URL = "https://api.neuralwatt.com/v1";
+
 export const DEFAULT_CONFIG: ResolvedNeuralwattConfig = {
   quotaCommand: {
     enabled: true,
@@ -12,5 +14,6 @@ export const DEFAULT_CONFIG: ResolvedNeuralwattConfig = {
   },
   provider: {
     api: "openai-completions",
+    apiBaseUrl: NEURALWATT_API_BASE_URL,
   },
 };

package/src/config/loader.ts

@@ -33,6 +33,7 @@ function normalizeResolvedConfig(
     },
     provider: {
       api: resolveApi(config.provider?.api),
+      apiBaseUrl: resolveApiBaseUrl(config.provider?.apiBaseUrl),
     },
   };
 }
@@ -44,6 +45,30 @@ export function resolveApi(value: string | undefined): NeuralwattApi {
   return DEFAULT_CONFIG.provider.api;
 }
 
+export function resolveApiBaseUrl(value: string | undefined): string {
+  const trimmed = value?.trim().replace(/\/+$/u, "");
+  if (!trimmed) {
+    return DEFAULT_CONFIG.provider.apiBaseUrl;
+  }
+  try {
+    const url = new URL(trimmed);
+    if (url.protocol !== "https:" && url.protocol !== "http:") {
+      return DEFAULT_CONFIG.provider.apiBaseUrl;
+    }
+  } catch {
+    return DEFAULT_CONFIG.provider.apiBaseUrl;
+  }
+  return trimmed;
+}
+
+export function configuredApiBaseUrl(): string {
+  try {
+    return configLoader.getConfig().provider.apiBaseUrl;
+  } catch {
+    return DEFAULT_CONFIG.provider.apiBaseUrl;
+  }
+}
+
 export const configLoader = new ConfigLoader<
   NeuralwattConfig,
   ResolvedNeuralwattConfig

package/src/config/types.ts

@@ -23,6 +23,14 @@ export type NeuralwattApi = "openai-completions" | "anthropic-messages";
 export interface NeuralwattProviderConfig {
   /** Which API serves model requests. */
   api?: NeuralwattApi;
+
+  /**
+   * Fallback base URL for Neuralwatt API calls, used only by models that
+   * carry no baseUrl of their own (models.json, custom providers, and proxy
+   * routing all take precedence) and by the extension's own catalog and
+   * quota fetches. Defaults to the direct upstream.
+   */
+  apiBaseUrl?: string;
 }
 
 export interface NeuralwattConfig {
@@ -54,5 +62,6 @@ export interface ResolvedNeuralwattConfig {
   };
   provider: {
     api: NeuralwattApi;
+    apiBaseUrl: string;
   };
 }

package/src/lib/neuralwatt-api.ts

@@ -1,3 +1,4 @@
+import { configuredApiBaseUrl } from "../config/loader";
 import type {
   NeuralwattApiModel,
   NeuralwattApiModelsResponse,
@@ -5,7 +6,6 @@ import type {
 import type { NeuralwattQuotas } from "../types/quota-api";
 import type { QuotasResult } from "../types/quota-result";
 
-const BASE_URL = "https://api.neuralwatt.com/v1";
 const FETCH_TIMEOUT_MS = 15_000;
 
 function authHeaders(apiKey: string | undefined): Record<string, string> {
@@ -31,7 +31,7 @@ async function neuralwattFetch(
   signal?: AbortSignal,
   headers?: Record<string, string>,
 ): Promise<Response> {
-  return fetch(`${BASE_URL}${path}`, {
+  return fetch(`${configuredApiBaseUrl()}${path}`, {
     headers: { ...authHeaders(apiKey), ...headers },
     signal: combineSignals(signal),
   });