diff --git a/JS/edgechains/arakoodev/src/vector-db/src/index.ts b/JS/edgechains/arakoodev/src/vector-db/src/index.ts index 557104a14..5ca60f642 100644 --- a/JS/edgechains/arakoodev/src/vector-db/src/index.ts +++ b/JS/edgechains/arakoodev/src/vector-db/src/index.ts @@ -1 +1,2 @@ export { Supabase } from "./lib/supabase/supabase.js"; +export { Qdrant, QdrantDistanceMetric } from "./lib/qdrant/qdrant.js"; diff --git a/JS/edgechains/arakoodev/src/vector-db/src/lib/qdrant/qdrant.ts b/JS/edgechains/arakoodev/src/vector-db/src/lib/qdrant/qdrant.ts new file mode 100644 index 000000000..b4af073fc --- /dev/null +++ b/JS/edgechains/arakoodev/src/vector-db/src/lib/qdrant/qdrant.ts @@ -0,0 +1,352 @@ +import axios, { AxiosInstance, AxiosRequestConfig } from "axios"; +import retry from "retry"; +import { config } from "dotenv"; +config(); + +type QdrantPointId = number | string; +type QdrantVector = number[] | Record; +type QdrantSearchVector = QdrantVector | { name: string; vector: number[] }; +type QdrantPayload = Record; + +interface QdrantPoint { + id: QdrantPointId; + vector: QdrantVector; + payload?: QdrantPayload; +} + +interface QdrantRequestOptions { + wait?: boolean; + ordering?: "weak" | "medium" | "strong"; +} + +interface CreateCollectionArgs { + client: AxiosInstance; + collectionName: string; + vectorSize?: number; + distance?: QdrantDistanceMetric; + vectorsConfig?: Record; + [key: string]: any; +} + +interface UpsertPointsArgs extends QdrantRequestOptions { + client: AxiosInstance; + collectionName: string; + points: QdrantPoint[]; +} + +interface InsertVectorDataArgs extends QdrantRequestOptions { + client: AxiosInstance; + collectionName: string; + points?: QdrantPoint[]; + id?: QdrantPointId; + vector?: QdrantVector; + payload?: QdrantPayload; +} + +interface SearchPointsArgs { + client: AxiosInstance; + collectionName: string; + vector: QdrantSearchVector; + limit?: number; + offset?: QdrantPointId; + filter?: Record; + params?: Record; + withPayload?: boolean | string[] | Record; + withVector?: boolean | string[] | Record; + scoreThreshold?: number; +} + +interface ScrollPointsArgs { + client: AxiosInstance; + collectionName: string; + limit?: number; + offset?: QdrantPointId; + filter?: Record; + withPayload?: boolean | string[] | Record; + withVector?: boolean | string[] | Record; +} + +interface GetPointByIdArgs { + client: AxiosInstance; + collectionName: string; + id: QdrantPointId; + withPayload?: boolean | string[]; + withVector?: boolean | string[]; +} + +interface DeletePointsArgs extends QdrantRequestOptions { + client: AxiosInstance; + collectionName: string; + points?: QdrantPointId[]; + filter?: Record; +} + +export class Qdrant { + QDRANT_URL: string; + QDRANT_API_KEY?: string; + + constructor(QDRANT_URL?: string, QDRANT_API_KEY?: string) { + this.QDRANT_URL = + QDRANT_URL || process.env.QDRANT_URL || "http://localhost:6333"; + this.QDRANT_API_KEY = QDRANT_API_KEY || process.env.QDRANT_API_KEY; + } + + createClient() { + const headers: Record = { + "Content-Type": "application/json", + }; + + if (this.QDRANT_API_KEY) { + headers["api-key"] = this.QDRANT_API_KEY; + } + + return axios.create({ + baseURL: this.QDRANT_URL.replace(/\/+$/, ""), + headers, + }); + } + + async createCollection({ + client, + collectionName, + vectorSize, + distance = QdrantDistanceMetric.COSINE, + vectorsConfig, + ...args + }: CreateCollectionArgs): Promise { + if (!vectorsConfig && !vectorSize) { + throw new Error( + "Either vectorSize or vectorsConfig is required to create a collection", + ); + } + + const body = { + vectors: vectorsConfig || { + size: vectorSize, + distance, + }, + ...args, + }; + + return this.requestWithRetry(async () => { + const response = await client.put( + this.collectionPath(collectionName), + body, + ); + return response.data; + }); + } + + async upsertPoints({ + client, + collectionName, + points, + wait, + ordering, + }: UpsertPointsArgs): Promise { + return this.requestWithRetry(async () => { + const response = await client.put( + `${this.collectionPath(collectionName)}/points`, + { points }, + { params: this.requestParams({ wait, ordering }) }, + ); + return response.data; + }); + } + + async insertVectorData({ + points, + id, + vector, + payload, + ...args + }: InsertVectorDataArgs): Promise { + const qdrantPoints = + points || + (id !== undefined && vector ? [{ id, vector, payload }] : undefined); + + if (!qdrantPoints) { + throw new Error( + "Either points or both id and vector are required to insert Qdrant data", + ); + } + + return this.upsertPoints({ + ...args, + points: qdrantPoints, + }); + } + + async searchPoints({ + client, + collectionName, + vector, + limit = 10, + offset, + filter, + params, + withPayload = true, + withVector = false, + scoreThreshold, + }: SearchPointsArgs): Promise { + const body = this.cleanBody({ + vector, + limit, + offset, + filter, + params, + with_payload: withPayload, + with_vector: withVector, + score_threshold: scoreThreshold, + }); + + return this.requestWithRetry(async () => { + const response = await client.post( + `${this.collectionPath(collectionName)}/points/search`, + body, + ); + return response.data; + }); + } + + async getData({ + client, + collectionName, + limit = 10, + offset, + filter, + withPayload = true, + withVector = false, + }: ScrollPointsArgs): Promise { + const body = this.cleanBody({ + limit, + offset, + filter, + with_payload: withPayload, + with_vector: withVector, + }); + + return this.requestWithRetry(async () => { + const response = await client.post( + `${this.collectionPath(collectionName)}/points/scroll`, + body, + ); + return response.data; + }); + } + + async getDataById({ + client, + collectionName, + id, + withPayload = true, + withVector = false, + }: GetPointByIdArgs): Promise { + return this.requestWithRetry(async () => { + const response = await client.get( + `${this.collectionPath(collectionName)}/points/${encodeURIComponent(id)}`, + { + params: { + with_payload: withPayload, + with_vector: withVector, + }, + }, + ); + return response.data; + }); + } + + async deleteById({ + client, + collectionName, + id, + wait, + ordering, + }: { + client: AxiosInstance; + collectionName: string; + id: QdrantPointId; + } & QdrantRequestOptions): Promise { + return this.deletePoints({ + client, + collectionName, + points: [id], + wait, + ordering, + }); + } + + async deletePoints({ + client, + collectionName, + points, + filter, + wait, + ordering, + }: DeletePointsArgs): Promise { + if (!points && !filter) { + throw new Error( + "Either points or filter is required to delete Qdrant points", + ); + } + + return this.requestWithRetry(async () => { + const response = await client.post( + `${this.collectionPath(collectionName)}/points/delete`, + points ? { points } : { filter }, + { params: this.requestParams({ wait, ordering }) }, + ); + return response.data; + }); + } + + private collectionPath(collectionName: string) { + return `/collections/${encodeURIComponent(collectionName)}`; + } + + private requestParams(options: QdrantRequestOptions) { + return this.cleanBody(options); + } + + private cleanBody>(body: T) { + return Object.fromEntries( + Object.entries(body).filter(([, value]) => value !== undefined), + ); + } + + private async requestWithRetry(request: () => Promise): Promise { + return new Promise((resolve, reject) => { + const operation = retry.operation({ + retries: 5, + factor: 3, + minTimeout: 1 * 1000, + maxTimeout: 60 * 1000, + randomize: true, + }); + + operation.attempt(async () => { + try { + resolve(await request()); + } catch (error: any) { + const retryError = this.toError(error); + if (operation.retry(retryError)) return; + reject(operation.mainError() || retryError); + } + }); + }); + } + + private toError(error: any) { + if (error instanceof Error) return error; + + const message = typeof error === "string" ? error : JSON.stringify(error); + return new Error(message); + } +} + +export enum QdrantDistanceMetric { + COSINE = "Cosine", + DOT = "Dot", + EUCLID = "Euclid", + MANHATTAN = "Manhattan", +} diff --git a/JS/edgechains/arakoodev/src/vector-db/src/tests/qdrant/qdrant.test.ts b/JS/edgechains/arakoodev/src/vector-db/src/tests/qdrant/qdrant.test.ts new file mode 100644 index 000000000..525fd3497 --- /dev/null +++ b/JS/edgechains/arakoodev/src/vector-db/src/tests/qdrant/qdrant.test.ts @@ -0,0 +1,149 @@ +import { + Qdrant, + QdrantDistanceMetric, +} from "../../../../../dist/vector-db/src/lib/qdrant/qdrant.js"; + +const MOCK_QDRANT_API_KEY = "mock-api-key"; +const MOCK_QDRANT_URL = "https://mock-qdrant.local/"; + +describe("Qdrant", () => { + let qdrant: Qdrant; + + beforeEach(() => { + qdrant = new Qdrant(MOCK_QDRANT_URL, MOCK_QDRANT_API_KEY); + }); + + it("should create a REST client with Qdrant defaults", () => { + const client = qdrant.createClient(); + + expect(client.defaults.baseURL).toEqual("https://mock-qdrant.local"); + expect(client.defaults.headers["Content-Type"]).toEqual("application/json"); + expect(client.defaults.headers["api-key"]).toEqual(MOCK_QDRANT_API_KEY); + }); + + it("should create a collection through the REST API", async () => { + const client = { + put: jest.fn().mockResolvedValue({ data: { status: "ok" } }), + } as any; + + const res = await qdrant.createCollection({ + client, + collectionName: "documents", + vectorSize: 1536, + distance: QdrantDistanceMetric.COSINE, + }); + + expect(res).toEqual({ status: "ok" }); + expect(client.put).toHaveBeenCalledWith("/collections/documents", { + vectors: { + size: 1536, + distance: "Cosine", + }, + }); + }); + + it("should upsert points through the REST API", async () => { + const client = { + put: jest + .fn() + .mockResolvedValue({ data: { result: { status: "acknowledged" } } }), + } as any; + + const points = [ + { + id: 1, + vector: [0.1, 0.2, 0.3], + payload: { raw_text: "hello" }, + }, + ]; + + const res = await qdrant.upsertPoints({ + client, + collectionName: "documents", + points, + wait: true, + }); + + expect(res).toEqual({ result: { status: "acknowledged" } }); + expect(client.put).toHaveBeenCalledWith( + "/collections/documents/points", + { points }, + { params: { wait: true } }, + ); + }); + + it("should search points through the REST API", async () => { + const client = { + post: jest + .fn() + .mockResolvedValue({ data: { result: [{ id: 1, score: 0.9 }] } }), + } as any; + + const res = await qdrant.searchPoints({ + client, + collectionName: "documents", + vector: [0.1, 0.2, 0.3], + limit: 5, + filter: { must: [{ key: "namespace", match: { value: "docs" } }] }, + }); + + expect(res).toEqual({ result: [{ id: 1, score: 0.9 }] }); + expect(client.post).toHaveBeenCalledWith( + "/collections/documents/points/search", + { + vector: [0.1, 0.2, 0.3], + limit: 5, + filter: { must: [{ key: "namespace", match: { value: "docs" } }] }, + with_payload: true, + with_vector: false, + }, + ); + }); + + it("should search points with a named vector through the REST API", async () => { + const client = { + post: jest + .fn() + .mockResolvedValue({ data: { result: [{ id: 1, score: 0.9 }] } }), + } as any; + + const vector = { name: "content", vector: [0.1, 0.2, 0.3] }; + + const res = await qdrant.searchPoints({ + client, + collectionName: "documents", + vector, + }); + + expect(res).toEqual({ result: [{ id: 1, score: 0.9 }] }); + expect(client.post).toHaveBeenCalledWith( + "/collections/documents/points/search", + { + vector, + limit: 10, + with_payload: true, + with_vector: false, + }, + ); + }); + + it("should delete points by id through the REST API", async () => { + const client = { + post: jest.fn().mockResolvedValue({ data: { status: "ok" } }), + } as any; + + const res = await qdrant.deleteById({ + client, + collectionName: "documents", + id: "point-1", + ordering: "strong", + }); + + expect(res).toEqual({ status: "ok" }); + expect(client.post).toHaveBeenCalledWith( + "/collections/documents/points/delete", + { points: ["point-1"] }, + { params: { ordering: "strong" } }, + ); + }); +});