SPB Git forge

spb/rareindex

Public
54commits 1branches 0releases
7.1 MBsize
maindefault branch
10 days agolast push
TypeScript 61.9% HTML 37.2% SQL 0.7%
10.0 KB · 207 lines typescript
Raw Blame History
1import OpenAI from 'openai';2import { z } from 'zod';3import { estimateUsd, addUsage, ZERO_USAGE } from '../pricing.js';4import { recordCost } from '../costs.js';5import { AiNotConfiguredError, type ChatMessage, type Completion, type CompletionRequest, type ContentPart, type EmbedRequest, type EmbedResult, type ExtractRequest, type ExtractResult, type ModelProvider, type Role, type StreamDelta, type ToolAgentRequest, type Usage } from '../types.js';67export interface OpenAICompatibleOptions {8  id?: string;9  apiKey: string;10  baseURL?: string;11  models?: Partial<Record<Role, string>>;12  embeddingModel?: string;13  embeddingDimensions?: number;14}1516const DEFAULT_MODELS: Record<Role, string> = {17  normalize: 'gpt-5-mini',18  classify: 'gpt-5-mini',19  summarize: 'gpt-5-mini',20  resolve: 'gpt-5-mini',21  research: 'gpt-5',22  vision: 'gpt-5',23  embed: 'text-embedding-3-small',24};2526function toContent(content: string | ContentPart[]): string | OpenAI.Chat.Completions.ChatCompletionContentPart[] {27  if (typeof content === 'string') return content;28  return content.map((p): OpenAI.Chat.Completions.ChatCompletionContentPart => {29    if (p.type === 'text') return { type: 'text', text: p.text };30    const url = p.image.url ?? `data:${p.image.mediaType};base64,${p.image.data ?? ''}`;31    return { type: 'image_url', image_url: { url } };32  });33}3435function toMessages(system: string | undefined, messages: ChatMessage[]): OpenAI.Chat.Completions.ChatCompletionMessageParam[] {36  const out: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = [];37  if (system) out.push({ role: 'system', content: system });38  for (const m of messages) {39    if (m.role === 'user') out.push({ role: 'user', content: toContent(m.content) });40    else out.push({ role: 'assistant', content: typeof m.content === 'string' ? m.content : m.content.filter((p) => p.type === 'text').map((p) => (p as { text: string }).text).join('') });41  }42  return out;43}4445function usageOf(u: OpenAI.Completions.CompletionUsage | null | undefined): Usage {46  return {47    inputTokens: (u?.prompt_tokens ?? 0) - (u?.prompt_tokens_details?.cached_tokens ?? 0),48    outputTokens: u?.completion_tokens ?? 0,49    cacheReadTokens: u?.prompt_tokens_details?.cached_tokens ?? 0,50    cacheWriteTokens: 0,51  };52}5354/**55 * Provider for any OpenAI-compatible endpoint: OpenAI itself, or the local MacLustr gateway56 * (LOCAL_LLM_BASE_URL). Also the only embeddings provider for now.57 */58export class OpenAICompatibleProvider implements ModelProvider {59  readonly id: string;60  private readonly client: OpenAI;61  private readonly models: Record<Role, string>;62  private readonly embeddingModel: string;63  private readonly embeddingDimensions: number;64  private readonly isLocal: boolean;6566  constructor(opts: OpenAICompatibleOptions) {67    if (!opts.apiKey && !opts.baseURL) throw new AiNotConfiguredError('OPENAI_API_KEY or LOCAL_LLM_BASE_URL missing');68    this.id = opts.id ?? (opts.baseURL ? 'openai-compatible' : 'openai');69    this.isLocal = Boolean(opts.baseURL);70    this.client = new OpenAI({ apiKey: opts.apiKey || 'local', baseURL: opts.baseURL, timeout: 5 * 60_000, maxRetries: 2 });71    this.models = { ...DEFAULT_MODELS, ...(opts.models ?? {}) };72    this.embeddingModel = opts.embeddingModel ?? this.models.embed;73    this.embeddingDimensions = opts.embeddingDimensions ?? 1536;74  }7576  supports(role: Role): boolean {77    if (role === 'embed') return !this.isLocal || Boolean(this.embeddingModel);78    return true;79  }8081  modelFor(role: Role): string {82    return this.models[role];83  }8485  private async account(role: Role, model: string, usage: Usage, cost?: CompletionRequest['cost'], units?: number): Promise<number> {86    const usdEst = this.isLocal ? 0 : estimateUsd(model, usage);87    await recordCost({ provider: this.id, model, role, usage, usdEst, context: cost, units });88    return usdEst;89  }9091  async complete(role: Role, req: CompletionRequest): Promise<Completion> {92    const model = this.modelFor(role);93    const res = await this.client.chat.completions.create({94      model,95      messages: toMessages(req.system, req.messages),96      max_completion_tokens: req.maxTokens ?? 4096,97      ...(req.temperature !== undefined ? { temperature: req.temperature } : {}),98      ...(req.json ? { response_format: { type: 'json_object' } } : {}),99      ...(req.stopSequences ? { stop: req.stopSequences } : {}),100    });101    const usage = usageOf(res.usage);102    const usdEst = await this.account(role, model, usage, req.cost);103    const choice = res.choices[0];104    return { text: choice?.message?.content ?? '', model, provider: this.id, usage, usdEst, stopReason: choice?.finish_reason ?? null, refusal: choice?.message?.refusal ? { category: null, explanation: choice.message.refusal } : null };105  }106107  async *stream(role: Role, req: CompletionRequest): AsyncIterable<StreamDelta> {108    yield* this.runTools(role, { ...req, tools: [], execute: async () => null });109  }110111  async *runTools(role: Role, req: ToolAgentRequest): AsyncIterable<StreamDelta> {112    const model = this.modelFor(role);113    const messages = toMessages(req.system, req.messages);114    const tools: OpenAI.Chat.Completions.ChatCompletionTool[] = req.tools.map((t) => ({ type: 'function', function: { name: t.name, description: t.description, parameters: t.inputSchema } }));115    let total: Usage = ZERO_USAGE;116    let stopReason: string | null = null;117    const maxIter = req.maxIterations ?? 8;118    try {119      for (let i = 0; i < maxIter; i++) {120        const stream = await this.client.chat.completions.create(121          { model, messages, max_completion_tokens: req.maxTokens ?? 8192, ...(tools.length ? { tools } : {}), stream: true, stream_options: { include_usage: true } },122          { signal: req.signal },123        );124        const calls = new Map<number, { id: string; name: string; args: string }>();125        let finish: string | null = null;126        for await (const chunk of stream) {127          if (chunk.usage) {128            const u = usageOf(chunk.usage);129            total = addUsage(total, u);130            yield { type: 'usage', usage: u, usdEst: this.isLocal ? 0 : estimateUsd(model, u), model };131          }132          const c = chunk.choices[0];133          if (!c) continue;134          if (c.delta.content) yield { type: 'text', text: c.delta.content };135          for (const tc of c.delta.tool_calls ?? []) {136            const cur = calls.get(tc.index) ?? { id: tc.id ?? '', name: '', args: '' };137            if (tc.id) cur.id = tc.id;138            if (tc.function?.name) cur.name += tc.function.name;139            if (tc.function?.arguments) cur.args += tc.function.arguments;140            calls.set(tc.index, cur);141          }142          if (c.finish_reason) finish = c.finish_reason;143        }144        stopReason = finish;145        if (finish !== 'tool_calls' || calls.size === 0) break;146        const toolCalls = [...calls.values()];147        messages.push({ role: 'assistant', content: null, tool_calls: toolCalls.map((tc) => ({ id: tc.id, type: 'function', function: { name: tc.name, arguments: tc.args } })) });148        for (const tc of toolCalls) {149          let input: unknown = {};150          try {151            input = tc.args ? JSON.parse(tc.args) : {};152          } catch {153            input = {};154          }155          yield { type: 'tool_call', id: tc.id, name: tc.name, input };156          const started = Date.now();157          let output: unknown;158          let isError = false;159          try {160            output = await req.execute(tc.name, input);161          } catch (err) {162            isError = true;163            output = { error: err instanceof Error ? err.message : String(err) };164          }165          yield { type: 'tool_result', id: tc.id, name: tc.name, output, isError, durationMs: Date.now() - started };166          messages.push({ role: 'tool', tool_call_id: tc.id, content: typeof output === 'string' ? output : JSON.stringify(output ?? null) });167        }168      }169    } finally {170      await this.account(role, model, total, req.cost);171    }172    yield { type: 'done', stopReason };173  }174175  async extract<T>(role: Role, req: ExtractRequest<T>): Promise<ExtractResult<T>> {176    const model = this.modelFor(role);177    const schema = z.toJSONSchema(req.schema as z.ZodType, { target: 'draft-2020-12', io: 'output' }) as Record<string, unknown>;178    const content = typeof req.input === 'string' ? `${req.prompt}\n\n${req.input}` : [{ type: 'text', text: req.prompt } as ContentPart, ...req.input];179    const res = await this.client.chat.completions.create({180      model,181      messages: toMessages(req.system, [{ role: 'user', content }]),182      max_completion_tokens: req.maxTokens ?? 4096,183      response_format: { type: 'json_schema', json_schema: { name: 'extraction', schema, strict: false } },184    });185    const usage = usageOf(res.usage);186    const usdEst = await this.account(role, model, usage, req.cost);187    const text = res.choices[0]?.message?.content ?? '';188    const parsed = req.schema.safeParse(JSON.parse(text));189    if (!parsed.success) throw new Error(`extract: model output did not match schema: ${parsed.error.message}`);190    const data = parsed.data as T & { confidence?: number };191    return { data: parsed.data, confidence: typeof data.confidence === 'number' ? Math.max(0, Math.min(1, data.confidence)) : 1, model, provider: this.id, usage, usdEst };192  }193194  vision<T>(req: ExtractRequest<T>): Promise<ExtractResult<T>> {195    return this.extract('vision', req);196  }197198  async embed(req: EmbedRequest): Promise<EmbedResult> {199    if (!this.embeddingModel) throw new AiNotConfiguredError('no embedding model configured');200    const res = await this.client.embeddings.create({ model: this.embeddingModel, input: req.texts, ...(this.embeddingModel.startsWith('text-embedding-3') ? { dimensions: this.embeddingDimensions } : {}) });201    const usage: Usage = { ...ZERO_USAGE, inputTokens: res.usage?.prompt_tokens ?? 0 };202    const usdEst = await this.account('embed', this.embeddingModel, usage, req.cost, req.texts.length);203    const vectors = res.data.sort((a, b) => a.index - b.index).map((d) => d.embedding);204    return { vectors, model: this.embeddingModel, provider: this.id, dimensions: vectors[0]?.length ?? this.embeddingDimensions, usdEst };205  }206}207