release: WRNexusJS 0.8.0
This commit is contained in:
@@ -0,0 +1,353 @@
|
||||
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,
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user