totem-llm 0.1.0

package/server/utils/vectorDbProviders/chroma/index.js

@@ -0,0 +1,484 @@
+const { ChromaClient } = require("chromadb");
+const { TextSplitter } = require("../../TextSplitter");
+const { SystemSettings } = require("../../../models/systemSettings");
+const { storeVectorResult, cachedVectorInformation } = require("../../files");
+const { v4: uuidv4 } = require("uuid");
+const { toChunks, getEmbeddingEngineSelection } = require("../../helpers");
+const { parseAuthHeader } = require("../../http");
+const { sourceIdentifier } = require("../../chats");
+const { VectorDatabase } = require("../base");
+const COLLECTION_REGEX = new RegExp(
+  /^(?!\d+\.\d+\.\d+\.\d+$)(?!.*\.\.)(?=^[a-zA-Z0-9][a-zA-Z0-9_-]{1,61}[a-zA-Z0-9]$).{3,63}$/
+);
+
+class Chroma extends VectorDatabase {
+  constructor() {
+    super();
+  }
+
+  get name() {
+    return "Chroma";
+  }
+
+  // Chroma DB has specific requirements for collection names:
+  // (1) Must contain 3-63 characters
+  // (2) Must start and end with an alphanumeric character
+  // (3) Can only contain alphanumeric characters, underscores, or hyphens
+  // (4) Cannot contain two consecutive periods (..)
+  // (5) Cannot be a valid IPv4 address
+  // We need to enforce these rules by normalizing the collection names
+  // before communicating with the Chroma DB.
+  normalize(inputString) {
+    if (COLLECTION_REGEX.test(inputString)) return inputString;
+    let normalized = inputString.replace(/[^a-zA-Z0-9_-]/g, "-");
+
+    // Replace consecutive periods with a single period (if any)
+    normalized = normalized.replace(/\.\.+/g, ".");
+
+    // Ensure the name doesn't start with a non-alphanumeric character
+    if (normalized[0] && !/^[a-zA-Z0-9]$/.test(normalized[0])) {
+      normalized = "anythingllm-" + normalized.slice(1);
+    }
+
+    // Ensure the name doesn't end with a non-alphanumeric character
+    if (
+      normalized[normalized.length - 1] &&
+      !/^[a-zA-Z0-9]$/.test(normalized[normalized.length - 1])
+    ) {
+      normalized = normalized.slice(0, -1);
+    }
+
+    // Ensure the length is between 3 and 63 characters
+    if (normalized.length < 3) {
+      normalized = `anythingllm-${normalized}`;
+    } else if (normalized.length > 63) {
+      // Recheck the norm'd name if sliced since its ending can still be invalid.
+      normalized = this.normalize(normalized.slice(0, 63));
+    }
+
+    // Ensure the name is not an IPv4 address
+    if (/^\d+\.\d+\.\d+\.\d+$/.test(normalized)) {
+      normalized = "-" + normalized.slice(1);
+    }
+
+    return normalized;
+  }
+
+  async connect() {
+    if (process.env.VECTOR_DB !== "chroma")
+      throw new Error("Chroma::Invalid ENV settings");
+
+    const client = new ChromaClient({
+      path: process.env.CHROMA_ENDPOINT, // if not set will fallback to localhost:8000
+      ...(!!process.env.CHROMA_API_HEADER && !!process.env.CHROMA_API_KEY
+        ? {
+            fetchOptions: {
+              headers: parseAuthHeader(
+                process.env.CHROMA_API_HEADER || "X-Api-Key",
+                process.env.CHROMA_API_KEY
+              ),
+            },
+          }
+        : {}),
+    });
+
+    const isAlive = await client.heartbeat();
+    if (!isAlive)
+      throw new Error(
+        "ChromaDB::Invalid Heartbeat received - is the instance online?"
+      );
+    return { client };
+  }
+
+  async heartbeat() {
+    const { client } = await this.connect();
+    return { heartbeat: await client.heartbeat() };
+  }
+
+  async totalVectors() {
+    const { client } = await this.connect();
+    const collections = await client.listCollections();
+    var totalVectors = 0;
+    for (const collectionObj of collections) {
+      const collection = await client
+        .getCollection({ name: collectionObj.name })
+        .catch(() => null);
+      if (!collection) continue;
+      totalVectors += await collection.count();
+    }
+    return totalVectors;
+  }
+
+  distanceToSimilarity(distance = null) {
+    if (distance === null || typeof distance !== "number") return 0.0;
+    if (distance >= 1.0) return 1;
+    if (distance < 0) return 1 - Math.abs(distance);
+    return 1 - distance;
+  }
+
+  async namespaceCount(_namespace = null) {
+    const { client } = await this.connect();
+    const namespace = await this.namespace(client, this.normalize(_namespace));
+    return namespace?.vectorCount || 0;
+  }
+
+  async similarityResponse({
+    client,
+    namespace,
+    queryVector,
+    similarityThreshold = 0.25,
+    topN = 4,
+    filterIdentifiers = [],
+  }) {
+    const collection = await client.getCollection({
+      name: this.normalize(namespace),
+    });
+    const result = {
+      contextTexts: [],
+      sourceDocuments: [],
+      scores: [],
+    };
+
+    const response = await collection.query({
+      queryEmbeddings: queryVector,
+      nResults: topN,
+    });
+
+    response.ids[0].forEach((_, i) => {
+      const similarity = this.distanceToSimilarity(response.distances[0][i]);
+      if (similarity < similarityThreshold) return;
+
+      if (
+        filterIdentifiers.includes(sourceIdentifier(response.metadatas[0][i]))
+      ) {
+        this.logger(
+          "A source was filtered from context as it's parent document is pinned."
+        );
+        return;
+      }
+
+      result.contextTexts.push(response.documents[0][i]);
+      result.sourceDocuments.push(response.metadatas[0][i]);
+      result.scores.push(similarity);
+    });
+
+    return result;
+  }
+
+  async namespace(client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const collection = await client
+      .getCollection({ name: this.normalize(namespace) })
+      .catch(() => null);
+    if (!collection) return null;
+
+    return {
+      ...collection,
+      vectorCount: await collection.count(),
+    };
+  }
+
+  async hasNamespace(namespace = null) {
+    if (!namespace) return false;
+    const { client } = await this.connect();
+    return await this.namespaceExists(client, this.normalize(namespace));
+  }
+
+  async namespaceExists(client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const collection = await client
+      .getCollection({ name: this.normalize(namespace) })
+      .catch((e) => {
+        this.logger("namespaceExists", e.message);
+        return null;
+      });
+    return !!collection;
+  }
+
+  async deleteVectorsInNamespace(client, namespace = null) {
+    await client.deleteCollection({ name: this.normalize(namespace) });
+    return true;
+  }
+
+  async addDocumentToNamespace(
+    namespace,
+    documentData = {},
+    fullFilePath = null,
+    skipCache = false
+  ) {
+    const { DocumentVectors } = require("../../../models/vectors");
+    try {
+      const { pageContent, docId, ...metadata } = documentData;
+      if (!pageContent || pageContent.length == 0) return false;
+
+      this.logger("Adding new vectorized document into namespace", namespace);
+      if (!skipCache) {
+        const cacheResult = await cachedVectorInformation(fullFilePath);
+        if (cacheResult.exists) {
+          const { client } = await this.connect();
+          const collection = await client.getOrCreateCollection({
+            name: this.normalize(namespace),
+            // returns [-1, 1] unit vector
+            metadata: { "hnsw:space": "cosine" },
+          });
+          const { chunks } = cacheResult;
+          const documentVectors = [];
+
+          for (const chunk of chunks) {
+            const submission = {
+              ids: [],
+              embeddings: [],
+              metadatas: [],
+              documents: [],
+            };
+
+            // Before sending to Chroma and saving the records to our db
+            // we need to assign the id of each chunk that is stored in the cached file.
+            chunk.forEach((chunk) => {
+              const id = uuidv4();
+              const { id: _id, ...metadata } = chunk.metadata;
+              documentVectors.push({ docId, vectorId: id });
+              submission.ids.push(id);
+              submission.embeddings.push(chunk.values);
+              submission.metadatas.push(metadata);
+              submission.documents.push(metadata.text);
+            });
+
+            await this.smartAdd(collection, submission);
+          }
+
+          await DocumentVectors.bulkInsert(documentVectors);
+          return { vectorized: true, error: null };
+        }
+      }
+
+      // If we are here then we are going to embed and store a novel document.
+      // We have to do this manually as opposed to using LangChains `Chroma.fromDocuments`
+      // because we then cannot atomically control our namespace to granularly find/remove documents
+      // from vectordb.
+      const EmbedderEngine = getEmbeddingEngineSelection();
+      const textSplitter = new TextSplitter({
+        chunkSize: TextSplitter.determineMaxChunkSize(
+          await SystemSettings.getValueOrFallback({
+            label: "text_splitter_chunk_size",
+          }),
+          EmbedderEngine?.embeddingMaxChunkLength
+        ),
+        chunkOverlap: await SystemSettings.getValueOrFallback(
+          { label: "text_splitter_chunk_overlap" },
+          20
+        ),
+        chunkHeaderMeta: TextSplitter.buildHeaderMeta(metadata),
+        chunkPrefix: EmbedderEngine?.embeddingPrefix,
+      });
+      const textChunks = await textSplitter.splitText(pageContent);
+
+      this.logger("Snippets created from document:", textChunks.length);
+      const documentVectors = [];
+      const vectors = [];
+      const vectorValues = await EmbedderEngine.embedChunks(textChunks);
+      const submission = {
+        ids: [],
+        embeddings: [],
+        metadatas: [],
+        documents: [],
+      };
+
+      if (!!vectorValues && vectorValues.length > 0) {
+        for (const [i, vector] of vectorValues.entries()) {
+          const vectorRecord = {
+            id: uuidv4(),
+            values: vector,
+            // [DO NOT REMOVE]
+            // LangChain will be unable to find your text if you embed manually and dont include the `text` key.
+            // https://github.com/hwchase17/langchainjs/blob/2def486af734c0ca87285a48f1a04c057ab74bdf/langchain/src/vectorstores/pinecone.ts#L64
+            metadata: { ...metadata, text: textChunks[i] },
+          };
+
+          submission.ids.push(vectorRecord.id);
+          submission.embeddings.push(vectorRecord.values);
+          submission.metadatas.push(metadata);
+          submission.documents.push(textChunks[i]);
+
+          vectors.push(vectorRecord);
+          documentVectors.push({ docId, vectorId: vectorRecord.id });
+        }
+      } else {
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
+        );
+      }
+
+      const { client } = await this.connect();
+      const collection = await client.getOrCreateCollection({
+        name: this.normalize(namespace),
+        metadata: { "hnsw:space": "cosine" },
+      });
+
+      if (vectors.length > 0) {
+        const chunks = [];
+        this.logger("Inserting vectorized chunks into Chroma collection.");
+        for (const chunk of toChunks(vectors, 500)) chunks.push(chunk);
+
+        try {
+          await this.smartAdd(collection, submission);
+          this.logger(
+            `Successfully added ${submission.ids.length} vectors to collection ${this.normalize(namespace)}`
+          );
+        } catch (error) {
+          this.logger("Error adding to ChromaDB:", error);
+          throw new Error(`Error embedding into ChromaDB: ${error.message}`);
+        }
+
+        await storeVectorResult(chunks, fullFilePath);
+      }
+
+      await DocumentVectors.bulkInsert(documentVectors);
+      return { vectorized: true, error: null };
+    } catch (e) {
+      this.logger("addDocumentToNamespace", e.message);
+      return { vectorized: false, error: e.message };
+    }
+  }
+
+  async deleteDocumentFromNamespace(namespace, docId) {
+    const { DocumentVectors } = require("../../../models/vectors");
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, namespace))) return;
+    const collection = await client.getCollection({
+      name: this.normalize(namespace),
+    });
+
+    const knownDocuments = await DocumentVectors.where({ docId });
+    if (knownDocuments.length === 0) return;
+
+    const vectorIds = knownDocuments.map((doc) => doc.vectorId);
+    await this.smartDelete(collection, vectorIds);
+
+    const indexes = knownDocuments.map((doc) => doc.id);
+    await DocumentVectors.deleteIds(indexes);
+    return true;
+  }
+
+  async performSimilaritySearch({
+    namespace = null,
+    input = "",
+    LLMConnector = null,
+    similarityThreshold = 0.25,
+    topN = 4,
+    filterIdentifiers = [],
+  }) {
+    if (!namespace || !input || !LLMConnector)
+      throw new Error("Invalid request to performSimilaritySearch.");
+
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, this.normalize(namespace)))) {
+      return {
+        contextTexts: [],
+        sources: [],
+        message: "Invalid query - no documents found for workspace!",
+      };
+    }
+
+    const queryVector = await LLMConnector.embedTextInput(input);
+    const { contextTexts, sourceDocuments, scores } =
+      await this.similarityResponse({
+        client,
+        namespace,
+        queryVector,
+        similarityThreshold,
+        topN,
+        filterIdentifiers,
+      });
+
+    const sources = sourceDocuments.map((metadata, i) => ({
+      metadata: {
+        ...metadata,
+        text: contextTexts[i],
+        score: scores?.[i] || null,
+      },
+    }));
+
+    return {
+      contextTexts,
+      sources: this.curateSources(sources),
+      message: false,
+    };
+  }
+
+  async "namespace-stats"(reqBody = {}) {
+    const { namespace = null } = reqBody;
+    if (!namespace) throw new Error("namespace required");
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, this.normalize(namespace))))
+      throw new Error("Namespace by that name does not exist.");
+    const stats = await this.namespace(client, this.normalize(namespace));
+    return stats
+      ? stats
+      : { message: "No stats were able to be fetched from DB for namespace" };
+  }
+
+  async "delete-namespace"(reqBody = {}) {
+    const { namespace = null } = reqBody;
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, this.normalize(namespace))))
+      throw new Error("Namespace by that name does not exist.");
+
+    const details = await this.namespace(client, this.normalize(namespace));
+    await this.deleteVectorsInNamespace(client, this.normalize(namespace));
+    return {
+      message: `Namespace ${namespace} was deleted along with ${details?.vectorCount} vectors.`,
+    };
+  }
+
+  async reset() {
+    const { client } = await this.connect();
+    await client.reset();
+    return { reset: true };
+  }
+
+  curateSources(sources = []) {
+    const documents = [];
+    for (const source of sources) {
+      const { metadata = {} } = source;
+      if (Object.keys(metadata).length > 0) {
+        documents.push({
+          ...metadata,
+          ...(source.hasOwnProperty("pageContent")
+            ? { text: source.pageContent }
+            : {}),
+        });
+      }
+    }
+
+    return documents;
+  }
+
+  /**
+   * This method is a wrapper around the ChromaCollection.add method.
+   * It will return true if the add was successful, false otherwise.
+   * For local deployments, this will be the same as calling the add method directly since there are no limitations.
+   * @param {import("chromadb").Collection} collection
+   * @param {{ids: string[], embeddings: number[], metadatas: Record<string, any>[], documents: string[]}[]} submissions
+   * @returns {Promise<boolean>} True if the add was successful, false otherwise.
+   */
+  async smartAdd(collection, submissions) {
+    await collection.add(submissions);
+    return true;
+  }
+
+  /**
+   * This method is a wrapper around the ChromaCollection.delete method.
+   * It will return the result of the delete method directly.
+   * For local deployments, this will be the same as calling the delete method directly since there are no limitations.
+   * @param {import("chromadb").Collection} collection
+   * @param {string[]} vectorIds
+   * @returns {Promise<boolean>} True if the delete was successful, false otherwise.
+   */
+  async smartDelete(collection, vectorIds) {
+    await collection.delete({ ids: vectorIds });
+    return true;
+  }
+}
+
+module.exports.Chroma = Chroma;

package/server/utils/vectorDbProviders/chromacloud/index.js

@@ -0,0 +1,158 @@
+const { CloudClient } = require("chromadb");
+const { Chroma } = require("../chroma");
+const { toChunks } = require("../../helpers");
+
+/**
+ * ChromaCloud works nearly the same as Chroma so we can just extend the
+ * Chroma class and override the connect method to use the CloudClient for major differences in API functionality.
+ */
+class ChromaCloud extends Chroma {
+  constructor() {
+    super();
+  }
+
+  get name() {
+    return "ChromaCloud";
+  }
+
+  /**
+   * Basic quota/limitations for Chroma Cloud for accounts. Does not lookup client-specific limits.
+   * @see https://docs.trychroma.com/cloud/quotas-limits
+   */
+  limits = {
+    maxEmbeddingDim: 4_096,
+    maxDocumentBytes: 16_384,
+    maxMetadataBytes: 4_096,
+    maxRecordsPerWrite: 300,
+  };
+
+  async connect() {
+    if (process.env.VECTOR_DB !== "chromacloud")
+      throw new Error("ChromaCloud::Invalid ENV settings");
+
+    const client = new CloudClient({
+      apiKey: process.env.CHROMACLOUD_API_KEY,
+      tenant: process.env.CHROMACLOUD_TENANT,
+      database: process.env.CHROMACLOUD_DATABASE,
+    });
+
+    const isAlive = await client.heartbeat();
+    if (!isAlive)
+      throw new Error(
+        "ChromaCloud::Invalid Heartbeat received - is the instance online?"
+      );
+    return { client };
+  }
+
+  /**
+   * Chroma Cloud has some basic limitations on upserts to protect performance and latency.
+   * Local deployments do not have these limitations since they are self-hosted.
+   *
+   * This method, if cloud, will do some simple logic/heuristics to ensure that the upserts are not too large.
+   * Otherwise, it may throw a 422.
+   * @param {import("chromadb").Collection} collection
+   * @param {{ids: string[], embeddings: number[], metadatas: Record<string, any>[], documents: string[]}[]} submissions
+   * @returns {Promise<boolean>} True if the upsert was successful, false otherwise.
+   * If the upsert was not successful, the error message will be returned.
+   */
+  async smartAdd(collection, submission) {
+    const testSubmission = {
+      id: submission.ids[0],
+      embedding: submission.embeddings[0],
+      metadata: submission.metadatas[0],
+      document: submission.documents[0],
+    };
+
+    if (testSubmission.embedding.length > this.limits.maxEmbeddingDim)
+      console.warn(
+        `ChromaCloud::Embedding dimension too large (default max is ${this.limits.maxEmbeddingDim}). Got ${testSubmission.embedding.length}. Upsert may fail!`
+      );
+    if (testSubmission.document.length > this.limits.maxDocumentBytes)
+      console.warn(
+        `ChromaCloud::Document length too large (default max is ${this.limits.maxDocumentBytes}). Got ${testSubmission.document.length}. Upsert may fail!`
+      );
+    if (
+      JSON.stringify(testSubmission.metadata).length >
+      this.limits.maxMetadataBytes
+    )
+      console.warn(
+        `ChromaCloud::Metadata length too large (default max is ${this.limits.maxMetadataBytes}). Got ${JSON.stringify(testSubmission.metadata).length}. Upsert may fail!`
+      );
+
+    // If the submissions are not too large, just add them directly.
+    if (submission.ids.length <= this.limits.maxRecordsPerWrite) {
+      await collection.add(submission);
+      return true;
+    }
+
+    this.logger(
+      `Upsert Payload is too large (max is ${this.limits.maxRecordsPerWrite} records). Splitting into chunks of ${this.limits.maxRecordsPerWrite} records.`
+    );
+    const chunks = [];
+    let chunkedSubmission = {
+      ids: [],
+      embeddings: [],
+      metadatas: [],
+      documents: [],
+    };
+    for (let i = 0; i < submission.ids.length; i++) {
+      chunkedSubmission.ids.push(submission.ids[i]);
+      chunkedSubmission.embeddings.push(submission.embeddings[i]);
+      chunkedSubmission.metadatas.push(submission.metadatas[i]);
+      chunkedSubmission.documents.push(submission.documents[i]);
+      if (chunkedSubmission.ids.length === this.limits.maxRecordsPerWrite) {
+        this.logger(
+          `ChromaCloud::Adding chunk payload ${chunks.length + 1} of ${Math.ceil(submission.ids.length / this.limits.maxRecordsPerWrite)}`
+        );
+        chunks.push(chunkedSubmission);
+        chunkedSubmission = {
+          ids: [],
+          embeddings: [],
+          metadatas: [],
+          documents: [],
+        };
+      }
+    }
+    // Push remaining submissions to the last chunk
+    if (chunkedSubmission.ids.length > 0) chunks.push(chunkedSubmission);
+
+    let counter = 1;
+    for (const chunk of chunks) {
+      await collection.add(chunk);
+      //eslint-disable-next-line
+      counter++;
+    }
+    return true;
+  }
+
+  /**
+   * This method is a wrapper around the ChromaCollection.delete method.
+   * It will return the result of the delete method directly.
+   * Chroma Cloud has some basic limitations on deletes to protect performance and latency.
+   * Local deployments do not have these limitations since they are self-hosted.
+   *
+   * This method, if cloud, will do some simple logic/heuristics to ensure that the deletes are not too large.
+   * Otherwise, it may throw a 422.
+   * @param {import("chromadb").Collection} collection
+   * @param {string[]} vectorIds
+   * @returns {Promise<boolean>} True if the delete was successful, false otherwise.
+   */
+  async smartDelete(collection, vectorIds) {
+    if (vectorIds.length <= this.limits.maxRecordsPerWrite)
+      return await collection.delete({ ids: vectorIds });
+
+    this.logger(
+      `Delete Payload is too large (max is ${this.limits.maxRecordsPerWrite} records). Splitting into chunks of ${this.limits.maxRecordsPerWrite} records.`
+    );
+    const chunks = toChunks(vectorIds, this.limits.maxRecordsPerWrite);
+    let counter = 1;
+    for (const chunk of chunks) {
+      this.logger(`Deleting chunk ${counter} of ${chunks.length}`);
+      await collection.delete({ ids: chunk });
+      counter++;
+    }
+    return true;
+  }
+}
+
+module.exports.ChromaCloud = ChromaCloud;