diff --git a/src/shared/__tests__/embeddingModels.spec.ts b/src/shared/__tests__/embeddingModels.spec.ts new file mode 100644 index 0000000000..7a0b718ede --- /dev/null +++ b/src/shared/__tests__/embeddingModels.spec.ts @@ -0,0 +1,46 @@ +import { describe, it, expect } from "vitest" +import { getModelDimension, EMBEDDING_MODEL_PROFILES } from "../embeddingModels" + +describe("embeddingModels", () => { + describe("EMBEDDING_MODEL_PROFILES", () => { + it("should have codestral-embed model defined for openai-compatible provider", () => { + const openAiCompatibleModels = EMBEDDING_MODEL_PROFILES["openai-compatible"] + expect(openAiCompatibleModels).toBeDefined() + expect(openAiCompatibleModels?.["codestral-embed"]).toBeDefined() + expect(openAiCompatibleModels?.["codestral-embed"].dimension).toBe(3072) + expect(openAiCompatibleModels?.["codestral-embed"].scoreThreshold).toBe(0.4) + }) + + it("should have codestral-embed model defined for mistral provider", () => { + const mistralModels = EMBEDDING_MODEL_PROFILES["mistral"] + expect(mistralModels).toBeDefined() + expect(mistralModels?.["codestral-embed-2505"]).toBeDefined() + expect(mistralModels?.["codestral-embed-2505"].dimension).toBe(1536) + expect(mistralModels?.["codestral-embed-2505"].scoreThreshold).toBe(0.4) + }) + }) + + describe("getModelDimension", () => { + it("should return correct dimension for codestral-embed in openai-compatible provider", () => { + const dimension = getModelDimension("openai-compatible", "codestral-embed") + expect(dimension).toBe(3072) + }) + + it("should return correct dimension for codestral-embed-2505 in mistral provider", () => { + const dimension = getModelDimension("mistral", "codestral-embed-2505") + expect(dimension).toBe(1536) + }) + + it("should return undefined for unknown model", () => { + const dimension = getModelDimension("openai-compatible", "unknown-model") + expect(dimension).toBeUndefined() + }) + + it("should return correct dimensions for other openai-compatible models", () => { + expect(getModelDimension("openai-compatible", "text-embedding-3-small")).toBe(1536) + expect(getModelDimension("openai-compatible", "text-embedding-3-large")).toBe(3072) + expect(getModelDimension("openai-compatible", "text-embedding-ada-002")).toBe(1536) + expect(getModelDimension("openai-compatible", "nomic-embed-code")).toBe(3584) + }) + }) +}) diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index a3cd61e659..e969919618 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -45,6 +45,7 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { scoreThreshold: 0.15, queryPrefix: "Represent this query for searching relevant code: ", }, + "codestral-embed": { dimension: 3072, scoreThreshold: 0.4 }, }, gemini: { "text-embedding-004": { dimension: 768 },