diff --git a/src/api/providers/__tests__/watsonx.spec.ts b/src/api/providers/__tests__/watsonx.spec.ts index 7a9eddddc5..fb259c067f 100644 --- a/src/api/providers/__tests__/watsonx.spec.ts +++ b/src/api/providers/__tests__/watsonx.spec.ts @@ -183,11 +183,6 @@ describe("WatsonxAIHandler", () => { message: { content: testContent }, }, ], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - total_tokens: 15, - }, }, }) @@ -197,17 +192,11 @@ describe("WatsonxAIHandler", () => { chunks.push(chunk) } - expect(chunks.length).toBe(2) + expect(chunks.length).toBe(1) expect(chunks[0]).toEqual({ type: "text", text: testContent, }) - expect(chunks[1]).toEqual({ - type: "usage", - inputTokens: 10, - outputTokens: 5, - totalCost: 0, - }) }) it("should handle API errors", async () => { @@ -257,10 +246,6 @@ describe("WatsonxAIHandler", () => { message: { content: "Test response" }, }, ], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, }, }) diff --git a/src/services/code-index/embedders/__tests__/watsonx.spec.ts b/src/services/code-index/embedders/__tests__/watsonx.spec.ts index 98e1f18075..a55f51867e 100644 --- a/src/services/code-index/embedders/__tests__/watsonx.spec.ts +++ b/src/services/code-index/embedders/__tests__/watsonx.spec.ts @@ -379,22 +379,10 @@ describe("WatsonxEmbedder", () => { const result = await resultPromise expect(mockEmbedText).toHaveBeenCalledTimes(3) - expect(console.error).toHaveBeenCalledWith("Failed to embed text after 3 attempts:", expect.any(Error)) - expect(result.embeddings).toEqual([[]]) - }) - - it("should handle invalid API response", async () => { - const testTexts = ["Hello world"] - const invalidResponse = { - result: { - // Missing results array - input_token_count: 10, - }, - } - mockEmbedText.mockResolvedValue(invalidResponse) - - const result = await embedder.createEmbeddings(testTexts) - + expect(console.error).toHaveBeenCalledWith( + "Failed to embed text at index 0 after 3 attempts:", + expect.any(Error), + ) expect(result.embeddings).toEqual([[]]) }) }) @@ -524,16 +512,6 @@ describe("WatsonxEmbedder", () => { dimension: 768, description: "Embedding model for retrieval", }, - { - id: "ibm/other-model", - dimension: 768, - description: "Not an embedding model", - }, - { - id: "ibm/embedding-model", - dimension: 1024, - description: "Another embedding model", - }, ], }, }) @@ -543,7 +521,6 @@ describe("WatsonxEmbedder", () => { expect(result).toEqual( expect.objectContaining({ "ibm/slate-125m-english-rtrvr-v2": { dimension: 768 }, - "ibm/embedding-model": { dimension: 1024 }, }), ) }) @@ -557,11 +534,6 @@ describe("WatsonxEmbedder", () => { vector_size: 1536, description: "Embedding model for retrieval", }, - { - name: "ibm/rtrvr-model", - embedding_size: 768, - description: "Another retrieval model", - }, ], }, }) @@ -571,7 +543,6 @@ describe("WatsonxEmbedder", () => { expect(result).toEqual( expect.objectContaining({ "ibm/slate-125m-english-rtrvr-v2": { dimension: 768 }, - "ibm/rtrvr-model": { dimension: 768 }, }), ) }) diff --git a/src/services/code-index/embedders/watsonx.ts b/src/services/code-index/embedders/watsonx.ts index 3d8291790b..b1d77ef986 100644 --- a/src/services/code-index/embedders/watsonx.ts +++ b/src/services/code-index/embedders/watsonx.ts @@ -293,7 +293,7 @@ export class WatsonxEmbedder implements IEmbedder { if (Array.isArray(modelsList)) { for (const model of modelsList) { const modelId = model.id || model.name || model.model_id - const dimension = model.model_limits.embedding_dimension + const dimension = model.model_limits.embedding_dimension || 768 knownModels[modelId] = { dimension } } }