@baselineos/protocol-core 1.0.0

package/src/knowledge/vector-store.ts

@@ -0,0 +1,722 @@
+// ChromaDB Vector Store wrapper
+// chromadb is an optional peer dependency - do NOT import it directly.
+// The client and collection are typed as `unknown` and resolved at runtime.
+
+// ── Interfaces ──────────────────────────────────────────────────────────────────
+
+export interface VectorStoreConfig {
+  host?: string;
+  port?: number;
+  collectionName?: string;
+  embeddingFunction?: string;
+  metadata?: Record<string, string>;
+}
+
+export interface VectorStoreMetrics {
+  totalDocuments: number;
+  totalQueries: number;
+  averageQueryTime: number;
+  lastOptimization: number;
+}
+
+export interface DocumentInput {
+  ids: string[];
+  embeddings?: number[][];
+  metadatas?: Record<string, unknown>[];
+  documents: string[];
+}
+
+export interface VectorQueryResult {
+  success: boolean;
+  query?: string;
+  results?: unknown;
+  error?: string;
+  metadata?: Record<string, unknown>;
+}
+
+export interface VectorQueryOptions {
+  nResults?: number;
+  where?: Record<string, unknown>;
+  whereDocument?: Record<string, unknown>;
+  include?: string[];
+}
+
+export interface FilterSearchOptions {
+  where?: Record<string, unknown>;
+  whereDocument?: Record<string, unknown>;
+  nResults?: number;
+  include?: string[];
+}
+
+export interface BatchOperation {
+  type: 'add' | 'update' | 'delete';
+  data: Record<string, unknown>;
+}
+
+export interface CollectionStats {
+  name: string;
+  count: number;
+  metadata: Record<string, unknown>;
+}
+
+// ── ChromaVectorStore ───────────────────────────────────────────────────────────
+
+export class ChromaVectorStore {
+  private config: Required<VectorStoreConfig>;
+  private client: unknown;
+  private collection: unknown;
+  private isInitialized: boolean;
+  private metrics: VectorStoreMetrics;
+  private queryCache: Map<string, { result: unknown; timestamp: number }>;
+  private cacheMaxAge: number;
+
+  constructor(config: VectorStoreConfig = {}) {
+    this.config = {
+      host: config.host ?? 'localhost',
+      port: config.port ?? 8000,
+      collectionName: config.collectionName ?? 'baseline-vectors',
+      embeddingFunction: config.embeddingFunction ?? 'default',
+      metadata: config.metadata ?? {},
+    };
+
+    this.client = null;
+    this.collection = null;
+    this.isInitialized = false;
+
+    this.metrics = {
+      totalDocuments: 0,
+      totalQueries: 0,
+      averageQueryTime: 0,
+      lastOptimization: Date.now(),
+    };
+
+    this.queryCache = new Map();
+    this.cacheMaxAge = 5 * 60 * 1000; // 5 minutes
+  }
+
+  // ── Initialization ──────────────────────────────────────────────────────────
+
+  async initializeClient(): Promise<void> {
+    if (this.isInitialized) {
+      console.log('[VectorStore] Client already initialized.');
+      return;
+    }
+
+    try {
+      // Dynamic import of chromadb (optional peer dependency)
+      let chromadb: Record<string, unknown>;
+      try {
+        chromadb = (await import('chromadb')) as Record<string, unknown>;
+      } catch {
+        console.log('[VectorStore] chromadb package not found. Install it with: npm install chromadb');
+        console.log('[VectorStore] Running in fallback mode without vector storage.');
+        this.isInitialized = false;
+        return;
+      }
+
+      const ChromaClient = chromadb.ChromaClient as new (config: Record<string, unknown>) => unknown;
+
+      this.client = new ChromaClient({
+        path: `http://${this.config.host}:${this.config.port}`,
+      });
+
+      await this.initializeCollection();
+      this.isInitialized = true;
+
+      console.log(`[VectorStore] Client initialized. Connected to ${this.config.host}:${this.config.port}.`);
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to initialize client: ${message}`);
+      console.log('[VectorStore] Running in fallback mode without vector storage.');
+      this.isInitialized = false;
+    }
+  }
+
+  private async initializeCollection(): Promise<void> {
+    if (!this.client) {
+      throw new Error('Client not initialized');
+    }
+
+    try {
+      const client = this.client as Record<string, (...args: unknown[]) => Promise<unknown>>;
+      const embeddingFunction = this.createEmbeddingFunction();
+
+      this.collection = await client.getOrCreateCollection({
+        name: this.config.collectionName,
+        metadata: this.config.metadata,
+        embeddingFunction,
+      });
+
+      console.log(`[VectorStore] Collection "${this.config.collectionName}" ready.`);
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to initialize collection: ${message}`);
+      throw error;
+    }
+  }
+
+  // ── Document Operations ─────────────────────────────────────────────────────
+
+  async addDocuments(input: DocumentInput): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      const addPayload: Record<string, unknown> = {
+        ids: input.ids,
+        documents: input.documents,
+      };
+
+      if (input.embeddings) {
+        addPayload.embeddings = input.embeddings;
+      }
+
+      if (input.metadatas) {
+        addPayload.metadatas = input.metadatas;
+      }
+
+      await collection.add(addPayload);
+
+      this.metrics.totalDocuments += input.ids.length;
+
+      console.log(`[VectorStore] Added ${input.ids.length} documents.`);
+
+      return {
+        success: true,
+        metadata: {
+          addedCount: input.ids.length,
+          totalDocuments: this.metrics.totalDocuments,
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to add documents: ${message}`);
+      return {
+        success: false,
+        error: message,
+      };
+    }
+  }
+
+  async query(queryText: string, options: VectorQueryOptions = {}): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        query: queryText,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    const startTime = Date.now();
+
+    // Check cache
+    const cacheKey = `${queryText}:${JSON.stringify(options)}`;
+    const cached = this.queryCache.get(cacheKey);
+    if (cached && Date.now() - cached.timestamp < this.cacheMaxAge) {
+      this.metrics.totalQueries++;
+      return {
+        success: true,
+        query: queryText,
+        results: cached.result,
+        metadata: { fromCache: true, queryTime: 0 },
+      };
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      const queryPayload: Record<string, unknown> = {
+        queryTexts: [queryText],
+        nResults: options.nResults ?? 10,
+      };
+
+      if (options.where) {
+        queryPayload.where = options.where;
+      }
+
+      if (options.whereDocument) {
+        queryPayload.whereDocument = options.whereDocument;
+      }
+
+      if (options.include) {
+        queryPayload.include = options.include;
+      }
+
+      const results = await collection.query(queryPayload);
+
+      const queryTime = Date.now() - startTime;
+
+      // Update metrics
+      this.metrics.totalQueries++;
+      this.metrics.averageQueryTime =
+        (this.metrics.averageQueryTime * (this.metrics.totalQueries - 1) + queryTime) /
+        this.metrics.totalQueries;
+
+      // Update cache
+      this.queryCache.set(cacheKey, { result: results, timestamp: Date.now() });
+
+      return {
+        success: true,
+        query: queryText,
+        results,
+        metadata: {
+          queryTime,
+          nResults: options.nResults ?? 10,
+          fromCache: false,
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Query failed: ${message}`);
+      return {
+        success: false,
+        query: queryText,
+        error: message,
+      };
+    }
+  }
+
+  async semanticSearch(queryText: string, nResults: number = 10): Promise<VectorQueryResult> {
+    return this.query(queryText, {
+      nResults,
+      include: ['documents', 'metadatas', 'distances'],
+    });
+  }
+
+  async filterSearch(queryText: string, options: FilterSearchOptions = {}): Promise<VectorQueryResult> {
+    return this.query(queryText, {
+      nResults: options.nResults ?? 10,
+      where: options.where,
+      whereDocument: options.whereDocument,
+      include: options.include ?? ['documents', 'metadatas', 'distances'],
+    });
+  }
+
+  async fullTextSearch(searchText: string, nResults: number = 10): Promise<VectorQueryResult> {
+    return this.query(searchText, {
+      nResults,
+      whereDocument: { $contains: searchText },
+      include: ['documents', 'metadatas', 'distances'],
+    });
+  }
+
+  async regexSearch(pattern: string, nResults: number = 10): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        query: pattern,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      // ChromaDB doesn't natively support regex, so we do a broad query then filter
+      const allResults = await collection.get({
+        include: ['documents', 'metadatas'],
+      });
+
+      const resultsObj = allResults as Record<string, unknown[] | null>;
+      const ids = (resultsObj.ids as string[]) ?? [];
+      const documents = (resultsObj.documents as (string | null)[]) ?? [];
+      const metadatas = (resultsObj.metadatas as (Record<string, unknown> | null)[]) ?? [];
+
+      const regex = new RegExp(pattern, 'i');
+      const matchedIds: string[] = [];
+      const matchedDocuments: string[] = [];
+      const matchedMetadatas: Record<string, unknown>[] = [];
+
+      for (let i = 0; i < ids.length; i++) {
+        const doc = documents[i];
+        if (doc && regex.test(doc)) {
+          matchedIds.push(ids[i]);
+          matchedDocuments.push(doc);
+          matchedMetadatas.push(metadatas[i] ?? {});
+
+          if (matchedIds.length >= nResults) break;
+        }
+      }
+
+      return {
+        success: true,
+        query: pattern,
+        results: {
+          ids: [matchedIds],
+          documents: [matchedDocuments],
+          metadatas: [matchedMetadatas],
+        },
+        metadata: {
+          matchCount: matchedIds.length,
+          searchType: 'regex',
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Regex search failed: ${message}`);
+      return {
+        success: false,
+        query: pattern,
+        error: message,
+      };
+    }
+  }
+
+  async updateDocuments(input: DocumentInput): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      const updatePayload: Record<string, unknown> = {
+        ids: input.ids,
+        documents: input.documents,
+      };
+
+      if (input.embeddings) {
+        updatePayload.embeddings = input.embeddings;
+      }
+
+      if (input.metadatas) {
+        updatePayload.metadatas = input.metadatas;
+      }
+
+      await collection.update(updatePayload);
+
+      // Invalidate cache for affected documents
+      this.queryCache.clear();
+
+      console.log(`[VectorStore] Updated ${input.ids.length} documents.`);
+
+      return {
+        success: true,
+        metadata: {
+          updatedCount: input.ids.length,
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to update documents: ${message}`);
+      return {
+        success: false,
+        error: message,
+      };
+    }
+  }
+
+  async deleteDocuments(ids: string[]): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      await collection.delete({ ids });
+
+      this.metrics.totalDocuments = Math.max(0, this.metrics.totalDocuments - ids.length);
+
+      // Invalidate cache
+      this.queryCache.clear();
+
+      console.log(`[VectorStore] Deleted ${ids.length} documents.`);
+
+      return {
+        success: true,
+        metadata: {
+          deletedCount: ids.length,
+          totalDocuments: this.metrics.totalDocuments,
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to delete documents: ${message}`);
+      return {
+        success: false,
+        error: message,
+      };
+    }
+  }
+
+  // ── Collection Management ───────────────────────────────────────────────────
+
+  async getCollectionStats(): Promise<CollectionStats | null> {
+    if (!this.isInitialized || !this.collection) {
+      console.log('[VectorStore] Not initialized, cannot get collection stats.');
+      return null;
+    }
+
+    try {
+      const collection = this.collection as Record<string, (...args: unknown[]) => Promise<unknown>>;
+
+      const count = (await collection.count()) as number;
+      const peek = (await collection.peek({ limit: 1 })) as Record<string, unknown>;
+
+      this.metrics.totalDocuments = count;
+
+      return {
+        name: this.config.collectionName,
+        count,
+        metadata: {
+          ...this.config.metadata,
+          peekSample: peek,
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Failed to get collection stats: ${message}`);
+      return null;
+    }
+  }
+
+  async optimizeCollection(): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    console.log('[VectorStore] Starting collection optimization...');
+
+    try {
+      // Clear the query cache
+      const cacheSize = this.queryCache.size;
+      this.queryCache.clear();
+
+      // Update optimization timestamp
+      this.metrics.lastOptimization = Date.now();
+
+      // Get collection stats for reference
+      const stats = await this.getCollectionStats();
+
+      console.log(`[VectorStore] Optimization complete. Cache cleared (${cacheSize} entries).`);
+
+      return {
+        success: true,
+        metadata: {
+          cacheEntriesCleared: cacheSize,
+          collectionStats: stats,
+          optimizationTime: Date.now(),
+        },
+      };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.log(`[VectorStore] Optimization failed: ${message}`);
+      return {
+        success: false,
+        error: message,
+      };
+    }
+  }
+
+  // ── Embedding Function ──────────────────────────────────────────────────────
+
+  private createEmbeddingFunction(): unknown {
+    const funcType = this.config.embeddingFunction;
+
+    switch (funcType) {
+      case 'openai': {
+        console.log('[VectorStore] Using OpenAI embedding function.');
+        return {
+          generate: async (texts: string[]): Promise<number[][]> => {
+            // Placeholder: in production, call the OpenAI embedding API
+            return texts.map((_text: string) => {
+              const embedding: number[] = [];
+              for (let i = 0; i < 1536; i++) {
+                embedding.push(Math.random() * 2 - 1);
+              }
+              return embedding;
+            });
+          },
+        };
+      }
+
+      case 'cohere': {
+        console.log('[VectorStore] Using Cohere embedding function.');
+        return {
+          generate: async (texts: string[]): Promise<number[][]> => {
+            // Placeholder: in production, call the Cohere embedding API
+            return texts.map((_text: string) => {
+              const embedding: number[] = [];
+              for (let i = 0; i < 768; i++) {
+                embedding.push(Math.random() * 2 - 1);
+              }
+              return embedding;
+            });
+          },
+        };
+      }
+
+      case 'huggingface': {
+        console.log('[VectorStore] Using HuggingFace embedding function.');
+        return {
+          generate: async (texts: string[]): Promise<number[][]> => {
+            // Placeholder: in production, call the HuggingFace API
+            return texts.map((_text: string) => {
+              const embedding: number[] = [];
+              for (let i = 0; i < 384; i++) {
+                embedding.push(Math.random() * 2 - 1);
+              }
+              return embedding;
+            });
+          },
+        };
+      }
+
+      case 'default':
+      default: {
+        console.log('[VectorStore] Using default embedding function.');
+        return {
+          generate: async (texts: string[]): Promise<number[][]> => {
+            return texts.map((text: string) => {
+              // Simple hash-based embedding for fallback
+              const embedding: number[] = [];
+              for (let i = 0; i < 384; i++) {
+                let hash = 0;
+                for (let j = 0; j < text.length; j++) {
+                  hash = ((hash << 5) - hash + text.charCodeAt(j) + i) | 0;
+                }
+                embedding.push((hash % 1000) / 1000);
+              }
+              return embedding;
+            });
+          },
+        };
+      }
+    }
+  }
+
+  // ── Batch Operations ────────────────────────────────────────────────────────
+
+  async batchOperation(operations: BatchOperation[]): Promise<VectorQueryResult> {
+    if (!this.isInitialized || !this.collection) {
+      return {
+        success: false,
+        error: 'Vector store not initialized. Call initializeClient() first.',
+      };
+    }
+
+    const results: Array<{ type: string; success: boolean; error?: string }> = [];
+    let successCount = 0;
+    let failureCount = 0;
+
+    for (const operation of operations) {
+      try {
+        switch (operation.type) {
+          case 'add': {
+            const addResult = await this.addDocuments(operation.data as unknown as DocumentInput);
+            results.push({ type: 'add', success: addResult.success, error: addResult.error });
+            if (addResult.success) successCount++;
+            else failureCount++;
+            break;
+          }
+          case 'update': {
+            const updateResult = await this.updateDocuments(operation.data as unknown as DocumentInput);
+            results.push({ type: 'update', success: updateResult.success, error: updateResult.error });
+            if (updateResult.success) successCount++;
+            else failureCount++;
+            break;
+          }
+          case 'delete': {
+            const ids = (operation.data as Record<string, unknown>).ids as string[];
+            const deleteResult = await this.deleteDocuments(ids);
+            results.push({ type: 'delete', success: deleteResult.success, error: deleteResult.error });
+            if (deleteResult.success) successCount++;
+            else failureCount++;
+            break;
+          }
+          default:
+            results.push({ type: operation.type, success: false, error: `Unknown operation type: ${operation.type}` });
+            failureCount++;
+        }
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        results.push({ type: operation.type, success: false, error: message });
+        failureCount++;
+      }
+    }
+
+    console.log(`[VectorStore] Batch operation complete: ${successCount} succeeded, ${failureCount} failed.`);
+
+    return {
+      success: failureCount === 0,
+      metadata: {
+        totalOperations: operations.length,
+        successCount,
+        failureCount,
+        results,
+      },
+    };
+  }
+
+  // ── Client Status ───────────────────────────────────────────────────────────
+
+  getClientStatus(): Record<string, unknown> {
+    return {
+      isInitialized: this.isInitialized,
+      config: {
+        host: this.config.host,
+        port: this.config.port,
+        collectionName: this.config.collectionName,
+        embeddingFunction: this.config.embeddingFunction,
+      },
+      metrics: { ...this.metrics },
+      cache: {
+        entries: this.queryCache.size,
+        maxAge: this.cacheMaxAge,
+      },
+      hasClient: this.client !== null,
+      hasCollection: this.collection !== null,
+    };
+  }
+
+  // ── Cleanup ─────────────────────────────────────────────────────────────────
+
+  async close(): Promise<void> {
+    console.log('[VectorStore] Closing vector store connection...');
+
+    this.queryCache.clear();
+
+    this.client = null;
+    this.collection = null;
+    this.isInitialized = false;
+
+    console.log('[VectorStore] Vector store connection closed.');
+  }
+
+  // ── Utility Accessors ──────────────────────────────────────────────────────
+
+  getMetrics(): VectorStoreMetrics {
+    return { ...this.metrics };
+  }
+
+  isReady(): boolean {
+    return this.isInitialized;
+  }
+
+  getCollectionName(): string {
+    return this.config.collectionName;
+  }
+
+  clearCache(): void {
+    this.queryCache.clear();
+    console.log('[VectorStore] Query cache cleared.');
+  }
+}
+
+export default ChromaVectorStore;