pi-provider-litellm 4.2.0 → 4.3.1

package/src/models-store.ts

@@ -0,0 +1,117 @@
+/**
+ * Reads Pi's persisted catalog, so that a start does not re-prove what Pi already kept.
+ *
+ * The factory discovers before it returns because Pi's startup refresh runs with `allowNetwork:false` and returns before the
+ * network phase, so a provider that waits for Pi has no models when the session picks one (issue #137). That holds on a first
+ * run. It does not hold once Pi has persisted the catalog: with discovery switched off and a populated `models-store.json`,
+ * Pi offers the stored models and a request works, measured at a 5 ms factory against 1.8 s to discover.
+ *
+ * So the seed reads the stored catalog and hands those models over, instead of discovering them again. It only reads, and
+ * only `models-store.json`, which this repository's AGENTS.md assigns to Pi. Nothing here writes a model cache; the removed
+ * `litellm-models.json` is what that rule is about.
+ *
+ * The models are returned rather than an empty list on purpose. Pi restores its own catalog, so an empty list would do in
+ * the ordinary case. It would not do when Pi's active store is somewhere else: `ModelRuntime` takes a `modelsStorePath`, and
+ * resolves the default beside whatever `modelsPath` it was given. Every caller inside Pi passes an agent-directory path, but
+ * an SDK embedder need not. A catalog left in the default file would then satisfy this check while Pi's active store was
+ * empty, and the provider would have no models at all -- the very fault the seed exists to prevent. Returning the models
+ * costs nothing and cannot fail that way.
+ *
+ * Four things make a stored entry unusable, and all are checked here rather than left to Pi:
+ *
+ *   - Age. Only a network refresh writes the store, and `pi -p` never runs one, so without a limit print mode would keep
+ *     whatever catalog an interactive session last left. Pi's own remote catalogs skip the network for 4 hours after
+ *     `checkedAt`, and so does this.
+ *   - A different proxy. Stored models carry the `baseUrl` they were discovered from, so a changed root must still discover.
+ *   - An older `litellmDiscoveryVersion`. Pi keeps a same-version stored entry over what startup discovery just proved
+ *     (AGENTS.md), so skipping discovery for an older entry would hide the metadata a new version adds.
+ *   - An entry Pi cannot restore. Pi's no-network restore publishes the stored models over the seed, so the seed cannot keep
+ *     a short entry out of a session; it decides the catalog only when that restore fails. The restore runs
+ *     `enrichCachedModel` (provider.ts), and when that throws Pi restores nothing, so the seed makes the same call and
+ *     discovers instead.
+ *
+ * What is deliberately not checked is the credential. Nothing in Pi's store records which key discovered it, and a key scopes
+ * which models a LiteLLM proxy shows. Within those 4 hours, a key swapped on the same proxy, or a model added to it, therefore
+ * shows the stored catalog until Pi's own network refresh runs: interactive and RPC sessions run one in the background,
+ * `pi -p` does not. Recording a key fingerprint would mean writing a file of our own, which is the thing the rule forbids.
+ */
+
+import { readFile } from "node:fs/promises";
+import { join } from "node:path";
+import { LITELLM_DISCOVERY_VERSION } from "./backend-identity.js";
+import { enrichCachedModel } from "./discover.js";
+import { LITELLM_IMAGES_API, type LiteLLMImageModel } from "./images.js";
+import { isLiteLLMApi, resolveModelBaseUrl } from "./protocols.js";
+import { imageBaseUrl } from "./provider.js";
+import type { LiteLLMModel } from "./types.js";
+
+/**
+ * Where Pi keeps it. `FileModelsStore` defaults to `<agentDir>/models-store.json`, and `ModelRuntime` resolves it beside
+ * whatever `modelsPath` it was given -- `modelsPath` itself is the model *config* (`models.json`), a different file with a
+ * different shape, which is deliberately not read here.
+ */
+const STORE_FILENAME = "models-store.json";
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+/** Pi's own remote catalogs skip the network for 4 hours after `checkedAt` (`REMOTE_CATALOG_REFRESH_INTERVAL_MS`). */
+const STORE_FRESH_MS = 4 * 60 * 60 * 1000;
+
+function isImageModel(value: Record<string, unknown>): boolean {
+  return value.type === "image" && value.api === LITELLM_IMAGES_API;
+}
+
+function isUsableModel(value: unknown, provider: string, baseUrl: string, allowInsecureHttp?: boolean): boolean {
+  if (!isRecord(value) || value.provider !== provider) return false;
+  if (value.litellmDiscoveryVersion !== LITELLM_DISCOVERY_VERSION) return false;
+  try {
+    // A stored baseUrl is per protocol -- `/v1` for the OpenAI APIs and images, the bare root for Messages -- so the root
+    // goes through the same resolution rather than being compared to it.
+    if (isImageModel(value)) return value.baseUrl === imageBaseUrl(baseUrl, allowInsecureHttp);
+    // `api` indexes LITELLM_PROTOCOLS in resolveModelBaseUrl, so an unknown one is a TypeError, not a bad model.
+    return isLiteLLMApi(value.api) && value.baseUrl === resolveModelBaseUrl(baseUrl, value.api, allowInsecureHttp);
+  } catch {
+    return false;
+  }
+}
+
+async function readStore(agentDir: string): Promise<Record<string, unknown> | undefined> {
+  try {
+    const parsed: unknown = JSON.parse(await readFile(join(agentDir, STORE_FILENAME), "utf8"));
+    return isRecord(parsed) ? parsed : undefined;
+  } catch {
+    // Absent, unreadable, or not JSON: discover instead.
+    return undefined;
+  }
+}
+
+/**
+ * Pi's stored models for `provider`, when every one of them can be used as it is, and `undefined` otherwise. Never throws;
+ * anything it cannot read means "discover".
+ */
+export async function readStoredModels(options: {
+  agentDir: string;
+  provider: string;
+  baseUrl: string;
+  allowInsecureHttp?: boolean;
+}): Promise<(LiteLLMModel | LiteLLMImageModel)[] | undefined> {
+  const entry = (await readStore(options.agentDir))?.[options.provider];
+  if (!isRecord(entry) || !Array.isArray(entry.models) || entry.models.length === 0) return undefined;
+  if (typeof entry.checkedAt !== "number" || Date.now() - entry.checkedAt >= STORE_FRESH_MS) return undefined;
+  const usable = entry.models.every((model) =>
+    isUsableModel(model, options.provider, options.baseUrl, options.allowInsecureHttp),
+  );
+  if (!usable) return undefined;
+  try {
+    // Pi's restore path (provider.ts) runs the same call; when it throws, nothing is restored and the seed is the catalog.
+    // Image models carry no request policy, so that path restores them as stored.
+    return entry.models.map((model) => (isImageModel(model) ? model : enrichCachedModel(model))) as (
+      | LiteLLMModel
+      | LiteLLMImageModel
+    )[];
+  } catch {
+    return undefined;
+  }
+}

package/src/provider.ts

@@ -2,9 +2,14 @@ import {
   type AnyModel,
   type Api,
   type ApiStreamOptions,
+  type AssistantImages,
   type Credential,
   createProvider,
   getCurrentTools,
+  type ImageApi,
+  type ImageModel,
+  type ImagesContext,
+  type ImagesOptions,
   type Model,
   type Provider,
   type ProviderAuth,
@@ -13,6 +18,12 @@ import {
 } from "@earendil-works/pi-ai";
 import { LITELLM_DISCOVERY_VERSION } from "./backend-identity.js";
 import { enrichCachedModel, normalizeBaseUrl, restoreCachedModelPolicy } from "./discover.js";
+import {
+  createLiteLLMImageApis,
+  type DiscoveredImageModel,
+  LITELLM_IMAGES_API,
+  type LiteLLMImageModel,
+} from "./images.js";
 import { createLiteLLMProtocolApis, isLiteLLMApi, LITELLM_API_NAMES, resolveModelBaseUrl } from "./protocols.js";
 import type { DiscoveredModel, DiscoveryResult, LiteLLMApi, LiteLLMModel } from "./types.js";
 
@@ -23,7 +34,7 @@ export type LiteLLMProviderOptions = {
   auth: ProviderAuth;
   /** Catalog discovered during activation. Pi's startup refresh never allows network access, so
    * without a seed the provider stays empty until the user opens /model. */
-  models?: readonly LiteLLMModel[];
+  models?: readonly (LiteLLMModel | LiteLLMImageModel)[];
   allowInsecureHttp?: boolean;
   resolveCredentialRoot: (credential?: Credential, requestBaseUrl?: string, apiKey?: string) => string | undefined;
   discover(credential: Credential, signal?: AbortSignal): Promise<DiscoveryResult & { baseUrl?: string }>;
@@ -34,6 +45,23 @@ function isChatModel(model: AnyModel): model is Model<Api> {
   return (model.type ?? "chat") === "chat";
 }
 
+function isLiteLLMImageModel(model: AnyModel): model is LiteLLMImageModel {
+  return model.type === "image" && model.api === LITELLM_IMAGES_API;
+}
+
+export function imageBaseUrl(root: string, allowInsecureHttp = false): string {
+  return `${normalizeBaseUrl(root, allowInsecureHttp)}/v1`;
+}
+
+export function toNativeImageModels(
+  provider: string,
+  baseUrl: string,
+  models: readonly DiscoveredImageModel[],
+  allowInsecureHttp = false,
+): LiteLLMImageModel[] {
+  return models.map((model) => ({ ...model, provider, baseUrl: imageBaseUrl(baseUrl, allowInsecureHttp) }));
+}
+
 export function toNativeModels(
   provider: string,
   baseUrl: string,
@@ -117,6 +145,27 @@ function modelRootError(
   }
 }
 
+function imageRootError(
+  model: ImageModel<ImageApi>,
+  active: { root: string; canonical: string; hostname: string },
+  allowInsecureHttp = false,
+): Error | undefined {
+  let stored: { root: string; canonical: string; hostname: string };
+  try {
+    stored = proxyRoot(model.baseUrl, "Cached image model", allowInsecureHttp);
+  } catch (error) {
+    return error instanceof Error ? error : refreshRequired("Cached image model has an invalid LiteLLM model URL");
+  }
+  if (isPlaceholderHost(stored.hostname)) {
+    return refreshRequired("Cached image model uses a placeholder LiteLLM model host");
+  }
+  if (stored.canonical !== active.canonical) {
+    return refreshRequired(
+      `Cached image model has stale LiteLLM model root ${stored.root}; active credentials use ${active.root}`,
+    );
+  }
+}
+
 function requestModel(
   provider: string,
   model: Model<LiteLLMApi>,
@@ -148,6 +197,40 @@ export function createLiteLLMProvider(options: LiteLLMProviderOptions): Provider
     reportedAvailabilityDiagnostics.add(message);
     process.stderr.write(`LiteLLM (${options.id}): ${message}\n`);
   };
+  type ActiveRoot = { root: string; canonical: string; hostname: string };
+  const resolveActiveRoot = (credential: Credential | undefined, hiddenCount: number): ActiveRoot | undefined => {
+    try {
+      const root = options.resolveCredentialRoot(credential);
+      if (!root) {
+        if (hiddenCount > 0) {
+          reportUnavailable(
+            `${hiddenCount} model(s) hidden because no LiteLLM base URL is configured; ` +
+              "set LITELLM_BASE_URL or run /login litellm",
+          );
+        }
+        return undefined;
+      }
+      return activeCredentialRoot(root, options.allowInsecureHttp);
+    } catch (error) {
+      reportUnavailable(error instanceof Error ? error.message : String(error));
+      return undefined;
+    }
+  };
+  const availableChat = (models: readonly Model<LiteLLMApi>[], active: ActiveRoot): Model<LiteLLMApi>[] => {
+    const available: Model<LiteLLMApi>[] = [];
+    for (const model of models) {
+      const error = modelRootError(model, active, options.allowInsecureHttp);
+      if (error) {
+        reportUnavailable(error.message);
+        continue;
+      }
+      available.push({
+        ...model,
+        baseUrl: resolveModelBaseUrl(active.root, model.api, options.allowInsecureHttp),
+      });
+    }
+    return available;
+  };
   const provider = createProvider<LiteLLMApi>({
     id: options.id,
     name: options.name,
@@ -157,41 +240,34 @@ export function createLiteLLMProvider(options: LiteLLMProviderOptions): Provider
     async fetchModels(context) {
       if (!context.credential) throw new Error("LiteLLM model discovery requires a credential");
       const result = await options.discover(context.credential, context.signal);
-      return toNativeModels(options.id, result.baseUrl ?? options.baseUrl, result.models, options.allowInsecureHttp);
+      const root = result.baseUrl ?? options.baseUrl;
+      return [
+        ...toNativeModels(options.id, root, result.models, options.allowInsecureHttp),
+        ...toNativeImageModels(options.id, root, result.imageModels ?? [], options.allowInsecureHttp),
+      ];
     },
     filterModels(models, credential) {
-      let active: { root: string; canonical: string; hostname: string };
-      try {
-        const root = options.resolveCredentialRoot(credential);
-        if (!root) {
-          if (models.length > 0) {
-            reportUnavailable(
-              `${models.length} model(s) hidden because no LiteLLM base URL is configured; ` +
-                "set LITELLM_BASE_URL or run /login litellm",
-            );
-          }
-          return [];
-        }
-        active = activeCredentialRoot(root, options.allowInsecureHttp);
-      } catch (error) {
-        reportUnavailable(error instanceof Error ? error.message : String(error));
-        return [];
-      }
-      const available: Model<LiteLLMApi>[] = [];
-      for (const model of models) {
-        const error = modelRootError(model, active, options.allowInsecureHttp);
+      const active = resolveActiveRoot(credential, models.length);
+      return active ? availableChat(models, active) : [];
+    },
+    filterAllModels(models, credential) {
+      const active = resolveActiveRoot(credential, models.length);
+      if (!active) return [];
+      const chat = availableChat(models.filter(isChatModel) as Model<LiteLLMApi>[], active);
+      // Image models follow the same rule as chat: a cached root that no longer matches the
+      // active credential is hidden until a network refresh proves the new proxy's catalog.
+      const images = models.filter(isLiteLLMImageModel).flatMap((model) => {
+        const error = imageRootError(model, active, options.allowInsecureHttp);
         if (error) {
           reportUnavailable(error.message);
-          continue;
+          return [];
         }
-        available.push({
-          ...model,
-          baseUrl: resolveModelBaseUrl(active.root, model.api, options.allowInsecureHttp),
-        });
-      }
-      return available;
+        return [{ ...model, baseUrl: imageBaseUrl(active.root, options.allowInsecureHttp) }];
+      });
+      return [...chat, ...images];
     },
     api: createLiteLLMProtocolApis(),
+    images: createLiteLLMImageApis(),
   });
   const refreshModels = provider.refreshModels;
   const guardedProvider: Provider<LiteLLMApi> = {
@@ -220,6 +296,38 @@ export function createLiteLLMProvider(options: LiteLLMProviderOptions): Provider
         context,
         requestOptions,
       ),
+    // Same request-time guard as chat: the credential's root, never a cached model URL, picks the
+    // host, so an off-catalog or stale entry cannot send this credential to another proxy.
+    generateImages: async (model: ImageModel<ImageApi>, context: ImagesContext, requestOptions?: ImagesOptions) => {
+      try {
+        const root = options.resolveCredentialRoot(
+          undefined,
+          requestOptions?.env?.LITELLM_BASE_URL,
+          requestOptions?.apiKey,
+        );
+        if (!root) throw refreshRequired("Active credentials do not identify a LiteLLM model host");
+        const active = activeCredentialRoot(root, options.allowInsecureHttp);
+        const error = imageRootError(model, active, options.allowInsecureHttp);
+        if (error) throw error;
+        return await provider.generateImages!(
+          { ...model, provider: options.id, baseUrl: imageBaseUrl(active.root, options.allowInsecureHttp) },
+          context,
+          requestOptions,
+        );
+      } catch (error) {
+        // `Provider.generateImages` never rejects.
+        const failed: AssistantImages = {
+          api: model.api,
+          provider: model.provider,
+          model: model.id,
+          output: [],
+          stopReason: requestOptions?.signal?.aborted ? "aborted" : "error",
+          errorMessage: error instanceof Error ? error.message : String(error),
+          timestamp: Date.now(),
+        };
+        return failed;
+      }
+    },
   };
   if (!refreshModels) return guardedProvider;
   const seedKey = (model: { id: string; baseUrl: string }) => `${model.id}\0${model.baseUrl}`;
@@ -227,16 +335,23 @@ export function createLiteLLMProvider(options: LiteLLMProviderOptions): Provider
   return {
     ...guardedProvider,
     refreshModels: async (context) => {
-      // Pi 0.99 persists models of every type; this provider only publishes chat models, and the
-      // policy helpers below read chat-only fields.
+      // Pi 0.99 persists models of every type; the policy helpers below read chat-only fields, so
+      // chat and image entries are restored separately. Other types are not this provider's.
       const storedModels = (context.stored?.models ?? []).filter(isChatModel);
+      const storedImageModels = (context.stored?.models ?? []).filter(isLiteLLMImageModel).flatMap((model) => {
+        if (model.litellmDiscoveryVersion === LITELLM_DISCOVERY_VERSION) return [model];
+        // An image entry carries no request policy to restore; a stale one is only kept when
+        // startup discovery proved the same route, otherwise the network refresh replaces it.
+        const fresh = seeded.get(seedKey(model));
+        return fresh && isLiteLLMImageModel(fresh) ? [fresh] : [];
+      });
       let legacyCount = 0;
       const models = storedModels.map((model) => {
         if ((model as LiteLLMModel).litellmDiscoveryVersion !== LITELLM_DISCOVERY_VERSION) {
           // Pi applies stored models over the seed, so a stale entry would otherwise mask what
           // startup discovery just proved for the same route.
           const fresh = seeded.get(seedKey(model));
-          if (fresh) return fresh;
+          if (fresh && isChatModel(fresh)) return fresh;
           legacyCount++;
           return restoreCachedModelPolicy(model);
         }
@@ -247,7 +362,7 @@ export function createLiteLLMProvider(options: LiteLLMProviderOptions): Provider
           ...context,
           // The current pi-ai provider has no freshness gate, but force is the documented refresh contract field.
           force: context.force || (context.allowNetwork && legacyCount > 0),
-          stored: context.stored && { ...context.stored, models },
+          stored: context.stored && { ...context.stored, models: [...models, ...storedImageModels] },
         });
       } catch (error) {
         if (context.allowNetwork && legacyCount > 0) {

package/src/types.ts

@@ -1,5 +1,6 @@
 import type { Model } from "@earendil-works/pi-ai";
 import type { BackendFamily, LITELLM_DISCOVERY_VERSION } from "./backend-identity.js";
+import type { DiscoveredImageModel } from "./images.js";
 import type { ProxyVersion } from "./proxy-version.js";
 
 export type DiscoverySource = "model_info" | "models_list" | "health";
@@ -64,6 +65,8 @@ export type LiteLLMModel = Model<LiteLLMApi> & {
 
 export interface DiscoveryResult {
   models: DiscoveredModel[];
+  /** `mode: "image_generation"` routes from `/model/info`; absent when there are none. */
+  imageModels?: DiscoveredImageModel[];
   source: DiscoverySource;
   // Present only when a `/model/info` discovery had to read it and the proxy
   // answered. Absent means unknown, not old: discovery asks only when it must.
@@ -98,6 +101,7 @@ export interface ModelInfoEntry {
     supported_openai_params?: string[];
     input_cost_per_token?: number;
     output_cost_per_token?: number;
+    output_cost_per_image_token?: number;
     cache_read_input_token_cost?: number;
     cache_creation_input_token_cost?: number;
     max_input_tokens?: number;