add support for multi-vector points

hasura · Jun 8, 2024 · 9a804d3 · 9a804d3
1 parent c07dbcf
commit 9a804d3
Show file tree

Hide file tree

Showing 12 changed files with 376 additions and 74 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,4 +3,5 @@ dist/
 .DS_Store
 configuration.json
 PRIVATE_REMOTE_CONFIGURATION.http
-config.json
+config.json
+TMP.md
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,9 @@
 # Qdrant Connector Changelog
 This changelog documents changes between release tags.
 
+## [0.2.0] - 2024-05-6
+* Added support for multi-vector points
+
 ## [0.1.9] - 2024-05-6
 * Add multi-arch build
 

diff --git a/connector-definition/connector-metadata.yaml b/connector-definition/connector-metadata.yaml
@@ -1,6 +1,6 @@
 packagingDefinition:
   type: PrebuiltDockerImage
-  dockerImage: ghcr.io/hasura/ndc-qdrant:v0.1.9
+  dockerImage: ghcr.io/hasura/ndc-qdrant:v0.2.0
 supportedEnvironmentVariables:
   - name: QDRANT_URL
     description: The url for the Qdrant database
@@ -9,7 +9,7 @@ supportedEnvironmentVariables:
 commands:
   update:
     type: Dockerized
-    dockerImage: ghcr.io/hasura/ndc-qdrant:v0.1.9
+    dockerImage: ghcr.io/hasura/ndc-qdrant:v0.2.0
     commandArgs:
       - update
 dockerComposeWatch:

diff --git a/generate-config.ts b/generate-config.ts
@@ -2,7 +2,7 @@ import { getQdrantClient } from "./src/qdrant";
 import fs from "fs";
 import { promisify } from "util";
 import { insertion } from "./src/utilities";
-import { RESTRICTED_OBJECTS, BASE_FIELDS, BASE_TYPES, INSERT_FIELDS } from "./src/constants";
+import { RESTRICTED_OBJECTS, BASE_FIELDS, BASE_TYPES, INSERT_FIELDS, INSERT_FIELDS_VECTOR } from "./src/constants";
 const readFile = promisify(fs.readFile);
 const writeFile = promisify(fs.writeFile);
 let HASURA_CONFIGURATION_DIRECTORY = process.env["HASURA_CONFIGURATION_DIRECTORY"] as string | undefined;
@@ -21,7 +21,7 @@ let client = getQdrantClient(QDRANT_URL, QDRANT_API_KEY);
 async function main() {
   const collections = await client.getCollections();
   const collectionNames = collections.collections.map((c) => c.name);
-
+  let collectionVectors: Record<string, boolean> = {};
   let objectTypes: Record<string, any> = {
     ...BASE_TYPES,
   };
@@ -32,6 +32,7 @@ async function main() {
     const { points: records } = await client.scroll(cn, {
       limit: 1,
       with_payload: true,
+      with_vector: true
     });
     let fieldDict = {};
     let baseFields = {};
@@ -50,16 +51,29 @@ async function main() {
           },
           ...BASE_FIELDS
         };
-        insertFields = {
-          id: {
-            description: null,
-            type: {
-              type: "named",
-              name: "Int",
+        if (Array.isArray(records[0].vector)){
+          insertFields = {
+            id: {
+              description: null,
+              type: {
+                type: "named",
+                name: "Int",
+              },
             },
-          },
-          ...INSERT_FIELDS
-        };
+            ...INSERT_FIELDS
+          };
+        } else {
+          insertFields = {
+            id: {
+              description: null,
+              type: {
+                type: "named",
+                name: "Int",
+              },
+            },
+            ...INSERT_FIELDS_VECTOR
+          };
+        }
       } else {
         baseFields = {
           id: {
@@ -71,16 +85,36 @@ async function main() {
           },
           ...BASE_FIELDS
         };
-        insertFields = {
-          id: {
-            description: null,
-            type: {
-              type: "named",
-              name: "String",
+
+        if (Array.isArray(records[0].vector)){
+          insertFields = {
+            id: {
+              description: null,
+              type: {
+                type: "named",
+                name: "String",
+              },
             },
-          },
-          ...INSERT_FIELDS
-        };
+            ...INSERT_FIELDS
+          };
+        } else {
+          insertFields = {
+            id: {
+              description: null,
+              type: {
+                type: "named",
+                name: "String",
+              },
+            },
+            ...INSERT_FIELDS_VECTOR
+          };
+        }
+      }
+
+      if (!Array.isArray(records[0].vector)){
+        collectionVectors[cn] = true;
+      } else {
+        collectionVectors[cn] = false;
       }
     }
 
@@ -112,6 +146,7 @@ async function main() {
     collection_names: collectionNames,
     object_fields: objectFields,
     object_types: objectTypes,
+    collection_vectors: collectionVectors,
     functions: [],
     procedures: [],
   }

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "ndc-qdrant",
-  "version": "0.1.9",
+  "version": "0.2.0",
   "main": "index.js",
   "scripts": {
     "start": "ts-node ./src/index.ts serve --configuration=.",

diff --git a/src/constants.ts b/src/constants.ts
@@ -136,6 +136,32 @@ export const INSERT_FIELDS: Record<string, ObjectField> = {
   },
 };
 
+export const INSERT_FIELDS_VECTOR: Record<string, ObjectField> = {
+  vectors: {
+    description: null,
+    type: {
+      type: "array",
+      element_type: {
+        type: "array",
+        element_type: {
+          type: "named",
+          name: "Float",
+        }
+      }
+    },
+  },
+  vector_names: {
+    description: null,
+    type: {
+      type: "array",
+      element_type: {
+        type: "named",
+        name: "String"
+      }
+    },
+  },
+};
+
 export const BASE_FIELDS: Record<string, ObjectField> = {
   score: {
     description: null,
@@ -268,6 +294,44 @@ export const BASE_TYPES: { [k: string]: ObjectType } = {
       },
     },
   },
+  _searchVector: {
+    description: "Search the vector database for similar vectors",
+    fields: {
+      vector: {
+        type: {
+          type: "array",
+          element_type: {
+            type: "named",
+            name: "Float",
+          },
+        },
+      },
+      name: {
+        type: {
+          type: "named",
+          name: "String"
+        }
+      },
+      params: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "_params",
+          },
+        },
+      },
+      score_threshold: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "Float",
+          },
+        },
+      },
+    },
+  },
   _recommendInt: {
     description:
       "Provide an array of positive and negative example points and get a recommendation",
@@ -357,6 +421,108 @@ export const BASE_TYPES: { [k: string]: ObjectType } = {
         },
       },
     },
+  },
+  _recommendIntVector: {
+    description:
+      "Provide an array of positive and negative example points and get a recommendation",
+    fields: {
+      positive: {
+        type: {
+          type: "array",
+          element_type: {
+            type: "named",
+            name: "Int",
+          },
+        },
+      },
+      negative: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "array",
+            element_type: {
+              type: "named",
+              name: "Int",
+            },
+          },
+        },
+      },
+      using: {
+        type: {
+          type: "named",
+          name: "String"
+        }
+      },
+      params: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "_params",
+          },
+        },
+      },
+      score_threshold: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "Float",
+          },
+        },
+      },
+    },
+  },
+  _recommendStringVector: {
+    description:
+      "Provide an array of positive and negative example points and get a recommendation",
+    fields: {
+      positive: {
+        type: {
+          type: "array",
+          element_type: {
+            type: "named",
+            name: "String",
+          },
+        },
+      },
+      negative: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "array",
+            element_type: {
+              type: "named",
+              name: "String",
+            },
+          },
+        },
+      },
+      using: {
+        type: {
+          type: "named",
+          name: "String"
+        }
+      },
+      params: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "_params",
+          },
+        },
+      },
+      score_threshold: {
+        type: {
+          type: "nullable",
+          underlying_type: {
+            type: "named",
+            name: "Float",
+          },
+        },
+      },
+    },
   }
 };
 

diff --git a/src/handlers/explain.ts b/src/handlers/explain.ts
@@ -1,10 +1,10 @@
 import { ExplainResponse, QueryRequest } from "@hasura/ndc-sdk-typescript";
 import { QueryPlan, planQueries } from "./query";
 
-export async function doExplain(query: QueryRequest, collectionNames: string[], collectionFields: {[key: string]: string[]}): Promise<ExplainResponse>{
+export async function doExplain(query: QueryRequest, collectionNames: string[], collectionFields: {[key: string]: string[]}, collectionVectors: {[key: string]: boolean}): Promise<ExplainResponse>{
     let explainResponse: ExplainResponse;
     try {
-        let queryPlan: QueryPlan = await planQueries(query, collectionNames, collectionFields);
+        let queryPlan: QueryPlan = await planQueries(query, collectionNames, collectionFields, collectionVectors);
         let isScroll: boolean = queryPlan.scrollQueries.length > 0;
         explainResponse = {details:{
             queryRequest: JSON.stringify(query),