totem-llm 0.1.0

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

@@ -0,0 +1,434 @@
+const {
+  DataType,
+  MetricType,
+  IndexType,
+  MilvusClient,
+} = require("@zilliz/milvus2-sdk-node");
+const { TextSplitter } = require("../../TextSplitter");
+const { SystemSettings } = require("../../../models/systemSettings");
+const { v4: uuidv4 } = require("uuid");
+const { storeVectorResult, cachedVectorInformation } = require("../../files");
+const { toChunks, getEmbeddingEngineSelection } = require("../../helpers");
+const { sourceIdentifier } = require("../../chats");
+const { VectorDatabase } = require("../base");
+
+class Milvus extends VectorDatabase {
+  constructor() {
+    super();
+  }
+
+  get name() {
+    return "Milvus";
+  }
+
+  // Milvus/Zilliz only allows letters, numbers, and underscores in collection names
+  // so we need to enforce that by re-normalizing the names when communicating with
+  // the DB.
+  // If the first char of the collection is not an underscore or letter the collection name will be invalid.
+  normalize(inputString) {
+    let normalized = inputString.replace(/[^a-zA-Z0-9_]/g, "_");
+    if (new RegExp(/^[a-zA-Z_]/).test(normalized.slice(0, 1)))
+      normalized = `anythingllm_${normalized}`;
+    return normalized;
+  }
+
+  async connect() {
+    if (process.env.VECTOR_DB !== "milvus")
+      throw new Error(`${this.name}::Invalid ENV settings`);
+
+    const client = new MilvusClient({
+      address: process.env.MILVUS_ADDRESS,
+      username: process.env.MILVUS_USERNAME,
+      password: process.env.MILVUS_PASSWORD,
+    });
+
+    const { isHealthy } = await client.checkHealth();
+    if (!isHealthy)
+      throw new Error(
+        `${this.name}::Invalid Heartbeat received - is the instance online?`
+      );
+
+    return { client };
+  }
+
+  async heartbeat() {
+    await this.connect();
+    return { heartbeat: Number(new Date()) };
+  }
+
+  async totalVectors() {
+    const { client } = await this.connect();
+    const { collection_names } = await client.listCollections();
+    const total = collection_names.reduce(async (acc, collection_name) => {
+      const statistics = await client.getCollectionStatistics({
+        collection_name: this.normalize(collection_name),
+      });
+      return Number(acc) + Number(statistics?.data?.row_count ?? 0);
+    }, 0);
+    return total;
+  }
+
+  async namespaceCount(_namespace = null) {
+    const { client } = await this.connect();
+    const statistics = await client.getCollectionStatistics({
+      collection_name: this.normalize(_namespace),
+    });
+    return Number(statistics?.data?.row_count ?? 0);
+  }
+
+  async namespace(client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const collection = await client
+      .getCollectionStatistics({ collection_name: this.normalize(namespace) })
+      .catch(() => null);
+    return collection;
+  }
+
+  async hasNamespace(namespace = null) {
+    if (!namespace) return false;
+    const { client } = await this.connect();
+    return await this.namespaceExists(client, namespace);
+  }
+
+  async namespaceExists(client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const { value } = await client
+      .hasCollection({ collection_name: this.normalize(namespace) })
+      .catch((e) => {
+        console.error(`${this.name}::namespaceExists`, e.message);
+        return { value: false };
+      });
+    return value;
+  }
+
+  async deleteVectorsInNamespace(client, namespace = null) {
+    await client.dropCollection({ collection_name: this.normalize(namespace) });
+    return true;
+  }
+
+  // Milvus requires a dimension aspect for collection creation
+  // we pass this in from the first chunk to infer the dimensions like other
+  // providers do.
+  async getOrCreateCollection(client, namespace, dimensions = null) {
+    const isExists = await this.namespaceExists(client, namespace);
+    if (!isExists) {
+      if (!dimensions)
+        throw new Error(
+          `${this.name}::getOrCreateCollection Unable to infer vector dimension from input. Open an issue on GitHub for support.`
+        );
+
+      await client.createCollection({
+        collection_name: this.normalize(namespace),
+        fields: [
+          {
+            name: "id",
+            description: "id",
+            data_type: DataType.VarChar,
+            max_length: 255,
+            is_primary_key: true,
+          },
+          {
+            name: "vector",
+            description: "vector",
+            data_type: DataType.FloatVector,
+            dim: dimensions,
+          },
+          {
+            name: "metadata",
+            description: "metadata",
+            data_type: DataType.JSON,
+          },
+        ],
+      });
+      await client.createIndex({
+        collection_name: this.normalize(namespace),
+        field_name: "vector",
+        index_type: IndexType.AUTOINDEX,
+        metric_type: MetricType.COSINE,
+      });
+      await client.loadCollectionSync({
+        collection_name: this.normalize(namespace),
+      });
+    }
+  }
+
+  async addDocumentToNamespace(
+    namespace,
+    documentData = {},
+    fullFilePath = null,
+    skipCache = false
+  ) {
+    const { DocumentVectors } = require("../../../models/vectors");
+    try {
+      let vectorDimension = null;
+      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 { chunks } = cacheResult;
+          const documentVectors = [];
+          vectorDimension = chunks[0][0].values.length || null;
+
+          await this.getOrCreateCollection(client, namespace, vectorDimension);
+          try {
+            for (const chunk of chunks) {
+              // Before sending to Milvus and saving the records to our db
+              // we need to assign the id of each chunk that is stored in the cached file.
+              const newChunks = chunk.map((chunk) => {
+                const id = uuidv4();
+                documentVectors.push({ docId, vectorId: id });
+                return { id, vector: chunk.values, metadata: chunk.metadata };
+              });
+              const insertResult = await client.insert({
+                collection_name: this.normalize(namespace),
+                data: newChunks,
+              });
+
+              if (insertResult?.status.error_code !== "Success") {
+                throw new Error(
+                  `Error embedding into ${this.name}! Reason:${insertResult?.status.reason}`
+                );
+              }
+            }
+            await DocumentVectors.bulkInsert(documentVectors);
+            await client.flushSync({
+              collection_names: [this.normalize(namespace)],
+            });
+            return { vectorized: true, error: null };
+          } catch (insertError) {
+            console.error(
+              "Error inserting cached chunks:",
+              insertError.message
+            );
+            return { vectorized: false, error: insertError.message };
+          }
+        }
+      }
+
+      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);
+
+      if (!!vectorValues && vectorValues.length > 0) {
+        for (const [i, vector] of vectorValues.entries()) {
+          if (!vectorDimension) vectorDimension = vector.length;
+          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.
+            metadata: { ...metadata, text: 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."
+        );
+      }
+
+      if (vectors.length > 0) {
+        const chunks = [];
+        const { client } = await this.connect();
+        await this.getOrCreateCollection(client, namespace, vectorDimension);
+
+        this.logger(`Inserting vectorized chunks into ${this.name}.`);
+        for (const chunk of toChunks(vectors, 100)) {
+          chunks.push(chunk);
+          const insertResult = await client.insert({
+            collection_name: this.normalize(namespace),
+            data: chunk.map((item) => ({
+              id: item.id,
+              vector: item.values,
+              metadata: item.metadata,
+            })),
+          });
+
+          if (insertResult?.status.error_code !== "Success") {
+            throw new Error(
+              `Error embedding into ${this.name}! Reason:${insertResult?.status.reason}`
+            );
+          }
+        }
+        await storeVectorResult(chunks, fullFilePath);
+        await client.flushSync({
+          collection_names: [this.normalize(namespace)],
+        });
+      }
+
+      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 knownDocuments = await DocumentVectors.where({ docId });
+    if (knownDocuments.length === 0) return;
+
+    const vectorIds = knownDocuments.map((doc) => doc.vectorId);
+    const queryIn = vectorIds.map((v) => `'${v}'`).join(",");
+    await client.deleteEntities({
+      collection_name: this.normalize(namespace),
+      expr: `id in [${queryIn}]`,
+    });
+
+    const indexes = knownDocuments.map((doc) => doc.id);
+    await DocumentVectors.deleteIds(indexes);
+
+    // Even after flushing Milvus can take some time to re-calc the count
+    // so all we can hope to do is flushSync so that the count can be correct
+    // on a later call.
+    await client.flushSync({ collection_names: [this.normalize(namespace)] });
+    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, namespace))) {
+      return {
+        contextTexts: [],
+        sources: [],
+        message: "Invalid query - no documents found for workspace!",
+      };
+    }
+
+    const queryVector = await LLMConnector.embedTextInput(input);
+    const { contextTexts, sourceDocuments } = await this.similarityResponse({
+      client,
+      namespace,
+      queryVector,
+      similarityThreshold,
+      topN,
+      filterIdentifiers,
+    });
+
+    const sources = sourceDocuments.map((doc, i) => {
+      return { metadata: doc, text: contextTexts[i] };
+    });
+
+    return {
+      contextTexts,
+      sources: this.curateSources(sources),
+      message: false,
+    };
+  }
+
+  async similarityResponse({
+    client,
+    namespace,
+    queryVector,
+    similarityThreshold = 0.25,
+    topN = 4,
+    filterIdentifiers = [],
+  }) {
+    const result = {
+      contextTexts: [],
+      sourceDocuments: [],
+      scores: [],
+    };
+    const response = await client.search({
+      collection_name: this.normalize(namespace),
+      vectors: queryVector,
+      limit: topN,
+    });
+    response.results.forEach((match) => {
+      if (match.score < similarityThreshold) return;
+      if (filterIdentifiers.includes(sourceIdentifier(match.metadata))) {
+        this.logger(
+          `${this.name}: A source was filtered from context as its parent document is pinned.`
+        );
+        return;
+      }
+
+      result.contextTexts.push(match.metadata.text);
+      result.sourceDocuments.push({
+        ...match.metadata,
+        score: match.score,
+      });
+      result.scores.push(match.score);
+    });
+    return result;
+  }
+
+  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, namespace)))
+      throw new Error("Namespace by that name does not exist.");
+    const stats = await this.namespace(client, 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, namespace)))
+      throw new Error("Namespace by that name does not exist.");
+
+    const statistics = await this.namespace(client, namespace);
+    await this.deleteVectorsInNamespace(client, namespace);
+    const vectorCount = Number(statistics?.data?.row_count ?? 0);
+    return {
+      message: `Namespace ${namespace} was deleted along with ${vectorCount} vectors.`,
+    };
+  }
+
+  curateSources(sources = []) {
+    const documents = [];
+    for (const source of sources) {
+      const { metadata = {} } = source;
+      if (Object.keys(metadata).length > 0) {
+        documents.push({
+          ...metadata,
+          ...(source.text ? { text: source.text } : {}),
+        });
+      }
+    }
+    return documents;
+  }
+}
+
+module.exports.Milvus = Milvus;

package/server/utils/vectorDbProviders/pgvector/SETUP.md

@@ -0,0 +1,125 @@
+# Setting up `PGVector` for AnythingLLM
+
+Setting up PGVector for anythingllm to use as your vector database is quite easy. At a minimum, you will need the following:
+
+- PostgreSQL v12+
+- [`pgvector`](https://github.com/pgvector/pgvector) extension installed on DB
+- User with DB table creation perms and READ access
+
+## Setup on Mac (example)
+
+### Install pgvector extension on PostgreSQL DB
+
+```bash
+brew install postgresql
+brew services start postgresql
+brew install pgvector
+
+# assuming you have a database already set up + a user
+psql <database-name>
+CREATE EXTENSION vector;
+```
+
+### Set PG as your vector db
+
+_this can be done via the UI or by directly editing the `.env` file_
+
+First, obtain a valid connection string for the user, credentials, and db you want to target.
+eg: `postgresql://dbuser:dbuserpass@localhost:5432/yourdb`
+
+> [!IMPORTANT]
+> If you have an existing table that you want to use as a vector database, AnythingLLM **requires** that the table be
+> at least minimally conform to the expected schema - this can be seen in the [index.js](./index.js) file.
+
+_optional_ - set a table name you wish to have AnythingLLM store vectors to. By default this is `anythingllm_vectors`
+
+## Common Questions
+
+### I cannot connect to the DB (Running AnythingLLM in Docker)
+
+If you are running AnythingLLM in Docker, you will need to ensure that the DB is accessible from the container.
+If you are running your DB in another Docker container **or** on the host machine, you will need to ensure that the container can access the DB.
+
+`localhost` will not work in this case as it will attempt to connect to the DB _inside the AnythingLLM container_ instead of the host machine or another container.
+
+You will need to use the `host.docker.internal` (or `172.17.0.1` on Linux/Ubuntu) address.
+
+```
+on Mac or Windows:
+postgresql://dbuser:dbuserpass@localhost:5432/yourdb => postgresql://dbuser:dbuserpass@host.docker.internal:5432/yourdb
+
+on Linux:
+postgresql://dbuser:dbuserpass@localhost:5432/yourdb => postgresql://dbuser:dbuserpass@172.17.0.1:5432/yourdb
+```
+
+### Can I use an existing table as a vector database?
+
+Yes, you can use an existing table as a vector database. However, AnythingLLM **requires** that the table be at least minimally conform to the expected schema - this can be seen in the [index.js](./index.js) file.
+
+It is **absolutely critical** that the `embedding` column's `VECTOR(XXXX)` dimensions match the dimension of the embedder in AnythingLLM. The default embedding model is 384 dimensions. However, if you are using a custom embedder, you will need to ensure that the dimension value is set correctly.
+
+### Validate the connection to the database
+
+When setting the connection string in or table name via the AnythingLLM UI, the following validations will be attempted:
+
+- Validate the connection string
+- Validate the table name
+- Run test connection to ensure the table exists and is accessible by the connection string used
+- Check if the table name already exists and if so, validate that it is an embedding table with the correct schema
+
+### My embedding table is not present in the DB
+
+The embedding storage table is created by AnythingLLM **on the first upsert** of a vector. If you have not yet embedding any documents, the table will not be present in the DB.
+
+### How do I reset my vector database?
+
+_at the workspace level in Settings > Vector Database_
+
+You can use the "Reset Vector Database" button in the AnythingLLM UI to reset your vector database. This will drop all vectors within that workspace, but the table will remain in the DB.
+
+_reset the vector database at the db level_
+
+For this, you will need to `DROP TABLE` from the command line or however you manage your DB. Once the table is dropped, it will be recreated by AnythingLLM on the next upsert.
+
+## Troubleshooting
+
+### Cannot connect to DB
+
+- Ensure the connection string is valid
+- Ensure the user has access to the database
+- Ensure the pgvector extension is installed
+
+### Cannot create table
+
+- Ensure the user has `CREATE TABLE` permissions
+
+### Cannot insert vector
+
+- Ensure the user has `INSERT` permissions in the database
+- Ensure the table has a dimension value set and this matches the dimension of the embedder in AnythingLLM
+- Ensure the table has a vector column set
+
+### Cannot query vector
+
+- Ensure the user has `SELECT` permissions in the database
+- Ensure the table has a vector column set
+- Ensure the table has a dimension value set and this matches the dimension of the embedder in AnythingLLM
+
+### "type 'vector' does not exist" issues with PGVector
+
+If you are using the PGVector as your vector database, you may encounter an error similar to the following when embedding documents:
+
+```
+type 'vector' does not exist
+```
+
+This is due to the fact that the `vector` type is not installed on the PG database.
+
+First, follow the instructions in the [PGVector README](https://github.com/pgvector/pgvector#installation) to install the `vector` type on your database.
+
+Then, you will need to create the extension on the database. This can be done by running the following command:
+
+```bash
+psql <database-name>
+CREATE EXTENSION vector;
+```