@sprucelabs/chroma-data-store 1.0.0 → 1.0.2

package/README.md

@@ -0,0 +1,42 @@
+# Chroma Data Store
+
+Give your skill the ability to store and retrieve data from a Chroma database.
+
+## Running the Chroma Database
+
+1. Clone this rep
+2. Run `yarn start.chroma.docker`
+
+## Setting an embedding model
+
+By default, the ChromaDabatase class will use llama3.2 hosted through Ollama to generate embeddings
+
+### Installing Ollama
+
+1. Visit https://ollama.com
+2. Click "Download"
+3. Select your OS
+
+### Installing Llama3.2
+
+Llama 3.2 is the newest version of Llama (as of this writing) that supports embeddings.
+
+1. Inside of terminal, run `ollama run llama3.2`
+2. You should be able to visit http://localhost:11434/api/embeddings and get a 405 response (this is because the route only accepts POST requests)
+
+### Improving embeddings with `nomic-embed-text`
+
+We have seen significantly better search performance when using `nomic-embed-text` to generate embeddings.
+
+Run `ollama pull nomic-embed-text`
+
+
+## Using in your skill
+
+Add the following to your env:
+
+```bash
+DB_CONNECTION_STRING="chromadb://localhost:8000"
+DB_ADAPTER="@sprucelabs/chroma-data-store"
+
+```

package/build/.spruce/settings.json

@@ -1,6 +1,8 @@
 {
   "scriptUpdater": {
-    "skipped": []
+    "skipped": [
+      "test"
+    ]
   },
   "skipped": [
     "skill"

package/build/ChromaDatabase.d.ts

@@ -1,14 +1,19 @@
 import { CreateOptions, Database, DatabaseInternalOptions, Index, IndexWithFilter, QueryOptions } from '@sprucelabs/data-stores';
+import { OllamaEmbeddingFunction } from '@chroma-core/ollama';
 export default class ChromaDatabase implements Database {
+    static Class?: new (connectionString: string) => Database;
+    static EmbeddingFunction: typeof OllamaEmbeddingFunction;
+    private static embeddingFields?;
     private connectionString;
     private client;
     private _isConnected;
     private embeddings;
     private collections;
-    private static embeddingFields?;
     constructor(connectionString: string);
+    static Database(connectionString: string): ChromaDatabase | Database;
     static setEmbeddingsFields(collectionName: string, fields: string[]): void;
     static clearEmbeddingsFields(): void;
+    static getEmbeddingsFields(): Record<string, string[]> | undefined;
     syncUniqueIndexes(_collectionName: string, _indexes: Index[]): Promise<void>;
     syncIndexes(_collectionName: string, _indexes: Index[]): Promise<void>;
     dropIndex(_collectionName: string, _index: Index): Promise<void>;

package/build/ChromaDatabase.js

@@ -6,6 +6,7 @@ Object.defineProperty(exports, "__esModule", { value: true });
 const data_stores_1 = require("@sprucelabs/data-stores");
 const schema_1 = require("@sprucelabs/schema");
 const test_utils_1 = require("@sprucelabs/test-utils");
+const ollama_1 = require("@chroma-core/ollama");
 const chromadb_1 = require("chromadb");
 const SpruceError_1 = __importDefault(require("./errors/SpruceError"));
 class ChromaDatabase {
@@ -13,9 +14,9 @@ class ChromaDatabase {
         this._isConnected = false;
         this.collections = {};
         (0, schema_1.assertOptions)({ connectionString }, ['connectionString']);
-        this.embeddings = new chromadb_1.OllamaEmbeddingFunction({
-            model: 'llama3.2',
-            url: 'http://localhost:11434/api/embeddings',
+        this.embeddings = new ChromaDatabase.EmbeddingFunction({
+            model: process.env.CHROMA_EMBEDDING_MODEL ?? 'llama3.2',
+            url: 'http://localhost:11434',
         });
         if (!connectionString.startsWith('chroma://')) {
             throw new data_stores_1.DataStoresError({
@@ -24,6 +25,9 @@ class ChromaDatabase {
         }
         this.connectionString = connectionString.replace('chroma://', 'http://');
     }
+    static Database(connectionString) {
+        return new (this.Class ?? this)(connectionString);
+    }
     static setEmbeddingsFields(collectionName, fields) {
         if (!this.embeddingFields) {
             this.embeddingFields = {};
@@ -33,6 +37,9 @@ class ChromaDatabase {
     static clearEmbeddingsFields() {
         delete this.embeddingFields;
     }
+    static getEmbeddingsFields() {
+        return this.embeddingFields;
+    }
     async syncUniqueIndexes(_collectionName, _indexes) {
         throw new SpruceError_1.default({
             code: 'FEATURE_NOT_SUPPORTED',
@@ -183,14 +190,19 @@ class ChromaDatabase {
     async find(collection, query, options, _dbOptions) {
         let { ids, where, skipIds } = this.buildQuery(query);
         let matches;
-        const { limit, includeFields } = options ?? {};
+        let { limit, includeFields } = options ?? {};
         const col = await this.getCollection(collection);
+        if (!limit && where?.$prompt) {
+            limit = 10;
+        }
         if (where?.$prompt) {
             const { $prompt, ...rest } = where;
             const queryResults = await col.query({
                 queryTexts: [$prompt],
-                nResults: 1,
-                where: Object.keys(rest).length > 0 ? rest : undefined,
+                nResults: limit,
+                where: (Object.keys(rest).length > 0
+                    ? rest
+                    : undefined),
             });
             matches = {
                 ids: queryResults.ids[0],
@@ -201,7 +213,7 @@ class ChromaDatabase {
             matches = await col.get({
                 ids,
                 include: ['metadatas'],
-                where,
+                where: where,
                 limit: limit == 0 ? 1 : limit,
             });
         }
@@ -255,7 +267,6 @@ class ChromaDatabase {
             where = rest;
         }
         if (where?.$or && where.$or.length === 1) {
-            //@ts-ignore
             where = where.$or[0];
         }
         if (where) {
@@ -348,6 +359,9 @@ class ChromaDatabase {
     }
     async delete(collection, query) {
         const matches = await this.find(collection, query);
+        if (matches.length === 0) {
+            return 0;
+        }
         const col = await this.getCollection(collection);
         await col.delete({
             ids: matches.map((m) => m.id),
@@ -382,4 +396,5 @@ class ChromaDatabase {
         });
     }
 }
+ChromaDatabase.EmbeddingFunction = ollama_1.OllamaEmbeddingFunction;
 exports.default = ChromaDatabase;

package/build/MockChromaDatabase.d.ts

@@ -0,0 +1,8 @@
+import { NeDbDatabase } from '@sprucelabs/data-stores';
+export default class MockChromaDatabase extends NeDbDatabase {
+    private connectionString;
+    constructor(connectionString: string);
+    assertConnectionStringEquals(connectionString: string): void;
+    connect(): Promise<void>;
+    assertIsConnected(): void;
+}

package/build/MockChromaDatabase.js

@@ -0,0 +1,20 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+const data_stores_1 = require("@sprucelabs/data-stores");
+const test_utils_1 = require("@sprucelabs/test-utils");
+class MockChromaDatabase extends data_stores_1.NeDbDatabase {
+    constructor(connectionString) {
+        super();
+        this.connectionString = connectionString;
+    }
+    assertConnectionStringEquals(connectionString) {
+        test_utils_1.assert.isEqual(this.connectionString, connectionString, 'The connection string you passed to your ChromaDb does not equal what I expected!');
+    }
+    async connect() {
+        await super.connect();
+    }
+    assertIsConnected() {
+        test_utils_1.assert.isTrue(this.isConnected(), 'Database is not connected! Try calling db.connect()');
+    }
+}
+exports.default = MockChromaDatabase;

package/build/chroma.types.d.ts

@@ -1,177 +1,24 @@
-import { AddRecordsParams, UpsertRecordsParams, GetResponse, UpdateRecordsParams, QueryRecordsParams, QueryResponse, CollectionMetadata, CollectionParams, PeekParams, DeleteParams, ID, IDs, Where, IncludeEnum, WhereDocument } from 'chromadb';
-export interface Collection {
-    add(params: AddRecordsParams): Promise<void>;
-    /**
-     * Upsert items to the collection
-     * @param {Object} params - The parameters for the query.
-     * @param {ID | IDs} [params.ids] - IDs of the items to add.
-     * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
-     * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
-     * @param {Document | Documents} [params.documents] - Optional documents of the items to add.
-     * @returns {Promise<void>}
-     *
-     * @example
-     * ```typescript
-     * const response = await collection.upsert({
-     *   ids: ["id1", "id2"],
-     *   embeddings: [[1, 2, 3], [4, 5, 6]],
-     *   metadatas: [{ "key": "value" }, { "key": "value" }],
-     *   documents: ["document1", "document2"],
-     * });
-     * ```
-     */
-    upsert(params: UpsertRecordsParams): Promise<void>;
-    /**
-     * Count the number of items in the collection
-     * @returns {Promise<number>} - The number of items in the collection.
-     *
-     * @example
-     * ```typescript
-     * const count = await collection.count();
-     * ```
-     */
-    count(): Promise<number>;
-    /**
-     * Get items from the collection
-     * @param {Object} params - The parameters for the query.
-     * @param {ID | IDs} [params.ids] - Optional IDs of the items to get.
-     * @param {Where} [params.where] - Optional where clause to filter items by.
-     * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
-     * @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
-     * @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response.
-     * @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by.
-     * @returns {Promise<GetResponse>} - The response from the server.
-     *
-     * @example
-     * ```typescript
-     * const response = await collection.get({
-     *   ids: ["id1", "id2"],
-     *   where: { "key": "value" },
-     *   limit: 10,
-     *   offset: 0,
-     *   include: ["embeddings", "metadatas", "documents"],
-     *   whereDocument: { $contains: "value" },
-     * });
-     * ```
-     */
-    get(options?: BaseGetParams): Promise<GetResponse>;
-    /**
-     * Update items in the collection
-     * @param {Object} params - The parameters for the query.
-     * @param {ID | IDs} [params.ids] - IDs of the items to add.
-     * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
-     * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
-     * @param {Document | Documents} [params.documents] - Optional documents of the items to add.
-     * @returns {Promise<void>}
-     *
-     * @example
-     * ```typescript
-     * const response = await collection.update({
-     *   ids: ["id1", "id2"],
-     *   embeddings: [[1, 2, 3], [4, 5, 6]],
-     *   metadatas: [{ "key": "value" }, { "key": "value" }],
-     *   documents: ["document1", "document2"],
-     * });
-     * ```
-     */
-    update(params: UpdateRecordsParams): Promise<void>;
-    /**
-     * Performs a query on the collection using the specified parameters.
-     *
-     * @param {Object} params - The parameters for the query.
-     * @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search.
-     * @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10).
-     * @param {Where} [params.where] - Optional query condition to filter results based on metadata values.
-     * @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection.
-     * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content.
-     * @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document".
-     *
-     * @returns {Promise<QueryResponse>} A promise that resolves to the query results.
-     * @throws {Error} If there is an issue executing the query.
-     * @example
-     * // Query the collection using embeddings
-     * const results = await collection.query({
-     *   queryEmbeddings: [[0.1, 0.2, ...], ...],
-     *   nResults: 10,
-     *   where: {"name": {"$eq": "John Doe"}},
-     *   include: ["metadata", "document"]
-     * });
-     * @example
-     * ```js
-     * // Query the collection using query text
-     * const results = await collection.query({
-     *   queryTexts: "some text",
-     *   nResults: 10,
-     *   where: {"name": {"$eq": "John Doe"}},
-     *   include: ["metadata", "document"]
-     * });
-     * ```
-     *
-     */
-    query({ nResults, where, whereDocument, include, queryTexts, queryEmbeddings, }: QueryRecordsParams): Promise<QueryResponse>;
-    /**
-     * Modify the collection name or metadata
-     * @param {Object} params - The parameters for the query.
-     * @param {string} [params.name] - Optional new name for the collection.
-     * @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection.
-     * @returns {Promise<void>} - The response from the API.
-     *
-     * @example
-     * ```typescript
-     * const response = await client.updateCollection({
-     *   name: "new name",
-     *   metadata: { "key": "value" },
-     * });
-     * ```
-     */
-    modify({ name, metadata, }: {
-        name?: string;
-        metadata?: CollectionMetadata;
-    }): Promise<CollectionParams>;
-    /**
-     * Peek inside the collection
-     * @param {Object} params - The parameters for the query.
-     * @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10).
-     * @returns {Promise<GetResponse>} A promise that resolves to the query results.
-     * @throws {Error} If there is an issue executing the query.
-     *
-     * @example
-     * ```typescript
-     * const results = await collection.peek({
-     *   limit: 10
-     * });
-     * ```
-     */
-    peek(options?: PeekParams): Promise<GetResponse>;
-    /**
-     * Deletes items from the collection.
-     * @param {Object} params - The parameters for deleting items from the collection.
-     * @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete.
-     * @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values.
-     * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content.
-     * @returns {Promise<string[]>} A promise that resolves to the IDs of the deleted items.
-     * @throws {Error} If there is an issue deleting items from the collection.
-     *
-     * @example
-     * ```typescript
-     * const results = await collection.delete({
-     *   ids: "some_id",
-     *   where: {"name": {"$eq": "John Doe"}},
-     *   whereDocument: {"$contains":"search_string"}
-     * });
-     * ```
-     */
-    delete(options?: DeleteParams): Promise<void>;
+import { IncludeEnum, Metadata } from 'chromadb';
+export type ID = string;
+export type Embedding = number[];
+export type Embeddings = Embedding[];
+export interface GetResponse {
+    ids: ID[];
+    embeddings: Embeddings | null;
+    documents: (Document | null)[];
+    metadatas: (Metadata | null)[];
+    included: IncludeEnum[];
 }
-interface BaseGetParams {
-    ids?: ID | IDs;
-    where?: Where;
-    limit?: number;
-    offset?: number;
-    include?: IncludeEnum[];
-    whereDocument?: WhereDocument;
-}
-export type WhereWithPrompt = Where & {
+export type WhereWithPrompt = WhereDocument & {
     $prompt?: string;
 };
-export {};
+export interface WhereDocument {
+    $contains?: string;
+    $not_contains?: string;
+    $matches?: string;
+    $not_matches?: string;
+    $regex?: string;
+    $not_regex?: string;
+    $and?: WhereDocument[];
+    $or?: WhereDocument[];
+}

package/build/chromaDbAssert.utility.d.ts

@@ -0,0 +1,4 @@
+declare const chromaDbAssert: {
+    embeddingFieldsForCollectionEqual(collection: string, fields: string[]): void;
+};
+export default chromaDbAssert;

package/build/chromaDbAssert.utility.js

@@ -0,0 +1,20 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const schema_1 = require("@sprucelabs/schema");
+const test_utils_1 = require("@sprucelabs/test-utils");
+const ChromaDatabase_1 = __importDefault(require("./ChromaDatabase"));
+const chromaDbAssert = {
+    embeddingFieldsForCollectionEqual(collection, fields) {
+        (0, schema_1.assertOptions)({
+            collection,
+            fields,
+        }, ['collection', 'fields']);
+        const actual = ChromaDatabase_1.default.getEmbeddingsFields()?.[collection];
+        test_utils_1.assert.isTruthy(actual, `I could not find any embedding fields for the collection called ${collection}`);
+        test_utils_1.assert.isEqualDeep(actual, fields, `The embedding fields for the collection did not match.`);
+    },
+};
+exports.default = chromaDbAssert;

package/build/esm/ChromaDatabase.d.ts

@@ -1,14 +1,19 @@
 import { CreateOptions, Database, DatabaseInternalOptions, Index, IndexWithFilter, QueryOptions } from '@sprucelabs/data-stores';
+import { OllamaEmbeddingFunction } from '@chroma-core/ollama';
 export default class ChromaDatabase implements Database {
+    static Class?: new (connectionString: string) => Database;
+    static EmbeddingFunction: typeof OllamaEmbeddingFunction;
+    private static embeddingFields?;
     private connectionString;
     private client;
     private _isConnected;
     private embeddings;
     private collections;
-    private static embeddingFields?;
     constructor(connectionString: string);
+    static Database(connectionString: string): ChromaDatabase | Database;
     static setEmbeddingsFields(collectionName: string, fields: string[]): void;
     static clearEmbeddingsFields(): void;
+    static getEmbeddingsFields(): Record<string, string[]> | undefined;
     syncUniqueIndexes(_collectionName: string, _indexes: Index[]): Promise<void>;
     syncIndexes(_collectionName: string, _indexes: Index[]): Promise<void>;
     dropIndex(_collectionName: string, _index: Index): Promise<void>;

package/build/esm/ChromaDatabase.js

@@ -21,16 +21,18 @@ var __rest = (this && this.__rest) || function (s, e) {
 import { DataStoresError, generateId, } from '@sprucelabs/data-stores';
 import { assertOptions, flattenValues, expandValues } from '@sprucelabs/schema';
 import { NULL_PLACEHOLDER } from '@sprucelabs/test-utils';
-import { ChromaClient, OllamaEmbeddingFunction, } from 'chromadb';
+import { OllamaEmbeddingFunction } from '@chroma-core/ollama';
+import { ChromaClient, } from 'chromadb';
 import SpruceError from './errors/SpruceError.js';
-export default class ChromaDatabase {
+class ChromaDatabase {
     constructor(connectionString) {
+        var _a;
         this._isConnected = false;
         this.collections = {};
         assertOptions({ connectionString }, ['connectionString']);
-        this.embeddings = new OllamaEmbeddingFunction({
-            model: 'llama3.2',
-            url: 'http://localhost:11434/api/embeddings',
+        this.embeddings = new ChromaDatabase.EmbeddingFunction({
+            model: (_a = process.env.CHROMA_EMBEDDING_MODEL) !== null && _a !== void 0 ? _a : 'llama3.2',
+            url: 'http://localhost:11434',
         });
         if (!connectionString.startsWith('chroma://')) {
             throw new DataStoresError({
@@ -39,6 +41,10 @@ export default class ChromaDatabase {
         }
         this.connectionString = connectionString.replace('chroma://', 'http://');
     }
+    static Database(connectionString) {
+        var _a;
+        return new ((_a = this.Class) !== null && _a !== void 0 ? _a : this)(connectionString);
+    }
     static setEmbeddingsFields(collectionName, fields) {
         if (!this.embeddingFields) {
             this.embeddingFields = {};
@@ -48,6 +54,9 @@ export default class ChromaDatabase {
     static clearEmbeddingsFields() {
         delete this.embeddingFields;
     }
+    static getEmbeddingsFields() {
+        return this.embeddingFields;
+    }
     syncUniqueIndexes(_collectionName, _indexes) {
         return __awaiter(this, void 0, void 0, function* () {
             throw new SpruceError({
@@ -227,14 +236,19 @@ export default class ChromaDatabase {
         return __awaiter(this, void 0, void 0, function* () {
             let { ids, where, skipIds } = this.buildQuery(query);
             let matches;
-            const { limit, includeFields } = options !== null && options !== void 0 ? options : {};
+            let { limit, includeFields } = options !== null && options !== void 0 ? options : {};
             const col = yield this.getCollection(collection);
+            if (!limit && (where === null || where === void 0 ? void 0 : where.$prompt)) {
+                limit = 10;
+            }
             if (where === null || where === void 0 ? void 0 : where.$prompt) {
                 const { $prompt } = where, rest = __rest(where, ["$prompt"]);
                 const queryResults = yield col.query({
                     queryTexts: [$prompt],
-                    nResults: 1,
-                    where: Object.keys(rest).length > 0 ? rest : undefined,
+                    nResults: limit,
+                    where: (Object.keys(rest).length > 0
+                        ? rest
+                        : undefined),
                 });
                 matches = {
                     ids: queryResults.ids[0],
@@ -245,7 +259,7 @@ export default class ChromaDatabase {
                 matches = yield col.get({
                     ids,
                     include: ['metadatas'],
-                    where,
+                    where: where,
                     limit: limit == 0 ? 1 : limit,
                 });
             }
@@ -298,7 +312,6 @@ export default class ChromaDatabase {
             where = rest;
         }
         if ((where === null || where === void 0 ? void 0 : where.$or) && where.$or.length === 1) {
-            //@ts-ignore
             where = where.$or[0];
         }
         if (where) {
@@ -398,6 +411,9 @@ export default class ChromaDatabase {
     delete(collection, query) {
         return __awaiter(this, void 0, void 0, function* () {
             const matches = yield this.find(collection, query);
+            if (matches.length === 0) {
+                return 0;
+            }
             const col = yield this.getCollection(collection);
             yield col.delete({
                 ids: matches.map((m) => m.id),
@@ -443,3 +459,5 @@ export default class ChromaDatabase {
         });
     }
 }
+ChromaDatabase.EmbeddingFunction = OllamaEmbeddingFunction;
+export default ChromaDatabase;

package/build/esm/MockChromaDatabase.d.ts

@@ -0,0 +1,8 @@
+import { NeDbDatabase } from '@sprucelabs/data-stores';
+export default class MockChromaDatabase extends NeDbDatabase {
+    private connectionString;
+    constructor(connectionString: string);
+    assertConnectionStringEquals(connectionString: string): void;
+    connect(): Promise<void>;
+    assertIsConnected(): void;
+}

package/build/esm/MockChromaDatabase.js

@@ -0,0 +1,31 @@
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+import { NeDbDatabase } from '@sprucelabs/data-stores';
+import { assert } from '@sprucelabs/test-utils';
+export default class MockChromaDatabase extends NeDbDatabase {
+    constructor(connectionString) {
+        super();
+        this.connectionString = connectionString;
+    }
+    assertConnectionStringEquals(connectionString) {
+        assert.isEqual(this.connectionString, connectionString, 'The connection string you passed to your ChromaDb does not equal what I expected!');
+    }
+    connect() {
+        const _super = Object.create(null, {
+            connect: { get: () => super.connect }
+        });
+        return __awaiter(this, void 0, void 0, function* () {
+            yield _super.connect.call(this);
+        });
+    }
+    assertIsConnected() {
+        assert.isTrue(this.isConnected(), 'Database is not connected! Try calling db.connect()');
+    }
+}