@ixo/data-store 1.0.1

package/src/chroma/chroma-data-store.ts

@@ -0,0 +1,254 @@
+import {
+  ChromaClient,
+  type Collection,
+  type Embeddings,
+  type Metadata,
+  type QueryRecordsParams,
+} from 'chromadb';
+import {
+  VectorDBDataStore,
+  type IVectorStoreDocument,
+  type IVectorStoreDocumentWithEmbeddings,
+  type IVectorStoreOptions,
+  type IVectorStoreQueryOptions,
+} from '../types/vector-db-data-store';
+import { type IChromaMetadataFilter } from './chroma.types';
+import { OpenAIEmbeddingFunction } from './embedding-function';
+import { createVectorStoreSearchResult } from './utils';
+
+/**
+ * ChromaDataStore class for managing vector storage and retrieval using ChromaDB
+ */
+class ChromaDataStore extends VectorDBDataStore {
+  private readonly client: ChromaClient;
+  private collection: Collection;
+
+  /**
+   * Creates a new ChromaDataStore instance
+   * @param options - Configuration options for the vector store
+   * @throws Error if OPENAI_API_KEY is not set and no embedding function is provided
+   */
+  constructor(options: IVectorStoreOptions) {
+    if (
+      typeof process.env.OPENAI_API_KEY !== 'string' &&
+      !options.embeddingFunction
+    ) {
+      throw new Error(
+        'OPENAI_API_KEY is not set and no embedding function is provided',
+      );
+    }
+    options.embeddingFunction =
+      options.embeddingFunction ||
+      new OpenAIEmbeddingFunction({
+        openai_api_key: process.env.OPENAI_API_KEY ?? '',
+        openai_model: 'text-embedding-3-small',
+      });
+    options.url = options.url || 'http://localhost:8000';
+    options.collectionName = options.collectionName || 'default-vector-store';
+    super(options);
+    this.client = new ChromaClient({
+      path: options.url,
+    });
+
+    // this is a workaround to avoid setting collection as optional
+    this.collection = undefined as unknown as Collection;
+  }
+
+  /**
+   * Checks if the ChromaDB collection is initialized
+   * @throws Error if collection is not initialized
+   */
+  private checkIsInitialized(): void {
+    if (!this.collection as unknown) {
+      throw new Error('ChromaDataStore is not initialized');
+    }
+  }
+
+  /**
+   * Initializes the ChromaDB collection
+   */
+  async init(): Promise<void> {
+    if (this.collection as unknown) {
+      return;
+    }
+    const collection = await this.client.getOrCreateCollection({
+      name: this.options.collectionName,
+      embeddingFunction: this.options.embeddingFunction,
+    });
+    this.collection = collection;
+  }
+
+  /**
+   * Queries the vector store using text
+   * @param query - Text query to search for
+   * @param options - Query options including filters and top-k results
+   * @returns Array of matching documents
+   */
+  async query(
+    query: string,
+    options?: IVectorStoreQueryOptions<IChromaMetadataFilter>,
+  ): Promise<IVectorStoreDocument[]> {
+    this.checkIsInitialized();
+    const params: QueryRecordsParams = {
+      queryTexts: [query],
+      nResults: options?.topK || 10,
+    };
+    if (options?.filters) {
+      params.where = options.filters;
+    }
+    const result = await this.collection.query(params);
+    return createVectorStoreSearchResult(result);
+  }
+
+  /**
+   * Queries the vector store and filters results by similarity threshold
+   * @param query - Text query to search for
+   * @param options - Query options including similarity threshold
+   * @returns Array of documents meeting the similarity threshold
+   */
+  async queryWithSimilarity(
+    query: string,
+    options?: IVectorStoreQueryOptions<IChromaMetadataFilter> & {
+      similarityThreshold: number;
+    },
+  ): Promise<IVectorStoreDocument[]> {
+    this.checkIsInitialized();
+    const result = await this.query(query, options);
+    return result.filter(
+      (doc) =>
+        (doc.score && doc.score >= (options?.similarityThreshold ?? 0.5)) ??
+        false,
+    );
+  }
+
+  /**
+   * Upserts (inserts or updates) documents into the vector store
+   * @param documents - Array of documents to upsert
+   * @throws Error if any document is missing an ID
+   */
+  async upsert(documents: IVectorStoreDocument[]): Promise<void> {
+    this.checkIsInitialized();
+    const [ids, contents, metadatas]: [string[], string[], Metadata[]] =
+      documents.reduce<[string[], string[], Metadata[]]>(
+        (acc, doc) => {
+          const id = doc.id;
+          const metadata = doc.metadata ?? {};
+          const content = doc.content;
+          if (!id) {
+            throw new Error('Document ID is required');
+          }
+          acc[0].push(id);
+          acc[1].push(content);
+          acc[2].push(metadata);
+          return acc;
+        },
+        [[], [], []],
+      );
+
+    await this.collection.upsert({
+      ids,
+      documents: contents,
+      metadatas,
+    });
+  }
+
+  async addDocumentsWithEmbeddings(
+    documents: IVectorStoreDocumentWithEmbeddings[],
+  ): Promise<void> {
+    this.checkIsInitialized();
+    const [ids, embeddings, metadatas, contents]: [
+      string[],
+      number[][],
+      Metadata[],
+      string[],
+    ] = documents.reduce<[string[], Embeddings, Metadata[], string[]]>(
+      (acc, doc) => {
+        const id = doc.id;
+        const metadata = doc.metadata ?? {};
+        const embedding = doc.embedding;
+        if (!id) {
+          throw new Error('Document ID is required');
+        }
+        if (embedding.length < 1536) {
+          throw new Error('Embedding is required');
+        }
+        acc[0].push(id);
+        acc[1].push(embedding);
+        acc[2].push(metadata);
+        acc[3].push(doc.content);
+        return acc;
+      },
+      [[], [], [], []],
+    );
+
+    await this.collection.add({
+      ids,
+      embeddings,
+      metadatas,
+      documents: contents,
+    });
+  }
+
+  /**
+   * Deletes documents from the vector store by their IDs
+   * @param ids - Array of document IDs to delete
+   */
+  async delete(ids: string[]): Promise<void> {
+    this.checkIsInitialized();
+    await this.collection.delete({
+      ids,
+    });
+  }
+
+  /**
+   * Retrieves a document by its ID
+   * @param id - Document ID to retrieve
+   * @returns The document if found, null otherwise
+   */
+  async getById(id: string): Promise<IVectorStoreDocument | null> {
+    const result = await this.collection.get({
+      ids: [id],
+    });
+    return result.documents.at(0)
+      ? {
+          id: result.ids.at(0)?.toString() ?? '',
+          content: result.documents.at(0) ?? '',
+          metadata: result.metadatas.at(0) ?? {},
+          score: undefined,
+        }
+      : null;
+  }
+
+  /**
+   * Queries the vector store using a pre-computed embedding vector
+   * @param vector - The embedding vector to search with
+   * @param options - Query options including top-k results
+   * @returns Array of matching documents
+   */
+  async queryByVector(
+    vector: number[],
+    options?: IVectorStoreQueryOptions,
+  ): Promise<IVectorStoreDocument[]> {
+    this.checkIsInitialized();
+    const result = await this.collection.query({
+      queryEmbeddings: [vector],
+      nResults: options?.topK || 10,
+    });
+    return createVectorStoreSearchResult(result);
+  }
+
+  /**
+   * Updates the metadata for a document by its ID
+   * @param ids - Array of document IDs to update
+   * @param metadata - Metadata to update
+   */
+  async updateMetadata(ids: string[], metadatas: Metadata[]): Promise<void> {
+    this.checkIsInitialized();
+    await this.collection.update({
+      ids,
+      metadatas,
+    });
+  }
+}
+
+export { ChromaDataStore };

package/src/chroma/chroma.types.ts

@@ -0,0 +1,28 @@
+/**
+ * Chroma-specific document content filter operators
+ */
+export interface IChromaDocumentFilter {
+  $contains?: string;
+  $not_contains?: string;
+}
+
+/**
+ * Chroma-specific metadata filter operators
+ */
+export type IChromaMetadataFilter = Record<
+  string,
+  string | number | boolean | null | IChromaMetadataOperators
+>;
+
+export interface IChromaMetadataOperators {
+  $eq?: unknown;
+  $ne?: unknown;
+  $gt?: number;
+  $gte?: number;
+  $lt?: number;
+  $lte?: number;
+  $in?: unknown[];
+  $nin?: unknown[];
+  $and?: IChromaMetadataFilter[];
+  $or?: IChromaMetadataFilter[];
+}

package/src/chroma/embedding-function.ts

@@ -0,0 +1,26 @@
+import { type IEmbeddingFunction } from 'chromadb';
+import OpenAI from 'openai';
+
+interface IOpenAIEmbeddingFunctionOptions {
+  openai_api_key: string;
+  openai_model: string;
+  openai_organization_id?: string;
+}
+
+export class OpenAIEmbeddingFunction implements IEmbeddingFunction {
+  private readonly openai: OpenAI;
+  constructor(private readonly options: IOpenAIEmbeddingFunctionOptions) {
+    this.openai = new OpenAI({
+      apiKey: options.openai_api_key,
+      organization: options.openai_organization_id,
+    });
+  }
+
+  async generate(texts: string[]): Promise<number[][]> {
+    const embeddings = await this.openai.embeddings.create({
+      model: this.options.openai_model,
+      input: texts,
+    });
+    return embeddings.data.map((embedding) => embedding.embedding);
+  }
+}

package/src/chroma/index.ts

@@ -0,0 +1,2 @@
+export * from './chroma-data-store';
+export * from './chroma.types';

package/src/chroma/utils.ts

@@ -0,0 +1,40 @@
+import { type Collection } from 'chromadb';
+import { type IVectorStoreDocument } from 'types/vector-db-data-store';
+
+type MultiQueryResponse = Awaited<
+  ReturnType<typeof Collection.prototype.query>
+>;
+
+export const createVectorStoreSearchResult = (
+  result: MultiQueryResponse,
+): IVectorStoreDocument[] => {
+  const scores = result.distances
+    ? convertCosineDistancesToScores(result.distances)
+    : undefined;
+  const metadatas = result.metadatas.flat().filter(Boolean);
+  return result.documents.flat().map((document, idx) => ({
+    id: result.ids.flat().at(idx)?.toString() ?? '',
+    content: document ?? '',
+    metadata: metadatas.at(idx) ?? {},
+    score: scores?.at(idx),
+  }));
+};
+
+/**
+ * Converts a 2D array of cosine distances into a flattened array of similarity scores.
+ * @param distances - A 2D array of cosine distances (e.g., [[1.04, 1.24], [0.8, 1.5]]).
+ * @returns A flattened array of similarity scores (e.g., [0.48, 0.38, 0.6, 0.25]).
+ */
+function convertCosineDistancesToScores(distances: number[][]): number[] {
+  return distances.flat().map((distance) => {
+    // 1. Validate that the distance is within the valid range [0, 2].
+    if (distance < 0 || distance > 2) {
+      throw new Error(
+        `Cosine distance must be between 0 and 2. Received: ${distance}`,
+      );
+    }
+
+    // 2. Convert cosine distance (d) to similarity score (s): s = 1 - (d / 2).
+    return 1 - distance / 2;
+  });
+}

package/src/index.ts

@@ -0,0 +1,4 @@
+export * from './airtable-store';
+export * from './chroma';
+
+export * from './types';

package/src/types/index.ts

@@ -0,0 +1,2 @@
+export * from './structured-data-store';
+export * from './vector-db-data-store';

package/src/types/structured-data-store.ts

@@ -0,0 +1,34 @@
+interface IQueryParams<T> {
+  filterByFormula?: string;
+  maxRecords?: number;
+  pageSize?: number;
+  fields?: (keyof T)[];
+  view?: string;
+}
+
+export interface IDataStore<T> {
+  getAllRecords: (
+    tableName: string,
+    selectOptions: IQueryParams<T>,
+  ) => Promise<T[]>;
+  /**
+   * Get a record by its ID or throw an error if not found
+   */
+  getRecord: (tableName: string, recordId: string) => Promise<T>;
+  createRecord: (tableName: string, recordData: T) => Promise<T>;
+  updateRecord: (
+    tableName: string,
+    recordId: string,
+    recordData: T,
+  ) => Promise<T>;
+  batchUpdateRecords: (
+    tableName: string,
+    records: { id: string; fields: T }[],
+  ) => Promise<T[]>;
+  deleteRecord: (tableName: string, recordId: string) => Promise<T>;
+  getRecordByField: (
+    tableName: string,
+    fieldName: string,
+    fieldValue: string,
+  ) => Promise<T[]>;
+}

package/src/types/vector-db-data-store.ts

@@ -0,0 +1,78 @@
+import { type IEmbeddingFunction, type Metadata } from 'chromadb';
+
+/**
+ * Represents a document in the vector store.
+ */
+export interface IVectorStoreDocument {
+  /**
+   * Unique identifier for the document.
+   */
+  id: string;
+
+  /**
+   * The content or data of the document.
+   */
+  content: string;
+
+  /**
+   * Optional metadata associated with the document.
+   */
+  metadata?: Metadata;
+
+  /**
+   * The score of the document.
+   */
+  score?: number;
+}
+
+export interface IVectorStoreDocumentWithEmbeddings
+  extends IVectorStoreDocument {
+  embedding: number[];
+}
+
+/**
+ * Options for querying the vector store.
+ */
+export interface IVectorStoreQueryOptions<
+  Filters extends Record<string, unknown> = Record<string, unknown>,
+> {
+  /**
+   * The number of top results to return.
+   */
+  topK?: number;
+
+  /**
+   * Filters to apply based on document metadata.
+   */
+  filters?: Filters;
+}
+
+export interface IVectorStoreOptions {
+  collectionName: string;
+  url: string;
+  embeddingFunction?: IEmbeddingFunction;
+}
+
+export abstract class VectorDBDataStore {
+  constructor(protected readonly options: IVectorStoreOptions) {}
+
+  abstract upsert(documents: IVectorStoreDocument[]): Promise<void>;
+  abstract delete(ids: string[]): Promise<void>;
+  abstract queryByVector(
+    vector: number[],
+    options?: IVectorStoreQueryOptions,
+  ): Promise<IVectorStoreDocument[]>;
+  abstract query(
+    query: string,
+    options?: IVectorStoreQueryOptions,
+  ): Promise<IVectorStoreDocument[]>;
+  abstract getById(id: string): Promise<IVectorStoreDocument | null>;
+  abstract queryWithSimilarity(
+    query: string,
+    options?: IVectorStoreQueryOptions & { similarityThreshold: number },
+  ): Promise<IVectorStoreDocument[]>;
+  abstract addDocumentsWithEmbeddings(
+    documents: IVectorStoreDocumentWithEmbeddings[],
+  ): Promise<void>;
+  abstract init(): Promise<void>;
+}

package/src/utils/index.ts

@@ -0,0 +1 @@
+export * from './with-report-error';

package/src/utils/with-report-error.ts

@@ -0,0 +1,18 @@
+import { Logger } from '@ixo/logger';
+
+/**
+ * should report using the logger and throw the error if the promise is rejected
+ */
+export const withReportError = async <T>(promise: Promise<T>): Promise<T> => {
+  try {
+    const res = await promise;
+    return res;
+  } catch (error) {
+    const errorMessage =
+      error instanceof Error
+        ? error.message
+        : 'Something went wrong in AirtableDataStore';
+    Logger.error(errorMessage, error);
+    throw error;
+  }
+};

package/tsconfig.json

@@ -0,0 +1,7 @@
+{
+  "extends": "@ixo/typescript-config/base.json",
+  "compilerOptions": {
+    "baseUrl": "./src",
+    "outDir": "./dist"
+  }
+}