feat: Qdrant support for RAG (#40)

* feat: Qdrant RAG Signed-off-by: Anush008 <[email protected]> * chore: Misc improvements Signed-off-by: Anush008 <[email protected]> --------- Signed-off-by: Anush008 <[email protected]>
vespperhq · Nov 8, 2024 · 1bba457 · 1bba457
1 parent 6907bc2
commit 1bba457
Show file tree

Hide file tree

Showing 11 changed files with 749 additions and 1,020 deletions.
diff --git a/packages/db/src/schemas/db-index.ts b/packages/db/src/schemas/db-index.ts
@@ -12,7 +12,7 @@ export const IndexSchema = new Schema<IIndex>(
     },
     type: {
       type: String,
-      enum: ["pinecone", "chromadb"],
+      enum: ["pinecone", "chromadb", "qdrant"],
       default: "pinecone",
     },
     dataSources: [

diff --git a/packages/db/src/types.ts b/packages/db/src/types.ts
@@ -104,7 +104,7 @@ export interface IVendor {
 export interface IIndex {
   _id: Types.ObjectId;
   name: string;
-  type: "pinecone" | "chromadb";
+  type: "pinecone" | "chromadb" | "qdrant";
   organization: Types.ObjectId | IOrganization;
   dataSources: VendorName[];
   state: {

diff --git a/packages/py-db/src/db/db_types.py b/packages/py-db/src/db/db_types.py
@@ -38,7 +38,7 @@ class Index(CommonDBModel):
     organization: Union[PyObjectId, Organization]
     dataSources: List[str]
     name: str
-    type: Literal["chromadb", "pinecone"]
+    type: Literal["chromadb", "pinecone", "qdrant"]
     stats: Optional[dict] = None
     state: IndexState
 

diff --git a/services/api/package.json b/services/api/package.json
@@ -4,15 +4,16 @@
   "private": true,
   "dependencies": {
     "@datadog/datadog-api-client": "^1.21.0",
-    "@vespper/db": "*",
-    "@vespper/utils": "*",
     "@octokit/types": "^12.4.0",
     "@ory/client": "^1.9.0",
     "@ory/kratos-client": "^1.2.0",
     "@pinecone-database/pinecone": "^2.0.1",
+    "@qdrant/js-client-rest": "^1.12.0",
     "@slack/bolt": "^3.16.0",
     "@types/cors": "^2.8.17",
     "@types/ioredis": "^5.0.0",
+    "@vespper/db": "*",
+    "@vespper/utils": "*",
     "axios": "^1.6.2",
     "body-parser": "^1.19.0",
     "chromadb": "1.9.1",

diff --git a/services/api/src/agent/rag/qdrant.ts b/services/api/src/agent/rag/qdrant.ts
@@ -0,0 +1,103 @@
+import { VectorStore, Document, QueryOptions } from "./types";
+import { QdrantClient, Schemas } from "@qdrant/js-client-rest";
+import { embedModel } from "../model";
+import { TextNode } from "llamaindex";
+
+export class QdrantVectorStore implements VectorStore {
+  private readonly client: QdrantClient;
+  private readonly collectionName: string;
+  private readonly vectorName?: string;
+
+  constructor(
+    collectionName: string,
+    url: string,
+    apiKey?: string,
+    vectorName?: string,
+  ) {
+    this.client = new QdrantClient({
+      url: url,
+      apiKey: apiKey,
+    });
+    this.collectionName = collectionName;
+    this.vectorName = vectorName;
+  }
+
+  async query({
+    query,
+    topK = 5,
+    metadata = {},
+  }: QueryOptions): Promise<Document[]> {
+    const vector = await embedModel.getTextEmbedding(query);
+
+    const points = (
+      await this.client.query(this.collectionName, {
+        query: vector,
+        filter: metadata,
+        limit: topK,
+        with_vector: true,
+        with_payload: true,
+      })
+    ).points;
+
+    return pointsToDocuments(points, this.vectorName);
+  }
+
+  async deleteIndex(): Promise<void> {
+    await this.client.deleteCollection(this.collectionName);
+  }
+}
+
+function pointsToDocuments(
+  points: Schemas["ScoredPoint"][],
+  vectorName?: string,
+) {
+  const documents = [];
+  for (const point of points) {
+    const { metadata = {}, ...data } = parseMetadata(point.payload);
+    let vector = null;
+
+    if (
+      point.vector &&
+      typeof point.vector === "object" &&
+      !Array.isArray(point.vector)
+    ) {
+      // Qdrant supports multiple vectors per entry
+      // https://qdrant.tech/documentation/concepts/vectors/#named-vectors
+      // @ts-expect-error We ignore other formats like sparse vectors
+      vector = point.vector[vectorName];
+    } else if (Array.isArray(point.vector)) {
+      vector = point.vector as number[];
+    }
+
+    const document: Document = {
+      id: point.id.toString(),
+      embedding: vector,
+      score: point.score,
+      text: data.text || "",
+      metadata,
+    };
+
+    documents.push(document);
+  }
+
+  return documents;
+}
+
+function parseMetadata(
+  payload?: Record<string, unknown> | null,
+): Partial<TextNode> {
+  if (!payload) {
+    throw new Error("payload is undefined.");
+  }
+
+  const nodeContent = payload["_node_content"];
+  if (!nodeContent) {
+    throw new Error("_node_content key is not found in the payload");
+  }
+
+  if (typeof nodeContent !== "string") {
+    throw new Error("_node_content value is not a string.");
+  }
+
+  return JSON.parse(nodeContent);
+}
diff --git a/services/api/src/agent/rag/utils.ts b/services/api/src/agent/rag/utils.ts
@@ -1,10 +1,11 @@
 import { PineconeVectorStore } from "./pinecone";
 import { ChromaDBVectorStore } from "./chromadb";
+import { QdrantVectorStore } from "./qdrant";
 import type { Document } from "./types";
 
 export function getVectorStore(
   indexName: string,
-  indexType: "pinecone" | "chromadb",
+  indexType: "pinecone" | "chromadb" | "qdrant",
 ) {
   switch (indexType) {
     case "pinecone":
@@ -23,6 +24,16 @@ export function getVectorStore(
         process.env.CHROMA_API_KEY as string,
         indexName,
       );
+    case "qdrant":
+      if (!process.env.QDRANT_URL) {
+        throw new Error("QDRANT_URL is required to use Qdrant");
+      }
+      return new QdrantVectorStore(
+        indexName,
+        process.env.QDRANT_URL as string,
+        process.env.QDRANT_API_KEY as string,
+        process.env.QDRANT_VECTOR_NAME as string,
+      );
     default:
       throw new Error(`Invalid index source: ${indexType}`);
   }