From 448dc2be37b8e08c1bedc4a1caba4706cc560ab8 Mon Sep 17 00:00:00 2001 From: Daniel <57051444+daniel-lxs@users.noreply.github.com> Date: Sat, 24 May 2025 12:12:13 -0500 Subject: [PATCH] Fix Qdrant initialization logic when vector size changes (#3925) --- .../__tests__/config-manager.test.ts | 24 - src/services/code-index/config-manager.ts | 13 - src/services/code-index/manager.ts | 17 +- .../__tests__/qdrant-client.test.ts | 862 ++++++++++++++++++ .../code-index/vector-store/qdrant-client.ts | 63 +- 5 files changed, 914 insertions(+), 65 deletions(-) create mode 100644 src/services/code-index/vector-store/__tests__/qdrant-client.test.ts diff --git a/src/services/code-index/__tests__/config-manager.test.ts b/src/services/code-index/__tests__/config-manager.test.ts index 87fc958003..7b27139943 100644 --- a/src/services/code-index/__tests__/config-manager.test.ts +++ b/src/services/code-index/__tests__/config-manager.test.ts @@ -42,7 +42,6 @@ describe("CodeIndexConfigManager", () => { searchMinScore: 0.4, }) expect(result.requiresRestart).toBe(false) - expect(result.requiresClear).toBe(false) }) it("should load configuration from globalState and secrets", async () => { @@ -98,29 +97,6 @@ describe("CodeIndexConfigManager", () => { const result = await configManager.loadConfiguration() expect(result.requiresRestart).toBe(true) }) - - it("should detect clear requirement when model dimensions change", async () => { - // Initial state with a model - mockContextProxy.getGlobalState.mockReturnValue({ - codebaseIndexEnabled: true, - codebaseIndexQdrantUrl: "http://qdrant.local", - codebaseIndexEmbedderProvider: "openai", - codebaseIndexEmbedderModelId: "text-embedding-3-small", - }) - - await configManager.loadConfiguration() - - // Change to a model with different dimensions - mockContextProxy.getGlobalState.mockReturnValue({ - codebaseIndexEnabled: true, - codebaseIndexQdrantUrl: "http://qdrant.local", - codebaseIndexEmbedderProvider: "openai", - codebaseIndexEmbedderModelId: "text-embedding-3-large", - }) - - const result = await configManager.loadConfiguration() - expect(result.requiresClear).toBe(true) - }) }) describe("isConfigured", () => { diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index e866c4c7a1..78e6b74431 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -1,7 +1,6 @@ import { ApiHandlerOptions } from "../../shared/api" import { ContextProxy } from "../../core/config/ContextProxy" import { EmbedderProvider } from "./interfaces/manager" -import { getModelDimension, getDefaultModelId } from "../../shared/embeddingModels" import { CodeIndexConfig, PreviousConfigSnapshot } from "./interfaces/config" import { SEARCH_MIN_SCORE } from "./constants" @@ -38,7 +37,6 @@ export class CodeIndexConfigManager { searchMinScore?: number } requiresRestart: boolean - requiresClear: boolean }> { const previousConfigSnapshot: PreviousConfigSnapshot = { enabled: this.isEnabled, @@ -84,16 +82,6 @@ export class CodeIndexConfigManager { ollamaBaseUrl: codebaseIndexEmbedderBaseUrl, } - const previousModelId = - previousConfigSnapshot.modelId ?? getDefaultModelId(previousConfigSnapshot.embedderProvider) - const currentModelId = this.modelId ?? getDefaultModelId(this.embedderProvider) - const previousDimension = previousModelId - ? getModelDimension(previousConfigSnapshot.embedderProvider, previousModelId) - : undefined - const currentDimension = currentModelId ? getModelDimension(this.embedderProvider, currentModelId) : undefined - const requiresClear = - previousDimension !== undefined && currentDimension !== undefined && previousDimension !== currentDimension - return { configSnapshot: previousConfigSnapshot, currentConfig: { @@ -108,7 +96,6 @@ export class CodeIndexConfigManager { searchMinScore: this.searchMinScore, }, requiresRestart: this._didConfigChangeRequireRestart(previousConfigSnapshot), - requiresClear, } } diff --git a/src/services/code-index/manager.ts b/src/services/code-index/manager.ts index 6af05f1e18..bbc26605e8 100644 --- a/src/services/code-index/manager.ts +++ b/src/services/code-index/manager.ts @@ -100,7 +100,7 @@ export class CodeIndexManager { public async initialize(contextProxy: ContextProxy): Promise<{ requiresRestart: boolean }> { // 1. ConfigManager Initialization and Configuration Loading this._configManager = new CodeIndexConfigManager(contextProxy) - const { requiresRestart, requiresClear } = await this._configManager.loadConfiguration() + const { requiresRestart } = await this._configManager.loadConfiguration() // 2. Check if feature is enabled if (!this.isFeatureEnabled) { @@ -170,20 +170,11 @@ export class CodeIndexManager { ) } - // 5. Handle Data Clearing - if (requiresClear) { - if (this._orchestrator) { - await this._orchestrator.clearIndexData() - } - if (this._cacheManager) { - await this._cacheManager.clearCacheFile() - } - } - - // Handle Indexing Start/Restart + // 5. Handle Indexing Start/Restart + // The enhanced vectorStore.initialize() in startIndexing() now handles dimension changes automatically + // by detecting incompatible collections and recreating them, so we rely on that for dimension changes const shouldStartOrRestartIndexing = requiresRestart || - requiresClear || (needsServiceRecreation && (!this._orchestrator || this._orchestrator.state !== "Indexing")) if (shouldStartOrRestartIndexing) { diff --git a/src/services/code-index/vector-store/__tests__/qdrant-client.test.ts b/src/services/code-index/vector-store/__tests__/qdrant-client.test.ts new file mode 100644 index 0000000000..ecce86e88f --- /dev/null +++ b/src/services/code-index/vector-store/__tests__/qdrant-client.test.ts @@ -0,0 +1,862 @@ +import { QdrantVectorStore } from "../qdrant-client" +import { QdrantClient } from "@qdrant/js-client-rest" +import { createHash } from "crypto" +import * as path from "path" +import { getWorkspacePath } from "../../../../utils/path" +import { MAX_SEARCH_RESULTS, SEARCH_MIN_SCORE } from "../../constants" +import { Payload, VectorStoreSearchResult } from "../../interfaces" + +// Mocks +jest.mock("@qdrant/js-client-rest") +jest.mock("crypto") +jest.mock("../../../../utils/path") +jest.mock("path", () => ({ + ...jest.requireActual("path"), + sep: "/", +})) + +const mockQdrantClientInstance = { + getCollection: jest.fn(), + createCollection: jest.fn(), + deleteCollection: jest.fn(), + createPayloadIndex: jest.fn(), + upsert: jest.fn(), + query: jest.fn(), + delete: jest.fn(), +} + +const mockCreateHashInstance = { + update: jest.fn().mockReturnThis(), + digest: jest.fn(), +} + +describe("QdrantVectorStore", () => { + let vectorStore: QdrantVectorStore + const mockWorkspacePath = "/test/workspace" + const mockQdrantUrl = "http://mock-qdrant:6333" + const mockApiKey = "test-api-key" + const mockVectorSize = 1536 + const mockHashedPath = "a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6" // Needs to be long enough + const expectedCollectionName = `ws-${mockHashedPath.substring(0, 16)}` + + beforeEach(() => { + jest.clearAllMocks() + + // Mock QdrantClient constructor + ;(QdrantClient as jest.Mock).mockImplementation(() => mockQdrantClientInstance) + + // Mock crypto.createHash + ;(createHash as jest.Mock).mockReturnValue(mockCreateHashInstance) + mockCreateHashInstance.update.mockReturnValue(mockCreateHashInstance) // Ensure it returns 'this' + mockCreateHashInstance.digest.mockReturnValue(mockHashedPath) + + // Mock getWorkspacePath + ;(getWorkspacePath as jest.Mock).mockReturnValue(mockWorkspacePath) + + vectorStore = new QdrantVectorStore(mockWorkspacePath, mockQdrantUrl, mockVectorSize, mockApiKey) + }) + + it("should correctly initialize QdrantClient and collectionName in constructor", () => { + expect(QdrantClient).toHaveBeenCalledTimes(1) + expect(QdrantClient).toHaveBeenCalledWith({ + url: mockQdrantUrl, + apiKey: mockApiKey, + headers: { + "User-Agent": "Roo-Code", + }, + }) + expect(createHash).toHaveBeenCalledWith("sha256") + expect(mockCreateHashInstance.update).toHaveBeenCalledWith(mockWorkspacePath) + expect(mockCreateHashInstance.digest).toHaveBeenCalledWith("hex") + // Access private member for testing constructor logic (not ideal, but necessary here) + expect((vectorStore as any).collectionName).toBe(expectedCollectionName) + expect((vectorStore as any).vectorSize).toBe(mockVectorSize) + }) + it("should handle constructor with default URL when none provided", () => { + const vectorStoreWithDefaults = new QdrantVectorStore(mockWorkspacePath, undefined as any, mockVectorSize) + + expect(QdrantClient).toHaveBeenLastCalledWith({ + url: "http://localhost:6333", // Should use default QDRANT_URL + apiKey: undefined, + headers: { + "User-Agent": "Roo-Code", + }, + }) + }) + + it("should handle constructor without API key", () => { + const vectorStoreWithoutKey = new QdrantVectorStore(mockWorkspacePath, mockQdrantUrl, mockVectorSize) + + expect(QdrantClient).toHaveBeenLastCalledWith({ + url: mockQdrantUrl, + apiKey: undefined, + headers: { + "User-Agent": "Roo-Code", + }, + }) + }) + + describe("initialize", () => { + it("should create a new collection if none exists and return true", async () => { + // Mock getCollection to throw a 404-like error + mockQdrantClientInstance.getCollection.mockRejectedValue({ + response: { status: 404 }, + message: "Not found", + }) + mockQdrantClientInstance.createCollection.mockResolvedValue(true as any) // Cast to any to satisfy QdrantClient types if strict + mockQdrantClientInstance.createPayloadIndex.mockResolvedValue({} as any) // Mock successful index creation + + const result = await vectorStore.initialize() + + expect(result).toBe(true) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledWith(expectedCollectionName) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledWith(expectedCollectionName, { + vectors: { + size: mockVectorSize, + distance: "Cosine", // Assuming 'Cosine' is the DISTANCE_METRIC + }, + }) + expect(mockQdrantClientInstance.deleteCollection).not.toHaveBeenCalled() + + // Verify payload index creation + for (let i = 0; i <= 4; i++) { + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledWith(expectedCollectionName, { + field_name: `pathSegments.${i}`, + field_schema: "keyword", + }) + } + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledTimes(5) + }) + it("should not create a new collection if one exists with matching vectorSize and return false", async () => { + // Mock getCollection to return existing collection info with matching vector size + mockQdrantClientInstance.getCollection.mockResolvedValue({ + config: { + params: { + vectors: { + size: mockVectorSize, // Matching vector size + }, + }, + }, + } as any) // Cast to any to satisfy QdrantClient types + mockQdrantClientInstance.createPayloadIndex.mockResolvedValue({} as any) + + const result = await vectorStore.initialize() + + expect(result).toBe(false) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledWith(expectedCollectionName) + expect(mockQdrantClientInstance.createCollection).not.toHaveBeenCalled() + expect(mockQdrantClientInstance.deleteCollection).not.toHaveBeenCalled() + + // Verify payload index creation still happens + for (let i = 0; i <= 4; i++) { + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledWith(expectedCollectionName, { + field_name: `pathSegments.${i}`, + field_schema: "keyword", + }) + } + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledTimes(5) + }) + it("should recreate collection if it exists but vectorSize mismatches and return true", async () => { + const differentVectorSize = 768 + // Mock getCollection to return existing collection info with different vector size + mockQdrantClientInstance.getCollection.mockResolvedValue({ + config: { + params: { + vectors: { + size: differentVectorSize, // Mismatching vector size + }, + }, + }, + } as any) + mockQdrantClientInstance.deleteCollection.mockResolvedValue(true as any) + mockQdrantClientInstance.createCollection.mockResolvedValue(true as any) + mockQdrantClientInstance.createPayloadIndex.mockResolvedValue({} as any) + jest.spyOn(console, "warn").mockImplementation(() => {}) // Suppress console.warn + + const result = await vectorStore.initialize() + + expect(result).toBe(true) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledWith(expectedCollectionName) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledWith(expectedCollectionName) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledWith(expectedCollectionName, { + vectors: { + size: mockVectorSize, // Should use the new, correct vector size + distance: "Cosine", + }, + }) + + // Verify payload index creation + for (let i = 0; i <= 4; i++) { + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledWith(expectedCollectionName, { + field_name: `pathSegments.${i}`, + field_schema: "keyword", + }) + } + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledTimes(5) + ;(console.warn as jest.Mock).mockRestore() // Restore console.warn + }) + it("should re-throw error from getCollection if it is not a 404 error", async () => { + const genericError = new Error("Generic Qdrant Error") + mockQdrantClientInstance.getCollection.mockRejectedValue(genericError) + jest.spyOn(console, "error").mockImplementation(() => {}) // Suppress console.error + + await expect(vectorStore.initialize()).rejects.toThrow(genericError) + + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.createCollection).not.toHaveBeenCalled() + expect(mockQdrantClientInstance.deleteCollection).not.toHaveBeenCalled() + expect(mockQdrantClientInstance.createPayloadIndex).not.toHaveBeenCalled() + expect(console.error).toHaveBeenCalledTimes(2) // Once in the try/catch for getCollection, once in the outer try/catch + ;(console.error as jest.Mock).mockRestore() + }) + it("should re-throw error from createCollection when no collection initially exists", async () => { + mockQdrantClientInstance.getCollection.mockRejectedValue({ + response: { status: 404 }, + message: "Not found", + }) + const createError = new Error("Create Collection Failed") + mockQdrantClientInstance.createCollection.mockRejectedValue(createError) + jest.spyOn(console, "error").mockImplementation(() => {}) // Suppress console.error + + await expect(vectorStore.initialize()).rejects.toThrow(createError) + + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).not.toHaveBeenCalled() + expect(mockQdrantClientInstance.createPayloadIndex).not.toHaveBeenCalled() // Should not be called if createCollection fails + expect(console.error).toHaveBeenCalledTimes(1) // Only the outer try/catch + ;(console.error as jest.Mock).mockRestore() + }) + it("should log but not fail if payload index creation errors occur", async () => { + // Mock successful collection creation + mockQdrantClientInstance.getCollection.mockRejectedValue({ + response: { status: 404 }, + message: "Not found", + }) + mockQdrantClientInstance.createCollection.mockResolvedValue(true as any) + + // Mock payload index creation to fail + const indexError = new Error("Index creation failed") + mockQdrantClientInstance.createPayloadIndex.mockRejectedValue(indexError) + jest.spyOn(console, "warn").mockImplementation(() => {}) // Suppress console.warn + + const result = await vectorStore.initialize() + + // Should still return true since main collection setup succeeded + expect(result).toBe(true) + expect(mockQdrantClientInstance.createCollection).toHaveBeenCalledTimes(1) + + // Verify all payload index creations were attempted + expect(mockQdrantClientInstance.createPayloadIndex).toHaveBeenCalledTimes(5) + + // Verify warnings were logged for each failed index + expect(console.warn).toHaveBeenCalledTimes(5) + for (let i = 0; i <= 4; i++) { + expect(console.warn).toHaveBeenCalledWith( + expect.stringContaining(`Could not create payload index for pathSegments.${i}`), + indexError, + ) + } + + ;(console.warn as jest.Mock).mockRestore() + }) + + it("should re-throw error from deleteCollection when recreating collection with mismatched vectorSize", async () => { + const differentVectorSize = 768 + mockQdrantClientInstance.getCollection.mockResolvedValue({ + config: { + params: { + vectors: { + size: differentVectorSize, + }, + }, + }, + } as any) + + const deleteError = new Error("Delete Collection Failed") + mockQdrantClientInstance.deleteCollection.mockRejectedValue(deleteError) + jest.spyOn(console, "error").mockImplementation(() => {}) + jest.spyOn(console, "warn").mockImplementation(() => {}) + + await expect(vectorStore.initialize()).rejects.toThrow(deleteError) + + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.createCollection).not.toHaveBeenCalled() + expect(mockQdrantClientInstance.createPayloadIndex).not.toHaveBeenCalled() + ;(console.error as jest.Mock).mockRestore() + ;(console.warn as jest.Mock).mockRestore() + }) + }) + + it("should return true when collection exists", async () => { + mockQdrantClientInstance.getCollection.mockResolvedValue({ + config: { + /* collection data */ + }, + } as any) + + const result = await vectorStore.collectionExists() + + expect(result).toBe(true) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledWith(expectedCollectionName) + }) + + it("should return false when collection does not exist (404 error)", async () => { + mockQdrantClientInstance.getCollection.mockRejectedValue({ + response: { status: 404 }, + message: "Not found", + }) + + const result = await vectorStore.collectionExists() + + expect(result).toBe(false) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledWith(expectedCollectionName) + }) + + it("should return false and log error for non-404 errors", async () => { + const genericError = new Error("Network error") + mockQdrantClientInstance.getCollection.mockRejectedValue(genericError) + jest.spyOn(console, "error").mockImplementation(() => {}) + + const result = await vectorStore.collectionExists() + + expect(result).toBe(false) + expect(mockQdrantClientInstance.getCollection).toHaveBeenCalledTimes(1) + expect(console.error).toHaveBeenCalledWith("Error checking collection existence:", genericError) + ;(console.error as jest.Mock).mockRestore() + }) + describe("collectionExists", () => { + // Test scenarios for collectionExists will go here + it("should delete collection when it exists", async () => { + // Mock collectionExists to return true + jest.spyOn(vectorStore, "collectionExists").mockResolvedValue(true) + mockQdrantClientInstance.deleteCollection.mockResolvedValue(true as any) + + await vectorStore.deleteCollection() + + expect(vectorStore.collectionExists).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledWith(expectedCollectionName) + }) + + it("should not attempt to delete collection when it does not exist", async () => { + // Mock collectionExists to return false + jest.spyOn(vectorStore, "collectionExists").mockResolvedValue(false) + + await vectorStore.deleteCollection() + + expect(vectorStore.collectionExists).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).not.toHaveBeenCalled() + }) + + it("should log and re-throw error when deletion fails", async () => { + jest.spyOn(vectorStore, "collectionExists").mockResolvedValue(true) + const deleteError = new Error("Deletion failed") + mockQdrantClientInstance.deleteCollection.mockRejectedValue(deleteError) + jest.spyOn(console, "error").mockImplementation(() => {}) + + await expect(vectorStore.deleteCollection()).rejects.toThrow(deleteError) + + expect(vectorStore.collectionExists).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.deleteCollection).toHaveBeenCalledTimes(1) + expect(console.error).toHaveBeenCalledWith( + `[QdrantVectorStore] Failed to delete collection ${expectedCollectionName}:`, + deleteError, + ) + ;(console.error as jest.Mock).mockRestore() + }) + }) + + describe("deleteCollection", () => { + // Test scenarios for deleteCollection will go here + }) + + describe("upsertPoints", () => { + it("should correctly call qdrantClient.upsert with processed points", async () => { + const mockPoints = [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + filePath: "src/components/Button.tsx", + content: "export const Button = () => {}", + startLine: 1, + endLine: 3, + }, + }, + { + id: "test-id-2", + vector: [0.4, 0.5, 0.6], + payload: { + filePath: "src/utils/helpers.ts", + content: "export function helper() {}", + startLine: 5, + endLine: 7, + }, + }, + ] + + mockQdrantClientInstance.upsert.mockResolvedValue({} as any) + + await vectorStore.upsertPoints(mockPoints) + + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledWith(expectedCollectionName, { + points: [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + filePath: "src/components/Button.tsx", + content: "export const Button = () => {}", + startLine: 1, + endLine: 3, + pathSegments: { + "0": "src", + "1": "components", + "2": "Button.tsx", + }, + }, + }, + { + id: "test-id-2", + vector: [0.4, 0.5, 0.6], + payload: { + filePath: "src/utils/helpers.ts", + content: "export function helper() {}", + startLine: 5, + endLine: 7, + pathSegments: { + "0": "src", + "1": "utils", + "2": "helpers.ts", + }, + }, + }, + ], + wait: true, + }) + }) + + it("should handle points without filePath in payload", async () => { + const mockPoints = [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + content: "some content without filePath", + startLine: 1, + endLine: 3, + }, + }, + ] + + mockQdrantClientInstance.upsert.mockResolvedValue({} as any) + + await vectorStore.upsertPoints(mockPoints) + + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledWith(expectedCollectionName, { + points: [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + content: "some content without filePath", + startLine: 1, + endLine: 3, + }, + }, + ], + wait: true, + }) + }) + + it("should handle empty input arrays", async () => { + mockQdrantClientInstance.upsert.mockResolvedValue({} as any) + + await vectorStore.upsertPoints([]) + + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledWith(expectedCollectionName, { + points: [], + wait: true, + }) + }) + + it("should correctly process pathSegments for nested file paths", async () => { + const mockPoints = [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + filePath: "src/components/ui/forms/InputField.tsx", + content: "export const InputField = () => {}", + startLine: 1, + endLine: 3, + }, + }, + ] + + mockQdrantClientInstance.upsert.mockResolvedValue({} as any) + + await vectorStore.upsertPoints(mockPoints) + + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledWith(expectedCollectionName, { + points: [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + filePath: "src/components/ui/forms/InputField.tsx", + content: "export const InputField = () => {}", + startLine: 1, + endLine: 3, + pathSegments: { + "0": "src", + "1": "components", + "2": "ui", + "3": "forms", + "4": "InputField.tsx", + }, + }, + }, + ], + wait: true, + }) + }) + + it("should handle error scenarios when qdrantClient.upsert fails", async () => { + const mockPoints = [ + { + id: "test-id-1", + vector: [0.1, 0.2, 0.3], + payload: { + filePath: "src/test.ts", + content: "test content", + startLine: 1, + endLine: 1, + }, + }, + ] + + const upsertError = new Error("Upsert failed") + mockQdrantClientInstance.upsert.mockRejectedValue(upsertError) + jest.spyOn(console, "error").mockImplementation(() => {}) + + await expect(vectorStore.upsertPoints(mockPoints)).rejects.toThrow(upsertError) + + expect(mockQdrantClientInstance.upsert).toHaveBeenCalledTimes(1) + expect(console.error).toHaveBeenCalledWith("Failed to upsert points:", upsertError) + ;(console.error as jest.Mock).mockRestore() + }) + }) + + describe("search", () => { + it("should correctly call qdrantClient.query and transform results", async () => { + const queryVector = [0.1, 0.2, 0.3] + const mockQdrantResults = { + points: [ + { + id: "test-id-1", + score: 0.85, + payload: { + filePath: "src/test.ts", + codeChunk: "test code", + startLine: 1, + endLine: 5, + pathSegments: { "0": "src", "1": "test.ts" }, + }, + }, + { + id: "test-id-2", + score: 0.75, + payload: { + filePath: "src/utils.ts", + codeChunk: "utility code", + startLine: 10, + endLine: 15, + pathSegments: { "0": "src", "1": "utils.ts" }, + }, + }, + ], + } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + const results = await vectorStore.search(queryVector) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledTimes(1) + expect(mockQdrantClientInstance.query).toHaveBeenCalledWith(expectedCollectionName, { + query: queryVector, + filter: undefined, + score_threshold: SEARCH_MIN_SCORE, + limit: MAX_SEARCH_RESULTS, + params: { + hnsw_ef: 128, + exact: false, + }, + with_payload: { + include: ["filePath", "codeChunk", "startLine", "endLine", "pathSegments"], + }, + }) + + expect(results).toEqual(mockQdrantResults.points) + }) + + it("should apply filePathPrefix filter correctly", async () => { + const queryVector = [0.1, 0.2, 0.3] + const directoryPrefix = "src/components" + const mockQdrantResults = { + points: [ + { + id: "test-id-1", + score: 0.85, + payload: { + filePath: "src/components/Button.tsx", + codeChunk: "button code", + startLine: 1, + endLine: 5, + pathSegments: { "0": "src", "1": "components", "2": "Button.tsx" }, + }, + }, + ], + } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + const results = await vectorStore.search(queryVector, directoryPrefix) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledWith(expectedCollectionName, { + query: queryVector, + filter: { + must: [ + { + key: "pathSegments.0", + match: { value: "src" }, + }, + { + key: "pathSegments.1", + match: { value: "components" }, + }, + ], + }, + score_threshold: SEARCH_MIN_SCORE, + limit: MAX_SEARCH_RESULTS, + params: { + hnsw_ef: 128, + exact: false, + }, + with_payload: { + include: ["filePath", "codeChunk", "startLine", "endLine", "pathSegments"], + }, + }) + + expect(results).toEqual(mockQdrantResults.points) + }) + + it("should use custom minScore when provided", async () => { + const queryVector = [0.1, 0.2, 0.3] + const customMinScore = 0.8 + const mockQdrantResults = { points: [] } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + await vectorStore.search(queryVector, undefined, customMinScore) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledWith(expectedCollectionName, { + query: queryVector, + filter: undefined, + score_threshold: customMinScore, + limit: MAX_SEARCH_RESULTS, + params: { + hnsw_ef: 128, + exact: false, + }, + with_payload: { + include: ["filePath", "codeChunk", "startLine", "endLine", "pathSegments"], + }, + }) + }) + + it("should filter out results with invalid payloads", async () => { + const queryVector = [0.1, 0.2, 0.3] + const mockQdrantResults = { + points: [ + { + id: "valid-result", + score: 0.85, + payload: { + filePath: "src/test.ts", + codeChunk: "test code", + startLine: 1, + endLine: 5, + }, + }, + { + id: "invalid-result-1", + score: 0.75, + payload: { + // Missing required fields + filePath: "src/invalid.ts", + }, + }, + { + id: "valid-result-2", + score: 0.55, + payload: { + filePath: "src/test2.ts", + codeChunk: "test code 2", + startLine: 10, + endLine: 15, + }, + }, + ], + } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + const results = await vectorStore.search(queryVector) + + // Should only return results with valid payloads + expect(results).toHaveLength(2) + expect(results[0].id).toBe("valid-result") + expect(results[1].id).toBe("valid-result-2") + }) + + it("should filter out results with null or undefined payloads", async () => { + const queryVector = [0.1, 0.2, 0.3] + const mockQdrantResults = { + points: [ + { + id: "valid-result", + score: 0.85, + payload: { + filePath: "src/test.ts", + codeChunk: "test code", + startLine: 1, + endLine: 5, + }, + }, + { + id: "null-payload-result", + score: 0.75, + payload: null, + }, + { + id: "undefined-payload-result", + score: 0.65, + payload: undefined, + }, + { + id: "valid-result-2", + score: 0.55, + payload: { + filePath: "src/test2.ts", + codeChunk: "test code 2", + startLine: 10, + endLine: 15, + }, + }, + ], + } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + const results = await vectorStore.search(queryVector) + + // Should only return results with valid payloads, filtering out null and undefined + expect(results).toHaveLength(2) + expect(results[0].id).toBe("valid-result") + expect(results[1].id).toBe("valid-result-2") + }) + + it("should handle scenarios where no results are found", async () => { + const queryVector = [0.1, 0.2, 0.3] + const mockQdrantResults = { points: [] } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + const results = await vectorStore.search(queryVector) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledTimes(1) + expect(results).toEqual([]) + }) + + it("should handle complex directory prefix with multiple segments", async () => { + const queryVector = [0.1, 0.2, 0.3] + const directoryPrefix = "src/components/ui/forms" + const mockQdrantResults = { points: [] } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + await vectorStore.search(queryVector, directoryPrefix) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledWith(expectedCollectionName, { + query: queryVector, + filter: { + must: [ + { + key: "pathSegments.0", + match: { value: "src" }, + }, + { + key: "pathSegments.1", + match: { value: "components" }, + }, + { + key: "pathSegments.2", + match: { value: "ui" }, + }, + { + key: "pathSegments.3", + match: { value: "forms" }, + }, + ], + }, + score_threshold: SEARCH_MIN_SCORE, + limit: MAX_SEARCH_RESULTS, + params: { + hnsw_ef: 128, + exact: false, + }, + with_payload: { + include: ["filePath", "codeChunk", "startLine", "endLine", "pathSegments"], + }, + }) + }) + + it("should handle error scenarios when qdrantClient.query fails", async () => { + const queryVector = [0.1, 0.2, 0.3] + const queryError = new Error("Query failed") + mockQdrantClientInstance.query.mockRejectedValue(queryError) + jest.spyOn(console, "error").mockImplementation(() => {}) + + await expect(vectorStore.search(queryVector)).rejects.toThrow(queryError) + + expect(mockQdrantClientInstance.query).toHaveBeenCalledTimes(1) + expect(console.error).toHaveBeenCalledWith("Failed to search points:", queryError) + ;(console.error as jest.Mock).mockRestore() + }) + + it("should use constants MAX_SEARCH_RESULTS and SEARCH_MIN_SCORE correctly", async () => { + const queryVector = [0.1, 0.2, 0.3] + const mockQdrantResults = { points: [] } + + mockQdrantClientInstance.query.mockResolvedValue(mockQdrantResults) + + await vectorStore.search(queryVector) + + const callArgs = mockQdrantClientInstance.query.mock.calls[0][1] + expect(callArgs.limit).toBe(MAX_SEARCH_RESULTS) + expect(callArgs.score_threshold).toBe(SEARCH_MIN_SCORE) + }) + }) +}) diff --git a/src/services/code-index/vector-store/qdrant-client.ts b/src/services/code-index/vector-store/qdrant-client.ts index 38d40d943e..247fa1f9ff 100644 --- a/src/services/code-index/vector-store/qdrant-client.ts +++ b/src/services/code-index/vector-store/qdrant-client.ts @@ -44,19 +44,47 @@ export class QdrantVectorStore implements IVectorStore { async initialize(): Promise { try { let created = false - const collections = await this.client.getCollections() - const collectionExists = collections.collections.some( - (collection) => collection.name === this.collectionName, - ) - if (!collectionExists) { - await this.client.createCollection(this.collectionName, { - vectors: { - size: this.vectorSize, - distance: this.DISTANCE_METRIC, - }, - }) - created = true + try { + // Directly attempt to fetch the specific collection + const collectionInfo = await this.client.getCollection(this.collectionName) + + // Collection exists - check if vector size matches + const existingVectorSize = collectionInfo.config?.params?.vectors?.size + + if (existingVectorSize === this.vectorSize) { + // Collection exists and has correct vector size + created = false + } else { + // Collection exists but has wrong vector size - recreate it + console.warn( + `[QdrantVectorStore] Collection ${this.collectionName} exists with vector size ${existingVectorSize}, but expected ${this.vectorSize}. Recreating collection.`, + ) + await this.client.deleteCollection(this.collectionName) + await this.client.createCollection(this.collectionName, { + vectors: { + size: this.vectorSize, + distance: this.DISTANCE_METRIC, + }, + }) + created = true + } + } catch (error: any) { + // Check if this is a "Not Found" error (collection doesn't exist) + if (error?.response?.status === 404) { + // Collection doesn't exist - create it + await this.client.createCollection(this.collectionName, { + vectors: { + size: this.vectorSize, + distance: this.DISTANCE_METRIC, + }, + }) + created = true + } else { + // Other error - log and re-throw + console.error(`[QdrantVectorStore] Error checking collection ${this.collectionName}:`, error) + throw error + } } // Create payload indexes for pathSegments up to depth 5 @@ -129,8 +157,13 @@ export class QdrantVectorStore implements IVectorStore { * @param payload Payload to check * @returns Boolean indicating if the payload is valid */ - private isPayloadValid(payload: Record): payload is Payload { - return "filePath" in payload && "codeChunk" in payload && "startLine" in payload && "endLine" in payload + private isPayloadValid(payload: Record | null | undefined): payload is Payload { + if (!payload) { + return false + } + const validKeys = ["filePath", "codeChunk", "startLine", "endLine"] + const hasValidKeys = validKeys.every((key) => key in payload) + return hasValidKeys } /** @@ -177,7 +210,7 @@ export class QdrantVectorStore implements IVectorStore { } const operationResult = await this.client.query(this.collectionName, searchRequest) - const filteredPoints = operationResult.points.filter((p) => this.isPayloadValid(p.payload!)) + const filteredPoints = operationResult.points.filter((p) => this.isPayloadValid(p.payload)) return filteredPoints as VectorStoreSearchResult[] } catch (error) {