354 lines
12 KiB
TypeScript
354 lines
12 KiB
TypeScript
import { AIError, type GenerateOptions, type Message } from "./index.ts";
|
|
import type { AIProvider, AIResult } from "./providers.ts";
|
|
|
|
export interface HttpAIProviderOptions {
|
|
apiKey?: string;
|
|
model: string;
|
|
baseUrl?: string;
|
|
fetch?: typeof fetch;
|
|
headers?: HeadersInit;
|
|
}
|
|
function key(configured: string | undefined, name: string): string {
|
|
const value = configured ?? process.env[name];
|
|
if (!value) throw new AIError(`Missing ${name}.`, 0, "authentication_error");
|
|
return value;
|
|
}
|
|
function promptMessages(prompt: string | Message[], options?: GenerateOptions) {
|
|
const messages =
|
|
options?.messages ??
|
|
(typeof prompt === "string" ? [{ role: "user" as const, content: prompt }] : prompt);
|
|
return options?.system ? [{ role: "system", content: options.system }, ...messages] : messages;
|
|
}
|
|
|
|
export function openAIProvider(options: HttpAIProviderOptions): AIProvider {
|
|
const send = options.fetch ?? fetch;
|
|
const base = (options.baseUrl ?? "https://api.openai.com/v1").replace(/\/$/, "");
|
|
return {
|
|
name: "openai",
|
|
capabilities: { streaming: false, structuredOutput: true, tools: true, usage: true },
|
|
async generate(prompt, call = {}) {
|
|
const response = await send(`${base}/chat/completions`, {
|
|
method: "POST",
|
|
headers: {
|
|
authorization: `Bearer ${key(options.apiKey, "OPENAI_API_KEY")}`,
|
|
"content-type": "application/json",
|
|
...Object.fromEntries(new Headers(options.headers)),
|
|
},
|
|
body: JSON.stringify({
|
|
model: call.model ?? options.model,
|
|
messages: promptMessages(prompt, call),
|
|
...(call.maxTokens ? { max_completion_tokens: call.maxTokens } : {}),
|
|
}),
|
|
signal: call.signal,
|
|
});
|
|
if (!response.ok)
|
|
throw new AIError(`OpenAI returned ${response.status}.`, response.status, "provider_error");
|
|
const body = (await response.json()) as any;
|
|
return {
|
|
value: String(body.choices?.[0]?.message?.content ?? ""),
|
|
provider: "openai",
|
|
model: body.model,
|
|
finishReason: body.choices?.[0]?.finish_reason,
|
|
usage: {
|
|
inputTokens: body.usage?.prompt_tokens,
|
|
outputTokens: body.usage?.completion_tokens,
|
|
totalTokens: body.usage?.total_tokens,
|
|
},
|
|
toolCalls: body.choices?.[0]?.message?.tool_calls?.map((tool: any) => ({
|
|
id: String(tool.id),
|
|
name: String(tool.function?.name),
|
|
arguments: JSON.parse(tool.function?.arguments ?? "{}"),
|
|
})),
|
|
raw: body,
|
|
};
|
|
},
|
|
};
|
|
}
|
|
|
|
export function googleAIProvider(options: HttpAIProviderOptions): AIProvider {
|
|
const send = options.fetch ?? fetch;
|
|
const base = (options.baseUrl ?? "https://generativelanguage.googleapis.com/v1beta").replace(
|
|
/\/$/,
|
|
"",
|
|
);
|
|
return {
|
|
name: "google",
|
|
capabilities: { streaming: false, structuredOutput: true, tools: true, usage: true },
|
|
async generate(prompt, call = {}) {
|
|
const response = await send(
|
|
`${base}/models/${encodeURIComponent(call.model ?? options.model)}:generateContent?key=${encodeURIComponent(key(options.apiKey, "GOOGLE_AI_API_KEY"))}`,
|
|
{
|
|
method: "POST",
|
|
headers: {
|
|
"content-type": "application/json",
|
|
...Object.fromEntries(new Headers(options.headers)),
|
|
},
|
|
body: JSON.stringify({
|
|
contents: promptMessages(prompt, call)
|
|
.filter((message) => message.role !== "system")
|
|
.map((message) => ({
|
|
role: message.role === "assistant" ? "model" : "user",
|
|
parts: [{ text: message.content }],
|
|
})),
|
|
...(call.system ? { systemInstruction: { parts: [{ text: call.system }] } } : {}),
|
|
}),
|
|
signal: call.signal,
|
|
},
|
|
);
|
|
if (!response.ok)
|
|
throw new AIError(
|
|
`Google AI returned ${response.status}.`,
|
|
response.status,
|
|
"provider_error",
|
|
);
|
|
const body = (await response.json()) as any;
|
|
return {
|
|
value: String(
|
|
body.candidates?.[0]?.content?.parts?.map((part: any) => part.text ?? "").join("") ?? "",
|
|
),
|
|
provider: "google",
|
|
model: call.model ?? options.model,
|
|
finishReason: body.candidates?.[0]?.finishReason,
|
|
usage: {
|
|
inputTokens: body.usageMetadata?.promptTokenCount,
|
|
outputTokens: body.usageMetadata?.candidatesTokenCount,
|
|
totalTokens: body.usageMetadata?.totalTokenCount,
|
|
},
|
|
raw: body,
|
|
};
|
|
},
|
|
};
|
|
}
|
|
|
|
/** OpenAI-compatible local servers including Ollama, llama.cpp and vLLM. */
|
|
export function localAIProvider(
|
|
options: Omit<HttpAIProviderOptions, "apiKey"> & { apiKey?: string },
|
|
): AIProvider {
|
|
const provider = openAIProvider({
|
|
...options,
|
|
apiKey: options.apiKey ?? "local",
|
|
baseUrl: options.baseUrl ?? "http://localhost:11434/v1",
|
|
});
|
|
return {
|
|
...provider,
|
|
name: "local",
|
|
generate: async (prompt, call) => ({
|
|
...(await provider.generate(prompt, call)),
|
|
provider: "local",
|
|
}),
|
|
};
|
|
}
|
|
|
|
export interface EmbeddingProvider {
|
|
name: string;
|
|
embed(
|
|
values: string[],
|
|
options?: { model?: string; signal?: AbortSignal },
|
|
): Promise<{ vectors: number[][]; usage?: { tokens?: number } }>;
|
|
}
|
|
export function openAIEmbeddings(options: HttpAIProviderOptions): EmbeddingProvider {
|
|
return {
|
|
name: "openai",
|
|
async embed(values, call = {}) {
|
|
const response = await (options.fetch ?? fetch)(
|
|
`${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/$/, "")}/embeddings`,
|
|
{
|
|
method: "POST",
|
|
headers: {
|
|
authorization: `Bearer ${key(options.apiKey, "OPENAI_API_KEY")}`,
|
|
"content-type": "application/json",
|
|
},
|
|
body: JSON.stringify({ model: call.model ?? options.model, input: values }),
|
|
signal: call.signal,
|
|
},
|
|
);
|
|
if (!response.ok)
|
|
throw new AIError(`Embedding provider returned ${response.status}.`, response.status);
|
|
const body = (await response.json()) as any;
|
|
return {
|
|
vectors: body.data.map((item: any) => item.embedding as number[]),
|
|
usage: { tokens: body.usage?.total_tokens },
|
|
};
|
|
},
|
|
};
|
|
}
|
|
|
|
export interface VectorRecord<T = Record<string, unknown>> {
|
|
id: string;
|
|
vector: number[];
|
|
text: string;
|
|
metadata: T;
|
|
}
|
|
export interface VectorMatch<T = Record<string, unknown>> extends VectorRecord<T> {
|
|
score: number;
|
|
}
|
|
export interface VectorStore<T = Record<string, unknown>> {
|
|
upsert(records: VectorRecord<T>[]): Promise<void>;
|
|
query(
|
|
vector: number[],
|
|
limit?: number,
|
|
filter?: (metadata: T) => boolean,
|
|
): Promise<VectorMatch<T>[]>;
|
|
delete(ids: string[]): Promise<void>;
|
|
}
|
|
export function memoryVectorStore<T = Record<string, unknown>>(): VectorStore<T> {
|
|
const records = new Map<string, VectorRecord<T>>();
|
|
const cosine = (a: number[], b: number[]) => {
|
|
if (a.length !== b.length || !a.length) return 0;
|
|
let dot = 0,
|
|
aa = 0,
|
|
bb = 0;
|
|
for (let i = 0; i < a.length; i++) {
|
|
dot += a[i]! * b[i]!;
|
|
aa += a[i]! ** 2;
|
|
bb += b[i]! ** 2;
|
|
}
|
|
return aa && bb ? dot / Math.sqrt(aa * bb) : 0;
|
|
};
|
|
return {
|
|
async upsert(values) {
|
|
for (const value of values) records.set(value.id, { ...value, vector: [...value.vector] });
|
|
},
|
|
async query(vector, limit = 5, filter) {
|
|
return [...records.values()]
|
|
.filter((record) => !filter || filter(record.metadata))
|
|
.map((record) => ({ ...record, score: cosine(vector, record.vector) }))
|
|
.sort((a, b) => b.score - a.score)
|
|
.slice(0, Math.max(0, limit));
|
|
},
|
|
async delete(ids) {
|
|
for (const id of ids) records.delete(id);
|
|
},
|
|
};
|
|
}
|
|
|
|
export function createRagPipeline<T>(options: {
|
|
embeddings: EmbeddingProvider;
|
|
store: VectorStore<T>;
|
|
generate: (prompt: string, options?: GenerateOptions) => Promise<AIResult<string>>;
|
|
maxContextCharacters?: number;
|
|
}) {
|
|
return {
|
|
async index(documents: Array<{ id: string; text: string; metadata: T }>) {
|
|
const embedded = await options.embeddings.embed(documents.map((document) => document.text));
|
|
await options.store.upsert(
|
|
documents.map((document, index) => ({ ...document, vector: embedded.vectors[index]! })),
|
|
);
|
|
return documents.length;
|
|
},
|
|
async ask(
|
|
question: string,
|
|
call: GenerateOptions & { limit?: number; filter?: (metadata: T) => boolean } = {},
|
|
) {
|
|
const embedded = await options.embeddings.embed([question], { signal: call.signal });
|
|
const matches = await options.store.query(embedded.vectors[0]!, call.limit, call.filter);
|
|
const context = matches
|
|
.map((match, index) => `[${index + 1}] ${match.text}`)
|
|
.join("\n\n")
|
|
.slice(0, options.maxContextCharacters ?? 12_000);
|
|
const result = await options.generate(
|
|
`Answer using only the supplied context. Cite sources as [n].\n\nContext:\n${context}\n\nQuestion: ${question}`,
|
|
call,
|
|
);
|
|
return { ...result, sources: matches };
|
|
},
|
|
};
|
|
}
|
|
|
|
export interface ConversationStore {
|
|
load(id: string): Promise<Message[]>;
|
|
append(id: string, messages: Message[]): Promise<void>;
|
|
clear(id: string): Promise<void>;
|
|
}
|
|
export function memoryConversationStore(maxMessages = 100): ConversationStore {
|
|
const conversations = new Map<string, Message[]>();
|
|
return {
|
|
async load(id) {
|
|
return [...(conversations.get(id) ?? [])];
|
|
},
|
|
async append(id, messages) {
|
|
conversations.set(id, [...(conversations.get(id) ?? []), ...messages].slice(-maxMessages));
|
|
},
|
|
async clear(id) {
|
|
conversations.delete(id);
|
|
},
|
|
};
|
|
}
|
|
export function promptTemplate(template: string) {
|
|
return (variables: Record<string, string | number>) =>
|
|
template.replace(/\{\{\s*([\w.-]+)\s*\}\}/g, (_match, name: string) => {
|
|
if (!(name in variables)) throw new Error(`WRN-AI-PROMPT-VARIABLE:${name}`);
|
|
return String(variables[name]);
|
|
});
|
|
}
|
|
export type AIGuardrail = (input: {
|
|
prompt: string | Message[];
|
|
output?: string;
|
|
}) => void | Promise<void>;
|
|
export const maxPromptLength =
|
|
(maximum: number): AIGuardrail =>
|
|
({ prompt }) => {
|
|
const length =
|
|
typeof prompt === "string"
|
|
? prompt.length
|
|
: prompt.reduce((sum, message) => sum + message.content.length, 0);
|
|
if (length > maximum)
|
|
throw new AIError("Prompt exceeds configured guardrail.", 400, "guardrail");
|
|
};
|
|
export function guardedProvider(provider: AIProvider, guardrails: AIGuardrail[]): AIProvider {
|
|
return {
|
|
...provider,
|
|
async generate(prompt, options) {
|
|
for (const guardrail of guardrails) await guardrail({ prompt });
|
|
const result = await provider.generate(prompt, options);
|
|
for (const guardrail of guardrails) await guardrail({ prompt, output: result.value });
|
|
return result;
|
|
},
|
|
};
|
|
}
|
|
export function aiRateLimiter(options: { limit: number; windowMs: number; now?: () => number }) {
|
|
const buckets = new Map<string, { count: number; reset: number }>();
|
|
const now = options.now ?? Date.now;
|
|
return (key: string) => {
|
|
const time = now();
|
|
const bucket = buckets.get(key);
|
|
if (!bucket || bucket.reset <= time) {
|
|
buckets.set(key, { count: 1, reset: time + options.windowMs });
|
|
return { allowed: true, remaining: options.limit - 1 };
|
|
}
|
|
if (bucket.count >= options.limit)
|
|
return { allowed: false, remaining: 0, retryAfterMs: bucket.reset - time };
|
|
bucket.count++;
|
|
return { allowed: true, remaining: options.limit - bucket.count };
|
|
};
|
|
}
|
|
export async function evaluateAI(
|
|
cases: Array<{
|
|
name: string;
|
|
prompt: string;
|
|
expected?: string;
|
|
score?: (output: string) => number | Promise<number>;
|
|
}>,
|
|
generate: (prompt: string) => Promise<string>,
|
|
) {
|
|
const results = [];
|
|
for (const item of cases) {
|
|
const started = performance.now();
|
|
const output = await generate(item.prompt);
|
|
const score = item.score
|
|
? await item.score(output)
|
|
: item.expected === undefined
|
|
? 1
|
|
: output.includes(item.expected)
|
|
? 1
|
|
: 0;
|
|
results.push({ name: item.name, output, score, durationMs: performance.now() - started });
|
|
}
|
|
return {
|
|
results,
|
|
score: results.length
|
|
? results.reduce((sum, result) => sum + result.score, 0) / results.length
|
|
: 0,
|
|
};
|
|
}
|