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%
17.4 KB · 345 lines typescript
Raw Blame History
1import { GoogleGenAI, ApiError, type Content, type Part, type GenerateContentConfig, type GenerateContentResponse, type Model as GeminiModel, type Tool, type FunctionDeclaration, FunctionCallingConfigMode, type ThinkingConfig } from "@google/genai";2import {3  type AIProviderAdapter,4  type PolyModel,5  type UnifiedChatRequest,6  type UnifiedChatResponse,7  type UnifiedStreamEvent,8  type ValidationResult,9  type TokenEstimate,10  type UnifiedMessage,11  type FinishReason,12  PolyProviderError,13  modelKey,14} from "@/lib/ai/core/types";15import { normalizeGenericError, refineByMessage, codeFromStatus, isRetryableStatus } from "@/lib/ai/core/errors";16import { filterSettings, heuristicTokens } from "@/lib/ai/core/normalize";17import { collectStream } from "@/lib/ai/core/stream-utils";18import { isTextLike, inlineTextFile } from "@/lib/ai/core/content";19import { GEMINI_CATALOG, GEMINI_NON_CHAT } from "./catalog";2021const DEFAULT_TIMEOUT_MS = 10 * 60_000;2223function client(apiKey: string, timeoutMs = DEFAULT_TIMEOUT_MS) {24  return new GoogleGenAI({ apiKey, httpOptions: { timeout: timeoutMs } });25}2627export function normalizeGeminiModel(m: GeminiModel): PolyModel | null {28  const id = (m.name ?? "").replace(/^models\//, "");29  if (!id || GEMINI_NON_CHAT.test(id)) return null;30  const actions = m.supportedActions ?? [];31  if (actions.length && !actions.includes("generateContent")) return null;32  const cat = GEMINI_CATALOG.get(id);33  const thinking = Boolean(m.thinking);34  const gone = /^gemini-(1|2)\./.test(id);35  return {36    key: modelKey("gemini", id),37    id,38    provider: "gemini",39    displayName: cat?.displayName ?? m.displayName ?? id,40    family: cat?.family ?? "Gemini",41    capabilities: cat?.capabilities ?? { text: true, vision: true, audioInput: true, audioOutput: false, imageGeneration: false, video: true, reasoning: thinking, tools: true, structuredOutput: true, streaming: true, files: true, webSearch: true },42    limits: { contextTokens: m.inputTokenLimit ?? cat?.limits?.contextTokens, maxOutputTokens: m.outputTokenLimit ?? cat?.limits?.maxOutputTokens },43    parameters: cat?.parameters ?? { temperature: true, topP: true, topK: true, seed: true, stop: true, maxTokens: true, frequencyPenalty: false, presencePenalty: false, reasoningEffort: thinking, reasoningEffortLevels: thinking ? ["low", "medium", "high"] : undefined, thinkingBudget: false, temperatureRange: { min: 0, max: Math.min(2, m.maxTemperature ?? 2) } },44    status: gone ? "deprecated" : cat?.status ?? (id.includes("preview") ? "preview" : "unknown"),45    pricing: cat?.pricing ?? null,46    metadata: { thinkingProfile: thinking ? { mode: "level", levels: ["low", "medium", "high"], none: null } : { mode: "none" }, ...(cat?.metadata ?? {}), version: m.version, description: m.description, defaultTemperature: m.temperature, sortWeight: cat?.sortWeight ?? 0 },47  };48}4950// ---------------------------------------------------------------------------51// Request translation52// ---------------------------------------------------------------------------53function toContents(messages: UnifiedMessage[]): Content[] {54  const out: Content[] = [];55  const push = (role: "user" | "model", parts: Part[]) => {56    if (!parts.length) return;57    const last = out[out.length - 1];58    if (last && last.role === role) last.parts = [...(last.parts ?? []), ...parts];59    else out.push({ role, parts });60  };61  for (const m of messages) {62    if (m.role === "system") continue;63    if (m.role === "tool") {64      const parts: Part[] = [];65      for (const p of m.content) if (p.type === "tool-result") parts.push({ functionResponse: { id: p.toolCallId, name: p.name, response: typeof p.result === "object" && p.result !== null ? (p.result as Record<string, unknown>) : { result: p.result } } });66      push("user", parts);67      continue;68    }69    const parts: Part[] = [];70    const toolResults: Part[] = [];71    for (const p of m.content) {72      switch (p.type) {73        case "text":74          if (p.text) parts.push({ text: p.text });75          break;76        case "image":77          parts.push({ inlineData: { mimeType: p.mimeType, data: p.data } });78          break;79        case "file":80          if (p.mimeType === "application/pdf") parts.push({ inlineData: { mimeType: "application/pdf", data: p.data } });81          else if (isTextLike(p.mimeType, p.name)) parts.push({ text: inlineTextFile(p.name, p.data) });82          else parts.push({ text: `[Attached file "${p.name}" (${p.mimeType}) — unsupported binary type]` });83          break;84        case "tool-call": {85          // Gemini 3.x rejects replayed function calls without a thought signature; the documented86          // escape hatch for calls that came from another provider is `skip_thought_signature_validator`.87          const sig = (p.providerData as { thoughtSignature?: string } | undefined)?.thoughtSignature ?? "skip_thought_signature_validator";88          parts.push({ functionCall: { id: p.id, name: p.name, args: p.arguments }, thoughtSignature: sig });89          break;90        }91        case "tool-result":92          toolResults.push({ functionResponse: { id: p.toolCallId, name: p.name, response: typeof p.result === "object" && p.result !== null ? (p.result as Record<string, unknown>) : { result: p.result } } });93          break;94        case "reasoning":95          // Thought summaries are not replayed; signatures travel on function-call parts.96          break;97      }98    }99    push(m.role === "assistant" ? "model" : "user", parts);100    if (toolResults.length) push("user", toolResults);101  }102  while (out.length && out[0].role !== "user") out.shift();103  return out;104}105106function buildConfig(req: UnifiedChatRequest): GenerateContentConfig {107  const { settings } = filterSettings(req.settings, req.modelInfo);108  const meta = (req.modelInfo?.metadata ?? {}) as Record<string, unknown>;109  const config: GenerateContentConfig = {};110  if (req.system?.trim()) config.systemInstruction = req.system;111  if (settings.temperature !== undefined) config.temperature = settings.temperature;112  if (settings.topP !== undefined) config.topP = settings.topP;113  if (settings.topK !== undefined) config.topK = settings.topK;114  if (settings.seed !== undefined) config.seed = settings.seed;115  if (settings.stop?.length) config.stopSequences = settings.stop.slice(0, 5);116  if (settings.maxTokens !== undefined) config.maxOutputTokens = settings.maxTokens;117118  // Thinking — per-family profile (see catalog.ts). Never send both level and budget.119  const profile = (meta.thinkingProfile as { mode: "level" | "budget" | "none"; levels?: string[]; none?: string | null } | undefined) ?? (req.modelInfo ? undefined : { mode: "none" as const });120  if (req.modelInfo?.capabilities.reasoning && profile && profile.mode !== "none") {121    const effort = settings.reasoningEffort;122    const thinking: ThinkingConfig = { includeThoughts: settings.includeReasoning !== false };123    let send = true;124    if (effort === "none") {125      thinking.includeThoughts = false;126      if (profile.none === "budget0") thinking.thinkingBudget = 0;127      else if (profile.none === "omit" || !profile.none) send = false; // off by default / cannot be disabled128      else (thinking as { thinkingLevel?: string }).thinkingLevel = profile.none;129    } else if (effort && profile.mode === "level") {130      const wanted = effort === "minimal" ? "minimal" : effort === "low" ? "low" : effort === "medium" ? "medium" : "high";131      const levels = profile.levels ?? ["low", "medium", "high"];132      const level = levels.includes(wanted) ? wanted : wanted === "minimal" ? levels[0] : levels[levels.length - 1];133      (thinking as { thinkingLevel?: string }).thinkingLevel = level.toUpperCase();134    } else if (effort && profile.mode === "budget") {135      thinking.thinkingBudget = effort === "minimal" || effort === "low" ? 1024 : effort === "medium" ? 8192 : 24_576;136    } else if (settings.thinkingBudget !== undefined && profile.mode === "budget") {137      thinking.thinkingBudget = settings.thinkingBudget;138    }139    if (send) config.thinkingConfig = thinking;140  }141142  if (settings.responseFormat?.type === "json_schema" && settings.responseFormat.schema) {143    config.responseMimeType = "application/json";144    config.responseJsonSchema = settings.responseFormat.schema;145  } else if (settings.responseFormat?.type === "json") {146    config.responseMimeType = "application/json";147  }148149  const tools: Tool[] = [];150  if (req.tools?.length) {151    tools.push({ functionDeclarations: req.tools.map<FunctionDeclaration>((t) => ({ name: t.name, description: t.description, parametersJsonSchema: t.parameters })) });152    const tc = settings.toolChoice;153    if (tc === "none") config.toolConfig = { functionCallingConfig: { mode: FunctionCallingConfigMode.NONE } };154    else if (tc === "required") config.toolConfig = { functionCallingConfig: { mode: FunctionCallingConfigMode.ANY } };155    else if (tc && typeof tc === "object") config.toolConfig = { functionCallingConfig: { mode: FunctionCallingConfigMode.ANY, allowedFunctionNames: [tc.name] } };156  }157  // Function declarations cannot be combined with built-in tools on most models; prefer the user's functions.158  if (!req.tools?.length) {159    if (settings.webSearch) tools.push({ googleSearch: {} });160    if (settings.codeExecution) tools.push({ codeExecution: {} });161  }162  if (tools.length) config.tools = tools;163  return config;164}165166function mapFinish(reason: string | undefined, sawTool: boolean): FinishReason {167  if (sawTool) return "tool-calls";168  switch (reason) {169    case "STOP":170      return "stop";171    case "MAX_TOKENS":172      return "length";173    case "SAFETY":174    case "RECITATION":175    case "PROHIBITED_CONTENT":176    case "SPII":177    case "BLOCKLIST":178    case "IMAGE_SAFETY":179      return "content-filter";180    case undefined:181      return "stop";182    default:183      return "other";184  }185}186187// ---------------------------------------------------------------------------188// Adapter189// ---------------------------------------------------------------------------190export const geminiAdapter: AIProviderAdapter = {191  id: "gemini",192  name: "Google Gemini",193  keyDocsUrl: "https://aistudio.google.com/apikey",194  keyPrefixHint: "AIza / AQ.",195196  async validateApiKey(apiKey, signal): Promise<ValidationResult> {197    const t0 = Date.now();198    try {199      const pager = await client(apiKey, 20_000).models.list({ config: { pageSize: 200, httpOptions: { timeout: 20_000 } } });200      let n = 0;201      for await (const m of pager) {202        if (signal?.aborted) break;203        if (normalizeGeminiModel(m)) n++;204      }205      return { ok: true, modelsAvailable: n, latencyMs: Date.now() - t0 };206    } catch (e) {207      return { ok: false, error: this.normalizeError(e).toJSON(), latencyMs: Date.now() - t0 };208    }209  },210211  async listModels(apiKey): Promise<PolyModel[]> {212    try {213      const out: PolyModel[] = [];214      const pager = await client(apiKey, 30_000).models.list({ config: { pageSize: 200 } });215      for await (const m of pager) {216        const n = normalizeGeminiModel(m);217        if (n) out.push(n);218      }219      return out;220    } catch (e) {221      throw this.normalizeError(e);222    }223  },224225  async chat(req): Promise<UnifiedChatResponse> {226    return collectStream("gemini", req.model, this.streamChat(req));227  },228229  async *streamChat(req: UnifiedChatRequest): AsyncIterable<UnifiedStreamEvent> {230    let stream: AsyncGenerator<GenerateContentResponse>;231    try {232      stream = await client(req.apiKey, req.timeoutMs).models.generateContentStream({ model: req.model, contents: toContents(req.messages), config: { ...buildConfig(req), abortSignal: req.signal } });233    } catch (e) {234      yield { type: "error", error: this.normalizeError(e).toJSON() };235      return;236    }237    let started = false;238    let sawTool = false;239    let finish: string | undefined;240    let lastUsage: GenerateContentResponse["usageMetadata"] | undefined;241    let toolIdx = 0;242    const seenCitations = new Set<string>();243    let blockReason: string | undefined;244    try {245      for await (const chunk of stream) {246        if (!started) {247          started = true;248          yield { type: "start", id: chunk.responseId, model: chunk.modelVersion };249        }250        if (chunk.promptFeedback?.blockReason) blockReason = chunk.promptFeedback.blockReason;251        const cand = chunk.candidates?.[0];252        for (const part of cand?.content?.parts ?? []) {253          if (part.text !== undefined && part.text !== "") {254            if (part.thought) yield { type: "reasoning-delta", text: part.text };255            else yield { type: "text-delta", text: part.text };256          }257          if (part.functionCall) {258            sawTool = true;259            const id = part.functionCall.id ?? `call_${Date.now()}_${toolIdx++}`;260            const name = part.functionCall.name ?? "";261            const argsText = JSON.stringify(part.functionCall.args ?? {});262            yield { type: "tool-start", id, name };263            yield { type: "tool-end", id, name, arguments: (part.functionCall.args ?? {}) as Record<string, unknown>, argumentsText: argsText, providerData: part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : undefined };264          }265          if (part.executableCode) yield { type: "server-tool", name: "code_execution", status: "started", data: { language: part.executableCode.language } };266          if (part.codeExecutionResult) yield { type: "server-tool", name: "code_execution", status: "completed", data: { outcome: part.codeExecutionResult.outcome, output: part.codeExecutionResult.output?.slice(0, 2000) } };267        }268        const gm = cand?.groundingMetadata;269        if (gm?.groundingChunks?.length) {270          if (gm.webSearchQueries?.length) yield { type: "server-tool", name: "web_search", status: "completed", data: { queries: gm.webSearchQueries } };271          for (const g of gm.groundingChunks) {272            const uri = g.web?.uri;273            if (uri && !seenCitations.has(uri)) {274              seenCitations.add(uri);275              yield { type: "citation", citation: { url: uri, title: g.web?.title ?? g.web?.domain ?? undefined, source: "google_search" } };276            }277          }278        }279        if (cand?.finishReason) finish = String(cand.finishReason);280        if (chunk.usageMetadata) lastUsage = chunk.usageMetadata;281      }282      if (blockReason && !finish) {283        yield { type: "error", error: { code: "CONTENT_REJECTED", message: `Prompt blocked (${blockReason})`, provider: "gemini", retryable: false, providerCode: blockReason } };284        return;285      }286      const u = lastUsage;287      const input = u?.promptTokenCount ?? 0;288      const thoughts = u?.thoughtsTokenCount ?? 0;289      const output = (u?.candidatesTokenCount ?? 0) + thoughts;290      yield { type: "usage", usage: { inputTokens: input, outputTokens: output, cachedInputTokens: u?.cachedContentTokenCount ?? undefined, reasoningTokens: thoughts || undefined, totalTokens: u?.totalTokenCount ?? input + output } };291      yield { type: "finish", reason: mapFinish(finish, sawTool) };292    } catch (e) {293      yield { type: "error", error: this.normalizeError(e).toJSON() };294    }295  },296297  async estimateTokens(req): Promise<TokenEstimate> {298    try {299      const res = await client(req.apiKey, 20_000).models.countTokens({ model: req.model, contents: toContents(req.messages) });300      return { inputTokens: res.totalTokens ?? 0, method: "provider" };301    } catch {302      const text = req.messages.map((m) => m.content.map((p) => (p.type === "text" ? p.text : "")).join(" ")).join(" ") + (req.system ?? "");303      return { inputTokens: heuristicTokens(text), method: "heuristic" };304    }305  },306307  normalizeError(error: unknown): PolyProviderError {308    if (error instanceof PolyProviderError) return error;309    if (error instanceof ApiError || (error && typeof error === "object" && "status" in error && typeof (error as { status: unknown }).status === "number")) {310      const e = error as { status: number; message: string };311      let message = e.message ?? "Gemini error";312      let googleStatus: string | undefined;313      // The SDK sometimes puts the raw JSON body in `message`.314      try {315        const parsed = JSON.parse(message) as { error?: { message?: string; status?: string } };316        if (parsed?.error) {317          message = parsed.error.message ?? message;318          googleStatus = parsed.error.status;319        }320      } catch {321        /* plain text */322      }323      let code = codeFromStatus(e.status);324      if (e.status === 401 || e.status === 403 || /api key not valid|api_key_invalid|unauthenticated|unregistered callers|permission_denied|API key/i.test(message)) code = "INVALID_API_KEY";325      else if (e.status === 404 && /no longer available|not found|is not supported/i.test(message)) code = "MODEL_NOT_FOUND";326      else if (e.status === 429) code = /limit: 0\b/.test(message) ? "INSUFFICIENT_CREDITS" : "RATE_LIMITED";327      else if (e.status === 503 || googleStatus === "UNAVAILABLE" || /high demand/i.test(message)) code = "PROVIDER_UNAVAILABLE";328      else if (e.status === 400) code = refineByMessage("INVALID_PARAMETER", message);329      else code = refineByMessage(code, message);330      const retryAfter = message.match(/retry in ([\d.]+)s/i);331      return new PolyProviderError({332        code,333        message: code === "INVALID_API_KEY" ? "Invalid API key" : message.replace(/\s+/g, " ").slice(0, 600),334        provider: "gemini",335        status: e.status,336        retryable: (isRetryableStatus(e.status) || code === "PROVIDER_UNAVAILABLE") && code !== "INSUFFICIENT_CREDITS",337        retryAfterMs: retryAfter ? Math.min(60_000, Math.round(Number(retryAfter[1]) * 1000)) : undefined,338        providerCode: googleStatus,339        cause: error,340      });341    }342    return normalizeGenericError("gemini", error);343  },344};345