Merge main into sweep/sweep_gha_fix_the_github_actions_run_fai_1

examples/chromadb/test.ts

@@ -0,0 +1,36 @@
+import {
+  ChromaVectorStore,
+  PapaCSVReader,
+  storageContextFromDefaults,
+  VectorStoreIndex
+} from "llamaindex"
+
+async function main() {
+  try {
+    console.log("Loading data from astradb/data/movie_reviews.csv")
+    const reader = new PapaCSVReader(false, ", ", "\n", {
+  header: true,
+      header: true
+    });
+    const docs = await reader.loadData("../astradb/data/movie_reviews.csv");
+
+    console.log("Creating ChromaDB vector store")
+    const collectionName = "movie_reviews";
+    const chromaVS = new ChromaVectorStore({ collectionName });
+    const ctx = await storageContextFromDefaults({ vectorStore: chromaVS });
+    console.log("Embedding documents and adding to index")
+    const index = await VectorStoreIndex.fromDocuments(docs, {
+  storageContext: ctx,
+      storageContext: ctx,
+    });
+
+    console.log("Querying index")
+    const queryEngine = index.asQueryEngine();
+    const response = await queryEngine.query("Tell me about Godfrey Cheshire's rating of La Sapienza.");
+    console.log(response.toString());
+  } catch (e) {
+    console.error(e);
+  }
+}
+
+main();

examples/package.json

@@ -3,12 +3,13 @@
   "private": true,
   "name": "examples",
   "dependencies": {
-    "@notionhq/client": "^2.2.13",
     "@datastax/astra-db-ts": "^0.1.2",
+    "@notionhq/client": "^2.2.13",
     "@pinecone-database/pinecone": "^1.1.2",
+    "chromadb": "^1.7.3",
     "commander": "^11.1.0",
-    "llamaindex": "latest",
     "dotenv": "^16.3.1",
+    "llamaindex": "latest",
     "mongodb": "^6.2.0"
   },
   "devDependencies": {

packages/core/package.json

@@ -10,6 +10,7 @@
     "@pinecone-database/pinecone": "^1.1.2",
     "@xenova/transformers": "^2.10.0",
     "assemblyai": "^4.0.0",
+    "chromadb": "^1.7.3",
     "compromise": "^14.10.1",
     "file-type": "^18.7.0",
     "js-tiktoken": "^1.0.8",

packages/core/src/storage/index.ts

@@ -8,6 +8,7 @@ export * from "./indexStore/types";
 export { SimpleKVStore } from "./kvStore/SimpleKVStore";
 export * from "./kvStore/types";
 export { AstraDBVectorStore } from "./vectorStore/AstraDBVectorStore";
+export { ChromaVectorStore } from "./vectorStore/ChromaVectorStore";
 export { MongoDBAtlasVectorSearch } from "./vectorStore/MongoDBAtlasVectorStore";
 export { PGVectorStore } from "./vectorStore/PGVectorStore";
 export { PineconeVectorStore } from "./vectorStore/PineconeVectorStore";

packages/core/src/storage/vectorStore/ChromaVectorStore.ts

@@ -0,0 +1,133 @@
+import { VectorStore, VectorStoreQuery, VectorStoreQueryMode, VectorStoreQueryResult } from "./types";
+import { BaseNode, Document, MetadataMode } from "../../Node";
+import { ChromaClient, Collection, IncludeEnum, Where, WhereDocument, Metadata, Embeddings } from 'chromadb';
+
+type ChromaDBParams = {
+  collectionName: string
+}
+
+type ChromaDeleteOptions = {
+  where?: Where;
+  whereDocument?: WhereDocument;
+}
+
+type ChromaQueryOptions = {
+  whereDocument?: WhereDocument
+}
+
+export class ChromaVectorStore implements VectorStore {
+  storesText: boolean = true;
+  private chromaClient: ChromaClient;
+  private collection: Collection | null;
+
+  constructor(params: ChromaDBParams) {
+    this.chromaClient = new ChromaClient();
+    this.collection = null;
+    this.chromaClient.createCollection({
+      name: params.collectionName,
+    }).then((collection) => {
+      this.collection = collection;
+    }).catch((err) => {
+      console.error(err);
+    }
+    );
+  }
+
+  client(): ChromaClient {
+    return this.chromaClient;
+  }
+
+  async add(embeddingResults: BaseNode[]): Promise<string[]> {
+    if (!this.collection) {
+        throw new Error("Must initialize collection before adding.");
+    }
+
+    // Check if all nodes have an embedding and only if they do, add the embeddings to the collection
+    const allNodesHaveEmbeddings = embeddingResults.every((node) => node.embedding !== undefined);
+    const embeddings = allNodesHaveEmbeddings ? embeddingResults.map((node) => node.getEmbedding()): undefined;
+    try {
+      await this.collection.add({
+        ids: embeddingResults.map((node) => node.id_),
+        metadatas: embeddingResults.map((node) => node.metadata),
+        documents: embeddingResults.map((node) => node.getContent(MetadataMode.NONE)),
+        embeddings: embeddings
+      });
+    }
+    catch (err) {
+      const msg = `${err}`;
+      console.log(msg, err);
+      throw err;
+    }
+    return embeddingResults.map((node) => node.id_)
+  }
+  
+  async delete(refDocId: string, deleteOptions?: ChromaDeleteOptions): Promise<void> {
+    if (!this.collection) {
+      throw new Error('Must initialize collection before deleting');
+    }
+
+    try {
+      await this.collection.delete({
+        ids: [refDocId],
+        where: deleteOptions?.where,
+        whereDocument: deleteOptions?.whereDocument,
+      });
+    } catch (err) {
+      const msg = `${err}`;
+      console.log(msg, err);
+      throw err;
+    }
+  }
+
+  async query(query: VectorStoreQuery, options?: ChromaQueryOptions): Promise<VectorStoreQueryResult> {
+    if (!this.collection) {
+      throw new Error("Must connect to collection before querying.");
+    }
+    if (query.docIds) {
+      throw new Error("ChromaDB does not support querying by docIDs");
+    }
+    if (query.mode != VectorStoreQueryMode.DEFAULT) {
+      throw new Error("ChromaDB does not support querying by mode");
+    }
+
+    const chromaWhere: { [x: string]: string | number | boolean; } = {};
+    if (query.filters) {
+      query.filters.filters.map((filter) => {
+        const filterKey = filter.key;
+        const filterValue = filter.value;
+        chromaWhere[filterKey] = filterValue;
+      }
+      );
+    }
+    try {
+      const queryResponse = await this.collection.query({
+        queryEmbeddings: query.queryEmbedding ?? undefined,
+        queryTexts: query.queryStr ?? undefined,
+        nResults: query.similarityTopK,
+        where: Object.keys(chromaWhere).length ? chromaWhere: undefined,
+        whereDocument: options?.whereDocument,
+        //ChromaDB doesn't return the result embeddings by default so we need to include them  
+        include: [IncludeEnum.Distances, IncludeEnum.Metadatas, IncludeEnum.Documents, IncludeEnum.Embeddings]
+      });
+      //Only given 1 embedding/text, so get 0th index for everything
+      const vectorStoreQueryResult: VectorStoreQueryResult = {
+        nodes: queryResponse.ids[0].map((id, index) => {
+          return new Document({
+            id_: id,
+            text: (queryResponse.documents as string[][])[0][index],
+            metadata: queryResponse.metadatas[0][index] ?? {},
+            embedding: (queryResponse.embeddings as Embeddings[])[0][index],
+          });
+        }),
+        similarities: (queryResponse.distances as number[][])[0].map((distance) => 1 - distance),
+        ids: queryResponse.ids[0],
+      };
+      return vectorStoreQueryResult;
+    }
+    catch (err) {
+      const msg = `${err}`;
+      console.log(msg, err);
+      throw err;
+    }
+  }
+}

pnpm-lock.yaml

@@ -115,6 +115,9 @@ importers:
       '@pinecone-database/pinecone':
         specifier: ^1.1.2
         version: 1.1.2
+      chromadb:
+        specifier: ^1.7.3
+        version: 1.7.3(openai@4.20.1)
       commander:
         specifier: ^11.1.0
         version: 11.1.0
@@ -158,6 +161,9 @@ importers:
       assemblyai:
         specifier: ^4.0.0
         version: 4.0.0
+      chromadb:
+        specifier: ^1.7.3
+        version: 1.7.3(openai@4.20.1)
       compromise:
         specifier: ^14.10.1
         version: 14.10.1
@@ -6255,6 +6261,28 @@ packages:
     resolution: {integrity: sha512-bIomtDF5KGpdogkLd9VspvFzk9KfpyyGlS8YFVZl7TGPBHL5snIOnxeshwVgPteQ9b4Eydl+pVbIyE1DcvCWgQ==}
     engines: {node: '>=10'}
 
+  /chromadb@1.7.3(openai@4.20.1):
+    resolution: {integrity: sha512-3GgvQjpqgk5C89x5EuTDaXKbfrdqYDJ5UVyLQ3ZmwxnpetNc+HhRDGjkvXa5KSvpQ3lmKoyDoqnN4tZepfFkbw==}
+    engines: {node: '>=14.17.0'}
+    peerDependencies:
+      '@google/generative-ai': ^0.1.1
+      cohere-ai: ^5.0.0 || ^6.0.0 || ^7.0.0
+      openai: ^3.0.0 || ^4.0.0
+    peerDependenciesMeta:
+      '@google/generative-ai':
+        optional: true
+      cohere-ai:
+        optional: true
+      openai:
+        optional: true
+    dependencies:
+      cliui: 8.0.1
+      isomorphic-fetch: 3.0.0
+      openai: 4.20.1
+    transitivePeerDependencies:
+      - encoding
+    dev: false
+
   /chrome-trace-event@1.0.3:
     resolution: {integrity: sha512-p3KULyQg4S7NIHixdwbGX+nFHkoBiA4YQmyWtjb8XngSKV124nJmRysgAeujbUVb15vh+RvFUfCPqU7rXk+hZg==}
     engines: {node: '>=6.0'}
@@ -6359,7 +6387,6 @@ packages:
       string-width: 4.2.3
       strip-ansi: 6.0.1
       wrap-ansi: 7.0.0
-    dev: true
 
   /clone-deep@4.0.1:
     resolution: {integrity: sha512-neHB9xuzh/wk0dIHweyAXv2aPGZIVk3pLMe+/RNzINf17fe0OG96QroktYAUm7SM1PBnzTabaLboqqxDyMU+SQ==}
@@ -10140,6 +10167,15 @@ packages:
     resolution: {integrity: sha512-WhB9zCku7EGTj/HQQRz5aUQEUeoQZH2bWcltRErOpymJ4boYE6wL9Tbr23krRPSZ+C5zqNSrSw+Cc7sZZ4b7vg==}
     engines: {node: '>=0.10.0'}
 
+  /isomorphic-fetch@3.0.0:
+    resolution: {integrity: sha512-qvUtwJ3j6qwsF3jLxkZ72qCgjMysPzDfeV240JHiGZsANBYd+EEuu35v7dfrJ9Up0Ak07D7GGSkGhCHTqg/5wA==}
+    dependencies:
+      node-fetch: 2.7.0(encoding@0.1.13)
+      whatwg-fetch: 3.6.20
+    transitivePeerDependencies:
+      - encoding
+    dev: false
+
   /isomorphic-timers-promises@1.0.1:
     resolution: {integrity: sha512-u4sej9B1LPSxTGKB/HiuzvEQnXH0ECYkSVQU39koSwmFAxhlEAFl9RdTvLv4TOTQUgBS5O3O5fwUxk6byBZ+IQ==}
     engines: {node: '>=10'}
@@ -16364,6 +16400,10 @@ packages:
     engines: {node: '>=0.8.0'}
     dev: false
 
+  /whatwg-fetch@3.6.20:
+    resolution: {integrity: sha512-EqhiFU6daOA8kpjOWTL0olhVOF3i7OrFzSYiGsEMB8GcXS+RrzauAERX65xMeNWVqxA6HXH2m69Z9LaKKdisfg==}
+    dev: false
+
   /whatwg-url@13.0.0:
     resolution: {integrity: sha512-9WWbymnqj57+XEuqADHrCJ2eSXzn8WXIW/YSGaZtb2WKAInQ6CHfaUUcTyyver0p8BDg5StLQq8h1vtZuwmOig==}
     engines: {node: '>=16'}