@@ -1,8 +1,122 @@
1-import { describe, expect, it } from "vitest";
2-import { DEFAULT_LOCAL_MODEL } from "./embeddings.js";
1+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+import { createLocalEmbeddingProvider, DEFAULT_LOCAL_MODEL } from "./embeddings.js";
3+4+const nodeLlamaMock = vi.hoisted(() => ({
5+importNodeLlamaCpp: vi.fn(),
6+}));
7+8+vi.mock("./node-llama.js", () => ({
9+importNodeLlamaCpp: nodeLlamaMock.importNodeLlamaCpp,
10+}));
11+12+beforeEach(() => {
13+nodeLlamaMock.importNodeLlamaCpp.mockReset();
14+});
15+16+afterEach(() => {
17+vi.resetAllMocks();
18+});
19+20+function mockLocalEmbeddingRuntime(vector = new Float32Array([2.35, 3.45, 0.63, 4.3])) {
21+const getEmbeddingFor = vi.fn().mockResolvedValue({ vector });
22+const createEmbeddingContext = vi.fn().mockResolvedValue({ getEmbeddingFor });
23+const loadModel = vi.fn().mockResolvedValue({ createEmbeddingContext });
24+const resolveModelFile = vi.fn(async (modelPath: string) => `/resolved/${modelPath}`);
25+26+nodeLlamaMock.importNodeLlamaCpp.mockResolvedValue({
27+getLlama: async () => ({ loadModel }),
28+ resolveModelFile,
29+LlamaLogLevel: { error: 0 },
30+} as never);
31+32+return { createEmbeddingContext, getEmbeddingFor, loadModel, resolveModelFile };
33+}
34+35+describe("local embedding provider", () => {
36+it("normalizes local embeddings and resolves the default local model", async () => {
37+const runtime = mockLocalEmbeddingRuntime();
38+39+const provider = await createLocalEmbeddingProvider({
40+config: {} as never,
41+provider: "local",
42+model: "",
43+fallback: "none",
44+});
45+46+const embedding = await provider.embedQuery("test query");
47+const magnitude = Math.sqrt(embedding.reduce((sum, value) => sum + value * value, 0));
3484-describe("package embeddings barrel", () => {
5-it("re-exports the source local embedding contract", () => {
649expect(DEFAULT_LOCAL_MODEL).toContain("embeddinggemma");
50+expect(magnitude).toBeCloseTo(1, 5);
51+expect(runtime.resolveModelFile).toHaveBeenCalledWith(DEFAULT_LOCAL_MODEL, undefined);
52+expect(runtime.getEmbeddingFor).toHaveBeenCalledWith("test query");
53+});
54+55+it("passes default contextSize (4096) to createEmbeddingContext when not configured", async () => {
56+const runtime = mockLocalEmbeddingRuntime();
57+58+const provider = await createLocalEmbeddingProvider({
59+config: {} as never,
60+provider: "local",
61+model: "",
62+fallback: "none",
63+});
64+65+await provider.embedQuery("context size default test");
66+67+expect(runtime.createEmbeddingContext).toHaveBeenCalledWith({ contextSize: 4096 });
68+});
69+70+it("passes configured contextSize to createEmbeddingContext", async () => {
71+const runtime = mockLocalEmbeddingRuntime();
72+73+const provider = await createLocalEmbeddingProvider({
74+config: {} as never,
75+provider: "local",
76+model: "",
77+fallback: "none",
78+local: { contextSize: 2048 },
79+});
80+81+await provider.embedQuery("context size custom test");
82+83+expect(runtime.createEmbeddingContext).toHaveBeenCalledWith({ contextSize: 2048 });
84+});
85+86+it('passes "auto" contextSize to createEmbeddingContext when explicitly set', async () => {
87+const runtime = mockLocalEmbeddingRuntime();
88+89+const provider = await createLocalEmbeddingProvider({
90+config: {} as never,
91+provider: "local",
92+model: "",
93+fallback: "none",
94+local: { contextSize: "auto" },
95+});
96+97+await provider.embedQuery("context size auto test");
98+99+expect(runtime.createEmbeddingContext).toHaveBeenCalledWith({ contextSize: "auto" });
100+});
101+102+it("trims explicit local model paths and cache directories", async () => {
103+const runtime = mockLocalEmbeddingRuntime(new Float32Array([1, 0]));
104+105+const provider = await createLocalEmbeddingProvider({
106+config: {} as never,
107+provider: "local",
108+model: "",
109+fallback: "none",
110+local: {
111+modelPath: " /models/embed.gguf ",
112+modelCacheDir: " /cache/models ",
113+},
114+});
115+116+await provider.embedBatch(["a", "b"]);
117+118+expect(provider.model).toBe("/models/embed.gguf");
119+expect(runtime.resolveModelFile).toHaveBeenCalledWith("/models/embed.gguf", "/cache/models");
120+expect(runtime.getEmbeddingFor).toHaveBeenCalledTimes(2);
7121});
8122});
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。