Merge pull request #7 from JacobLinCool/feat-add-qdrant-vec

feat: add qdrant vector store backend

README.md

@@ -1,6 +1,6 @@
 # KVec
 
-Use Cloudflare KV or Upstash Redis with OpenAI/Cohere Text Embedding and Pinecone Vector Database.
+Use Cloudflare KV or Upstash Redis with OpenAI/Cohere Text Embedding and Pinecone/Qdrant Vector Database.
 
 ![icon](static/icon.png)
 
@@ -244,10 +244,12 @@ Currently, the following implementations are available:
   - [x] `MemoryObjStore`: Only for local development
 - **VecStore**
   - [x] `PineconeVecStore`: Use [Pinecone](https://www.pinecone.io/) as the vector store backend
+  - [x] `QdrantVecStore`: Use [Qdrant](https://qdrant.tech/) as the vector store backend
   - [x] `MemoryVecStore`: Only for local development
 - **Cache**
   - [x] `CloudflareCache`: Use Cloudflare's Cache API
   - [x] `MemoryCache`: Only for local development, it just "don't cache anything"
 
 > The auto module will automatically load the correct implementation based on the environment variables.
 > See [src/lib/server/auto/index.ts](./src/lib/server/auto/index.ts)
+> You can also take a look at [the README of each module](./src/lib/server) to see how to configure them.

src/lib/server/auto/index.ts

@@ -7,6 +7,7 @@ import { DEFAULT_K } from "../constants";
 import { MemoryObjStore, MemoryVecStore, MemoryCache, JustEncoder } from "../local";
 import { OpenAIEncoder } from "../openai";
 import { PineconeVecStore } from "../pinecone";
+import { QdrantVecStore } from "../qdrant";
 import { UpstashRedisObjStore } from "../upstash";
 
 export const AutoAdapter = BaseTextAdapter;
@@ -19,7 +20,11 @@ export const AutoObjStore =
 		: MemoryObjStore;
 
 export const AutoVecStore =
-	env.PINECONE_API_KEY && env.PINECONE_ENDPOINT ? PineconeVecStore : MemoryVecStore;
+	env.PINECONE_API_KEY && env.PINECONE_ENDPOINT
+		? PineconeVecStore
+		: env.QDRANT_SERVER && env.QDRANT_COLLECTION
+		? QdrantVecStore
+		: MemoryVecStore;
 
 export const AutoEncoder = env.OPENAI_API_KEY
 	? OpenAIEncoder

src/lib/server/qdrant/README.md

@@ -0,0 +1,9 @@
+# Qdrant
+
+## QdrantVecStore
+
+Enable when `QDRANT_SERVER` and `QDRANT_COLLECTION` are set.
+
+Configurable through environment variables:
+
+- `QDRANT_API_KEY`: The API key to use for authentication. It's not required if the authentication is disabled on the server.

src/lib/server/qdrant/index.ts

@@ -0,0 +1,149 @@
+import { env } from "$env/dynamic/private";
+import type { VecStoreItem, VecStore, VecFindOption } from "$lib/types";
+import { error } from "@sveltejs/kit";
+import { DEFAULT_K, DEFAULT_THRESHOLD } from "../constants";
+import { hash } from "../hash";
+
+/**
+ * The URL of the Qdrant server.
+ * @example "https://xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx.us-east-1-0.aws.cloud.qdrant.io:6333"
+ * @example "http://localhost:6333"
+ */
+const URL = env.QDRANT_SERVER;
+
+/**
+ * The name of the collection.
+ * @example "test"
+ */
+const COLLECTION = env.QDRANT_COLLECTION;
+
+/**
+ * The (optional) API key.
+ */
+const API_KEY = env.QDRANT_API_KEY;
+
+export class QdrantVecStore implements VecStore {
+	async put(item: VecStoreItem): Promise<void> {
+		const { url, collection } = this.env();
+
+		const res = await fetch(`${url}/collections/${collection}/points`, {
+			method: "PUT",
+			headers: this.headers(),
+			body: JSON.stringify({
+				points: [
+					{
+						id: await this.id(item.id),
+						vector: item.v,
+						payload: { ...item.meta, item_id: item.id },
+					},
+				],
+			}),
+		});
+
+		if (!res.ok) {
+			const body = await res.text().catch(() => "");
+			console.log("qdrant put error", body);
+			throw error(500, `Qdrant: ${res.statusText} ${body}`);
+		}
+	}
+
+	async has(id: string): Promise<boolean> {
+		const { url, collection } = this.env();
+
+		const res = await fetch(`${url}/collections/${collection}/points/${await this.id(id)}`, {
+			headers: this.headers(),
+		});
+
+		return res.ok;
+	}
+
+	async find(
+		item: VecStoreItem,
+		option: VecFindOption,
+	): Promise<{ id: string; score: number }[]> {
+		const { url, collection } = this.env();
+
+		const res = await fetch(`${url}/collections/${collection}/points/search`, {
+			method: "POST",
+			headers: this.headers(),
+			body: JSON.stringify({
+				vector: item.v,
+				limit: option.k ?? DEFAULT_K,
+				score_threshold: option.threshold ?? DEFAULT_THRESHOLD,
+				filter: {
+					must: Object.entries(item.meta).map(([key, value]) => ({
+						key,
+						match: { value },
+					})),
+				},
+				with_payload: ["item_id"],
+			}),
+		});
+
+		if (!res.ok) {
+			throw error(500, `Qdrant: ${res.statusText} ${await res.text().catch(() => "")}`);
+		}
+
+		const result = await res.json<{
+			time: number;
+			status: string;
+			result: [
+				{
+					id: number;
+					version: number;
+					score: number;
+					payload: { item_id: string };
+					vector: string;
+				},
+			];
+		}>();
+		console.log("qdrant result", result);
+
+		return result.result.map((r) => ({
+			id: r.payload.item_id,
+			score: r.score,
+		}));
+	}
+
+	async del(id: string): Promise<void> {
+		const { url, collection } = this.env();
+
+		const res = await fetch(`${url}/collections/${collection}/points/delete`, {
+			method: "POST",
+			headers: this.headers(),
+			body: JSON.stringify({
+				points: [await this.id(id)],
+			}),
+		});
+
+		if (!res.ok) {
+			throw error(500, `Qdrant: ${res.statusText} ${await res.text().catch(() => "")}`);
+		}
+	}
+
+	env(): { url: string; collection: string } {
+		if (!URL || !COLLECTION) {
+			throw error(500, "Qdrant URL or collection is not set");
+		}
+
+		return { url: URL, collection: COLLECTION };
+	}
+
+	headers(): Record<string, string> {
+		const h: Record<string, string> = {
+			accept: "application/json",
+			"content-type": "application/json",
+		};
+
+		if (API_KEY) {
+			h["api-key"] = API_KEY;
+		}
+
+		return h;
+	}
+
+	async id(raw: string): Promise<string> {
+		const h = await hash(raw);
+		return h.slice(0, 32);
+	}
+}