97 lines
3.1 KiB
TypeScript
97 lines
3.1 KiB
TypeScript
import { describe, expect, test } from "bun:test";
|
|
import {
|
|
aiRateLimiter,
|
|
createRagPipeline,
|
|
evaluateAI,
|
|
googleAIProvider,
|
|
guardedProvider,
|
|
localAIProvider,
|
|
maxPromptLength,
|
|
memoryConversationStore,
|
|
memoryVectorStore,
|
|
openAIProvider,
|
|
promptTemplate,
|
|
} from "../src/platform.ts";
|
|
|
|
describe("AI platform", () => {
|
|
test("adapts OpenAI, Google and local providers", async () => {
|
|
const openFetch = (async () =>
|
|
Response.json({
|
|
model: "gpt",
|
|
choices: [{ message: { content: "hello" }, finish_reason: "stop" }],
|
|
usage: { total_tokens: 3 },
|
|
})) as unknown as typeof fetch;
|
|
expect(
|
|
(await openAIProvider({ model: "gpt", apiKey: "key", fetch: openFetch }).generate("hi"))
|
|
.value,
|
|
).toBe("hello");
|
|
const googleFetch = (async () =>
|
|
Response.json({
|
|
candidates: [{ content: { parts: [{ text: "hola" }] } }],
|
|
usageMetadata: { totalTokenCount: 2 },
|
|
})) as unknown as typeof fetch;
|
|
expect(
|
|
(
|
|
await googleAIProvider({ model: "gemini", apiKey: "key", fetch: googleFetch }).generate(
|
|
"hi",
|
|
)
|
|
).value,
|
|
).toBe("hola");
|
|
expect(
|
|
(await localAIProvider({ model: "llama", fetch: openFetch }).generate("hi")).provider,
|
|
).toBe("local");
|
|
});
|
|
|
|
test("indexes and retrieves a RAG answer", async () => {
|
|
const store = memoryVectorStore<{ source: string }>();
|
|
const embeddings = {
|
|
name: "test",
|
|
embed: async (values: string[]) => ({
|
|
vectors: values.map((value) => (value.includes("Bun") ? [1, 0] : [0, 1])),
|
|
}),
|
|
};
|
|
const rag = createRagPipeline({
|
|
embeddings,
|
|
store,
|
|
generate: async (prompt) => ({
|
|
value: prompt.includes("fast") ? "Bun [1]" : "none",
|
|
provider: "test",
|
|
}),
|
|
});
|
|
await rag.index([
|
|
{ id: "bun", text: "Bun is fast", metadata: { source: "docs" } },
|
|
{ id: "other", text: "Other", metadata: { source: "other" } },
|
|
]);
|
|
const result = await rag.ask("Bun runtime");
|
|
expect(result.value).toBe("Bun [1]");
|
|
expect(result.sources[0]?.id).toBe("bun");
|
|
});
|
|
|
|
test("persists conversations, renders prompts, guards and rate limits", async () => {
|
|
const conversations = memoryConversationStore(2);
|
|
await conversations.append("one", [
|
|
{ role: "user", content: "a" },
|
|
{ role: "assistant", content: "b" },
|
|
{ role: "user", content: "c" },
|
|
]);
|
|
expect(await conversations.load("one")).toHaveLength(2);
|
|
expect(promptTemplate("Hello {{name}}")({ name: "WrNexus" })).toBe("Hello WrNexus");
|
|
const provider = guardedProvider(
|
|
{ name: "test", generate: async () => ({ value: "ok", provider: "test" }) },
|
|
[maxPromptLength(3)],
|
|
);
|
|
await expect(provider.generate("long")).rejects.toThrow("guardrail");
|
|
const limit = aiRateLimiter({ limit: 1, windowMs: 100, now: () => 0 });
|
|
expect(limit("u").allowed).toBe(true);
|
|
expect(limit("u").allowed).toBe(false);
|
|
});
|
|
|
|
test("evaluates model output", async () => {
|
|
const report = await evaluateAI(
|
|
[{ name: "answer", prompt: "question", expected: "42" }],
|
|
async () => "42",
|
|
);
|
|
expect(report.score).toBe(1);
|
|
});
|
|
});
|