SPB Git forge

spb/polyllm

Public
15commits 1branches 0releases
2.2 MBsize
maindefault branch
13 days agolast push
TypeScript 97.4% SQL 1% JavaScript 0.9% CSS 0.6%
10.5 KB · 263 lines typescript
Raw Blame History
1import type OpenAI from "openai";2import { type AIProviderAdapter, type PolyModel, PolyProviderError } from "@/lib/ai/core/types";3import { createOpenAICompatAdapter } from "../shared/openai-compat/factory";45/**6 * Custom OpenAI-compatible endpoints (Ollama, LM Studio, vLLM, llama.cpp server, MLX-LM, TGI,7 * LocalAI, any `/v1/chat/completions` server). One adapter instance per endpoint, built at request8 * time from the user's stored configuration (`@/lib/endpoints/service`).9 *10 * Model keys are `custom/<endpointId>:<modelId>` — the provider prefix stays `custom` so every11 * `parseModelKey()` call site keeps working; the endpoint id and the (possibly `:`-containing)12 * model id are split by `parseCustomModelKey()` on the first `:` after the prefix.13 */1415export const CUSTOM_KEY_PREFIX = "custom/";16/** Sent as `Authorization: Bearer …` when the endpoint has no key (Ollama & co. ignore it). */17export const NO_KEY_SENTINEL = "polyllm-no-key";1819export interface CustomEndpointConfig {20  id: string;21  name: string;22  /** Up to and including `/v1` (e.g. http://localhost:11434/v1). No trailing slash. */23  baseUrl: string;24  apiKey?: string | null;25  headers?: Record<string, string> | null;26  /** Relative path appended to `baseUrl` for discovery (default `/models`). Empty = discovery disabled. */27  modelsPath?: string | null;28  /** Manual declarations (also act as capability overrides for discovered ids). */29  manualModels?: ManualModel[] | null;30  /** Discovery request timeout (ms). */31  timeoutMs?: number;32}3334export interface ManualModel {35  id: string;36  displayName?: string;37  contextTokens?: number;38  vision?: boolean;39  tools?: boolean;40  reasoning?: boolean;41}4243export function customModelKey(endpointId: string, modelId: string): string {44  return `${CUSTOM_KEY_PREFIX}${endpointId}:${modelId}`;45}4647/** `custom/cep_abc:llama3.1:8b` → `{ endpointId: "cep_abc", modelId: "llama3.1:8b" }`. */48export function parseCustomModelKey(key: string): { endpointId: string; modelId: string } | null {49  if (!key.startsWith(CUSTOM_KEY_PREFIX)) return null;50  const rest = key.slice(CUSTOM_KEY_PREFIX.length);51  const idx = rest.indexOf(":");52  if (idx <= 0 || idx === rest.length - 1) return null;53  return { endpointId: rest.slice(0, idx), modelId: rest.slice(idx + 1) };54}5556export function isCustomModelKey(key: string): boolean {57  return parseCustomModelKey(key) !== null;58}5960export function normalizeBaseUrl(url: string): string {61  return url.trim().replace(/\/+$/, "");62}6364export function normalizeModelsPath(p: string | null | undefined): string {65  const t = (p ?? "").trim();66  if (!t) return "";67  return t.startsWith("/") ? t : `/${t}`;68}6970export interface DiscoveredModel {71  id: string;72  owned_by?: string;73  created?: number;74  /** Some servers (LM Studio, llama.cpp) add metadata; kept opaque. */75  [k: string]: unknown;76}7778/** Build a PolyModel for one model id on an endpoint; `manual` supplies capability overrides. */79export function toCustomPolyModel(cfg: Pick<CustomEndpointConfig, "id" | "name">, modelId: string, manual?: ManualModel, extra?: { ownedBy?: string; created?: number; source: "discovery" | "manual" }): PolyModel {80  return {81    key: customModelKey(cfg.id, modelId),82    id: modelId,83    provider: "custom",84    displayName: manual?.displayName?.trim() || modelId,85    family: cfg.name,86    capabilities: {87      text: true,88      vision: manual?.vision ?? false,89      audioInput: false,90      audioOutput: false,91      imageGeneration: false,92      video: false,93      reasoning: manual?.reasoning ?? false,94      tools: manual?.tools ?? false,95      structuredOutput: false,96      streaming: true,97      files: false,98      webSearch: false,99    },100    limits: manual?.contextTokens ? { contextTokens: manual.contextTokens } : {},101    parameters: {102      temperature: true,103      topP: true,104      maxTokens: true,105      stop: true,106      seed: true,107      frequencyPenalty: true,108      presencePenalty: true,109      temperatureRange: { min: 0, max: 2 },110    },111    status: "active",112    // Cost is unknown for arbitrary endpoints (local = free, hosted = billed elsewhere): never invent a price.113    pricing: null,114    metadata: {115      endpointId: cfg.id,116      endpointName: cfg.name,117      ownedBy: extra?.ownedBy,118      createdAt: extra?.created ? new Date(extra.created * 1000).toISOString() : undefined,119      source: extra?.source ?? "manual",120      sortWeight: -40,121    },122  };123}124125/** Merge discovered ids with manual declarations (manual entries win on capabilities; unknown ids are appended). */126export function mergeCustomModels(cfg: Pick<CustomEndpointConfig, "id" | "name">, discovered: DiscoveredModel[], manual: ManualModel[] = []): PolyModel[] {127  const manualById = new Map(manual.filter((m) => m.id?.trim()).map((m) => [m.id.trim(), m]));128  const seen = new Set<string>();129  const out: PolyModel[] = [];130  for (const d of discovered) {131    if (!d?.id || typeof d.id !== "string" || seen.has(d.id)) continue;132    seen.add(d.id);133    out.push(toCustomPolyModel(cfg, d.id, manualById.get(d.id), { ownedBy: typeof d.owned_by === "string" ? d.owned_by : undefined, created: typeof d.created === "number" ? d.created : undefined, source: "discovery" }));134  }135  for (const m of manual) {136    const id = m.id?.trim();137    if (!id || seen.has(id)) continue;138    seen.add(id);139    out.push(toCustomPolyModel(cfg, id, m, { source: "manual" }));140  }141  return out;142}143144function authHeaders(cfg: CustomEndpointConfig): Record<string, string> {145  const h: Record<string, string> = { Accept: "application/json", ...(cfg.headers ?? {}) };146  if (cfg.apiKey && cfg.apiKey !== NO_KEY_SENTINEL) h.Authorization = `Bearer ${cfg.apiKey}`;147  return h;148}149150/** `GET {baseUrl}{modelsPath}` — accepts OpenAI `{data:[{id}]}`, bare arrays and Ollama-style `{models:[{name}]}` bodies. */151export async function discoverCustomModels(cfg: CustomEndpointConfig, signal?: AbortSignal): Promise<DiscoveredModel[]> {152  const path = normalizeModelsPath(cfg.modelsPath ?? "/models");153  if (!path) return [];154  const url = `${normalizeBaseUrl(cfg.baseUrl)}${path}`;155  const ctrl = new AbortController();156  const timer = setTimeout(() => ctrl.abort(), cfg.timeoutMs ?? 8_000);157  signal?.addEventListener("abort", () => ctrl.abort(), { once: true });158  try {159    const res = await fetch(url, { headers: authHeaders(cfg), signal: ctrl.signal, redirect: "manual" });160    if (!res.ok) {161      const text = await res.text().catch(() => "");162      let message = `HTTP ${res.status}`;163      try {164        const body = JSON.parse(text) as { error?: { message?: string } | string; message?: string };165        message = (typeof body.error === "object" ? body.error?.message : typeof body.error === "string" ? body.error : undefined) ?? body.message ?? message;166      } catch {167        if (text) message = `${message}: ${text.slice(0, 200)}`;168      }169      throw Object.assign(new Error(message), { status: res.status, headers: res.headers });170    }171    const body = (await res.json()) as unknown;172    return parseModelListing(body);173  } finally {174    clearTimeout(timer);175  }176}177178export function parseModelListing(body: unknown): DiscoveredModel[] {179  const pick = (arr: unknown[]): DiscoveredModel[] =>180    arr181      .map((m) => {182        if (typeof m === "string") return { id: m };183        if (m && typeof m === "object") {184          const o = m as Record<string, unknown>;185          const id = typeof o.id === "string" ? o.id : typeof o.name === "string" ? o.name : typeof o.model === "string" ? o.model : null;186          return id ? { ...o, id } : null;187        }188        return null;189      })190      .filter((m): m is DiscoveredModel => Boolean(m));191  if (Array.isArray(body)) return pick(body);192  if (body && typeof body === "object") {193    const o = body as Record<string, unknown>;194    if (Array.isArray(o.data)) return pick(o.data);195    if (Array.isArray(o.models)) return pick(o.models);196  }197  return [];198}199200/**201 * Adapter for one endpoint. Chat goes through the shared Chat Completions implementation with the202 * endpoint's base URL and headers; discovery uses `discoverCustomModels`.203 */204export function createCustomEndpointAdapter(cfg: CustomEndpointConfig): AIProviderAdapter {205  const listModels = async (_client: OpenAI, _apiKey: string, signal?: AbortSignal): Promise<PolyModel[]> => {206    const discovered = await discoverCustomModels(cfg, signal);207    return mergeCustomModels(cfg, discovered, cfg.manualModels ?? []);208  };209  return createOpenAICompatAdapter({210    id: "custom",211    name: cfg.name,212    baseURL: normalizeBaseUrl(cfg.baseUrl),213    keyDocsUrl: "",214    defaultHeaders: cfg.headers ?? undefined,215    listModels,216    // Most local servers only understand `max_tokens`; the newer `max_completion_tokens` 400s on several of them.217    useMaxTokens: true,218    messageOptions: { inlineFiles: true },219    tweakParams: (params, settings, req) => {220      const p = params as unknown as Record<string, unknown>;221      // Endpoints marked "reasoning" by the user follow the de-facto `reasoning_effort` field (vLLM, llama.cpp, LM Studio).222      if (req.modelInfo?.capabilities.reasoning && settings.reasoningEffort && settings.reasoningEffort !== "none") p.reasoning_effort = settings.reasoningEffort === "minimal" ? "low" : settings.reasoningEffort;223    },224    refineError: (status, _code, message) => {225      if (status === 404 && /model/i.test(message)) return "MODEL_NOT_FOUND";226      if (status === 401 || status === 403) return "INVALID_API_KEY";227      if (status === 503 || status === 502) return "PROVIDER_UNAVAILABLE";228      return undefined;229    },230    timeoutMs: 10 * 60_000,231  });232}233234/**235 * Registered under `ADAPTERS.custom` so `getAdapter("custom")` never throws. It cannot talk to any236 * server: real requests must go through `resolveCustomEndpoint(userId, modelKey)` which builds a237 * per-endpoint adapter. Every method fails loudly with an actionable message.238 */239export const customPlaceholderAdapter: AIProviderAdapter = {240  id: "custom",241  name: "Custom endpoint",242  keyDocsUrl: "/app/settings/endpoints",243  async validateApiKey() {244    return { ok: false, latencyMs: 0, error: unresolved().toJSON() };245  },246  async listModels() {247    return [];248  },249  async chat() {250    throw unresolved();251  },252  async *streamChat() {253    yield { type: "error", error: unresolved().toJSON() };254  },255  normalizeError(error: unknown) {256    return error instanceof PolyProviderError ? error : unresolved();257  },258};259260function unresolved(): PolyProviderError {261  return new PolyProviderError({ code: "MODEL_NOT_FOUND", message: "Custom endpoint not resolved. Route this request through resolveCustomEndpoint().", provider: "custom", retryable: false });262}263