@retrivora-ai/rag-engine 0.1.6 → 0.1.8

package/src/core/ProviderHealthCheck.ts

@@ -3,6 +3,8 @@
  *
  * Performs connectivity tests, credential validation, and capability checks
  * before initializing providers.
+ *
+ * Note: Option keys used here mirror each Provider constructor's expected keys.
  */
 
 import { VectorDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig';
@@ -20,16 +22,15 @@ export interface HealthCheckResult {
 export class ProviderHealthCheck {
   /**
    * Validates vector database configuration before initialization.
-   * Performs connectivity checks and verifies required resources exist.
    */
   static async checkVectorProvider(config: VectorDBConfig): Promise<HealthCheckResult> {
     const timestamp = Date.now();
 
-    // First validate configuration
+    // Validate configuration first
     const configErrors = ConfigValidator.validate({
       projectId: 'health-check',
       vectorDb: config,
-      llm: { provider: 'openai', model: 'gpt-4o' }, // dummy
+      llm: { provider: 'openai', model: 'gpt-4o' },
       embedding: { provider: 'openai', model: 'text-embedding-3-small' },
     });
 
@@ -85,15 +86,12 @@ export class ProviderHealthCheck {
 
   private static async checkPinecone(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
-    
     try {
-      // Dynamic import to avoid requiring all dependencies
       const { Pinecone } = await import('@pinecone-database/pinecone');
       const client = new Pinecone({ apiKey: opts.apiKey });
-      
       const indexes = await client.listIndexes();
       const indexNames = indexes.indexes?.map((i) => i.name) ?? [];
-      
+
       if (!indexNames.includes(config.indexName)) {
         return {
           healthy: false,
@@ -102,14 +100,10 @@ export class ProviderHealthCheck {
           timestamp,
         };
       }
-
       return {
         healthy: true,
         provider: 'pinecone',
-        capabilities: {
-          indexes: indexNames.length,
-          targetIndex: config.indexName,
-        },
+        capabilities: { indexes: indexNames.length, targetIndex: config.indexName },
         timestamp,
       };
     } catch (error) {
@@ -124,24 +118,15 @@ export class ProviderHealthCheck {
 
   private static async checkPostgres(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
-
     try {
       const { Client } = await import('pg');
       const client = new Client({ connectionString: opts.connectionString });
-      
       await client.connect();
-      
-      // Check if pgvector extension is installed
       const result = await client.query(`
-        SELECT EXISTS(
-          SELECT 1 FROM pg_extension WHERE extname = 'vector'
-        );
+        SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector');
       `);
-      
       const hasVector = result.rows[0].exists;
-      
       await client.end();
-
       return {
         healthy: true,
         provider: 'postgresql',
@@ -160,27 +145,21 @@ export class ProviderHealthCheck {
 
   private static async checkMongoDB(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
-
     try {
       const { MongoClient } = await import('mongodb');
       const client = new MongoClient(opts.uri);
-      
       await client.connect();
-      
       const db = client.db(opts.database);
       const collections = await db.listCollections().toArray();
       const collectionNames = collections.map((c) => c.name);
-      
       const hasCollection = collectionNames.includes(opts.collection);
-      
       await client.close();
-
       return {
         healthy: true,
         provider: 'mongodb',
         capabilities: {
           collections: collectionNames.length,
-          targetCollection: hasCollection ? opts.collection : 'NOT FOUND',
+          targetCollection: hasCollection ? opts.collection : 'NOT FOUND (will be created on first insert)',
         },
         timestamp,
       };
@@ -194,30 +173,31 @@ export class ProviderHealthCheck {
     }
   }
 
+  /**
+   * Milvus health check — uses baseUrl/uri (matching MilvusProvider constructor).
+   */
   private static async checkMilvus(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, unknown>;
-    const host = (opts.host as string) || 'localhost';
-    const port = (opts.port as number) || 19530;
+
+    // Resolve endpoint — mirrors MilvusProvider constructor logic
+    const baseUrl =
+      (opts.baseUrl as string) ||
+      (opts.uri as string) ||
+      (opts.host ? `http://${opts.host}:${(opts.port as number) || 19530}` : 'http://localhost:19530');
 
     try {
-      // Try HTTP health check endpoint
       const controller = new AbortController();
       const timeoutId = setTimeout(() => controller.abort(), 5000);
-      
-      const response = await fetch(`http://${host}:${port}/healthz`, {
-        signal: controller.signal,
-      });
-      
+      const response = await fetch(`${baseUrl}/v1/health`, { signal: controller.signal });
       clearTimeout(timeoutId);
 
       if (!response.ok) {
         throw new Error(`Health check returned ${response.status}`);
       }
-
       return {
         healthy: true,
         provider: 'milvus',
-        capabilities: { endpoint: `${host}:${port}` },
+        capabilities: { endpoint: baseUrl },
         timestamp,
       };
     } catch (error) {
@@ -230,19 +210,19 @@ export class ProviderHealthCheck {
     }
   }
 
+  /**
+   * Qdrant health check — uses baseUrl (matching QdrantProvider constructor).
+   */
   private static async checkQdrant(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, unknown>;
-    const baseUrl = (opts.baseUrl as string) || (opts.url as string) || 'http://localhost:6333';
-
+    const baseUrl = ((opts.baseUrl as string) || (opts.url as string) || 'http://localhost:6333').replace(/\/$/, '');
     try {
-      const response = await fetch(`${baseUrl}/health`);
-      
-      if (!response.ok) {
-        throw new Error(`Health check returned ${response.status}`);
-      }
-
+      const apiKey = opts.apiKey as string | undefined;
+      const response = await fetch(`${baseUrl}/`, {
+        headers: apiKey ? { 'api-key': apiKey } : {},
+      });
+      if (!response.ok) throw new Error(`Health check returned ${response.status}`);
       const health = await response.json();
-
       return {
         healthy: true,
         provider: 'qdrant',
@@ -259,22 +239,24 @@ export class ProviderHealthCheck {
     }
   }
 
+  /**
+   * ChromaDB health check — uses baseUrl/host (matching ChromaDBProvider constructor).
+   */
   private static async checkChromaDB(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, unknown>;
-    const host = (opts.host as string) || 'localhost';
-    const port = (opts.port as number) || 8000;
 
-    try {
-      const response = await fetch(`http://${host}:${port}/api/v1/heartbeat`);
-      
-      if (!response.ok) {
-        throw new Error(`Health check returned ${response.status}`);
-      }
+    // Mirrors ChromaDBProvider constructor
+    const baseUrl =
+      (opts.baseUrl as string) ||
+      (opts.host ? `http://${opts.host}:${(opts.port as number) || 8000}` : 'http://localhost:8000');
 
+    try {
+      const response = await fetch(`${baseUrl}/api/v1/heartbeat`);
+      if (!response.ok) throw new Error(`Health check returned ${response.status}`);
       return {
         healthy: true,
         provider: 'chromadb',
-        capabilities: { endpoint: `${host}:${port}` },
+        capabilities: { endpoint: baseUrl },
         timestamp,
       };
     } catch (error) {
@@ -287,47 +269,49 @@ export class ProviderHealthCheck {
     }
   }
 
+  /**
+   * Redis health check — Redis is TCP-only (no HTTP endpoint).
+   * We report healthy=true with a note; actual connectivity is validated at first operation.
+   */
   private static async checkRedis(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
-    const url = opts.url || `redis://${opts.host}:${opts.port}`;
-
-    try {
-      // Try to connect via simple TCP check using fetch
-      // Note: Full Redis connection would require redis package as optional dependency
-      await fetch(url);
-      
-      return {
-        healthy: true,
-        provider: 'redis',
-        capabilities: { endpoint: url },
-        timestamp,
-      };
-    } catch (error) {
-      // If fetch fails, try to gracefully handle (redis may not have HTTP interface)
-      return {
-        healthy: false,
-        provider: 'redis',
-        error: `Connection check failed: ${error instanceof Error ? error.message : String(error)}. Ensure Redis is running at ${url}`,
-        timestamp,
-      };
+    const url =
+      opts.baseUrl || opts.url || (opts.host && opts.port ? `redis://${opts.host}:${opts.port}` : 'redis://localhost:6379');
+
+    // For Upstash REST endpoints, try an HTTP ping
+    if (url.startsWith('http')) {
+      try {
+        await fetch(url);
+        return { healthy: true, provider: 'redis', capabilities: { endpoint: url }, timestamp };
+      } catch {
+        // Ignore — Upstash may reject GET but accept POST
+      }
     }
+
+    return {
+      healthy: true,
+      provider: 'redis',
+      capabilities: {
+        endpoint: url,
+        note: 'Redis uses TCP; connectivity is validated on first operation.',
+      },
+      timestamp,
+    };
   }
 
+  /**
+   * Weaviate health check — uses baseUrl/url (matching WeaviateProvider constructor).
+   */
   private static async checkWeaviate(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
-    const url = (opts.url || 'http://localhost:8080').replace(/\/$/, '');
-
+    const baseUrl = ((opts.baseUrl || opts.url) || 'http://localhost:8080').replace(/\/$/, '');
     try {
-      const response = await fetch(`${url}/.well-known/ready`);
-      
-      if (!response.ok) {
-        throw new Error(`Health check returned ${response.status}`);
-      }
-
+      const response = await fetch(`${baseUrl}/v1/.well-known/ready`);
+      if (!response.ok) throw new Error(`Health check returned ${response.status}`);
       return {
         healthy: true,
         provider: 'weaviate',
-        capabilities: { endpoint: url },
+        capabilities: { endpoint: baseUrl },
         timestamp,
       };
     } catch (error) {
@@ -343,21 +327,13 @@ export class ProviderHealthCheck {
   private static async checkRestAPI(config: VectorDBConfig, timestamp: number): Promise<HealthCheckResult> {
     const opts = config.options as Record<string, string>;
     const baseUrl = (opts.baseUrl || '').replace(/\/$/, '');
-
     if (!baseUrl) {
-      return {
-        healthy: false,
-        provider: 'rest',
-        error: 'baseUrl is required',
-        timestamp,
-      };
+      return { healthy: false, provider: 'rest', error: 'baseUrl is required', timestamp };
     }
-
     try {
       const response = await fetch(`${baseUrl}/health`, {
         headers: opts.headers ? JSON.parse(opts.headers) : {},
       });
-
       return {
         healthy: response.ok,
         provider: 'rest',
@@ -379,7 +355,6 @@ export class ProviderHealthCheck {
    */
   static async checkLLMProvider(config: LLMConfig): Promise<HealthCheckResult> {
     const timestamp = Date.now();
-
     try {
       switch (config.provider) {
         case 'openai':
@@ -413,20 +388,12 @@ export class ProviderHealthCheck {
     try {
       const OpenAI = await import('openai');
       const client = new OpenAI.default({ apiKey: config.apiKey });
-
-      // Test models endpoint
       const models = await client.models.list();
-      
       const hasModel = models.data.some((m) => m.id === config.model);
-
       return {
         healthy: true,
         provider: 'openai',
-        capabilities: {
-          model: config.model,
-          available: hasModel,
-          totalModels: models.data.length,
-        },
+        capabilities: { model: config.model, available: hasModel, totalModels: models.data.length },
         timestamp,
       };
     } catch (error) {
@@ -443,20 +410,12 @@ export class ProviderHealthCheck {
     try {
       const { default: Anthropic } = await import('@anthropic-ai/sdk');
       const client = new Anthropic({ apiKey: config.apiKey });
-
-      // Try a minimal API call to verify credentials
       await client.messages.create({
         model: config.model,
         max_tokens: 10,
         messages: [{ role: 'user', content: 'ping' }],
       });
-
-      return {
-        healthy: true,
-        provider: 'anthropic',
-        capabilities: { model: config.model },
-        timestamp,
-      };
+      return { healthy: true, provider: 'anthropic', capabilities: { model: config.model }, timestamp };
     } catch (error) {
       return {
         healthy: false,
@@ -469,26 +428,16 @@ export class ProviderHealthCheck {
 
   private static async checkOllama(config: LLMConfig, timestamp: number): Promise<HealthCheckResult> {
     const baseUrl = (config.baseUrl || 'http://localhost:11434').replace(/\/$/, '');
-
     try {
       const response = await fetch(`${baseUrl}/api/tags`);
-      
-      if (!response.ok) {
-        throw new Error(`Health check returned ${response.status}`);
-      }
-
+      if (!response.ok) throw new Error(`Health check returned ${response.status}`);
       const data = await response.json() as Record<string, unknown>;
       const models = (data.models as Array<{ name: string }>) || [];
       const hasModel = models.some((m) => m.name === config.model);
-
       return {
         healthy: true,
         provider: 'ollama',
-        capabilities: {
-          model: config.model,
-          available: hasModel,
-          totalModels: models.length,
-        },
+        capabilities: { model: config.model, available: hasModel, totalModels: models.length },
         timestamp,
       };
     } catch (error) {
@@ -504,22 +453,12 @@ export class ProviderHealthCheck {
   private static async checkLLMRestAPI(config: LLMConfig, timestamp: number): Promise<HealthCheckResult> {
     const baseUrlRaw = config.baseUrl || (config.options as Record<string, unknown>)?.baseUrl;
     const baseUrl = (typeof baseUrlRaw === 'string' ? baseUrlRaw : '').replace(/\/$/, '');
-
     if (!baseUrl) {
-      return {
-        healthy: false,
-        provider: config.provider,
-        error: 'baseUrl is required',
-        timestamp,
-      };
+      return { healthy: false, provider: config.provider, error: 'baseUrl is required', timestamp };
     }
-
     try {
       const headers = (config.options as Record<string, unknown>)?.headers as Record<string, string> | undefined;
-      const response = await fetch(`${baseUrl}/health`, {
-        headers: headers || undefined,
-      });
-
+      const response = await fetch(`${baseUrl}/health`, { headers: headers || undefined });
       return {
         healthy: response.ok,
         provider: config.provider,
@@ -537,7 +476,7 @@ export class ProviderHealthCheck {
   }
 
   /**
-   * Runs comprehensive health checks on all configured providers.
+   * Runs comprehensive health checks on all configured providers in parallel.
    */
   static async checkAll(
     vectorDbConfig: VectorDBConfig,

package/src/core/ProviderRegistry.ts

@@ -3,17 +3,22 @@ import { VectorDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig'
 import { LLMFactory } from '../llm/LLMFactory';
 import { ILLMProvider } from '../llm/ILLMProvider';
 
-// We will import these as we create them
-// import { PineconeProvider } from '../providers/vectordb/PineconeProvider';
-
 /**
  * ProviderRegistry — dynamic provider loader for Vector DBs and LLMs.
+ *
+ * Supports:
+ * - Built-in providers loaded on-demand via dynamic imports (tree-shaking friendly)
+ * - Custom provider registration via registerVectorProvider()
  */
 export class ProviderRegistry {
   private static vectorProviders: Record<string, new (config: VectorDBConfig) => BaseVectorProvider> = {};
 
   /**
-   * Register a custom vector provider.
+   * Register a custom vector provider class by name.
+   * The name must match the provider value used in VectorDBConfig.provider.
+   *
+   * @example
+   * ProviderRegistry.registerVectorProvider('my-db', MyCustomProvider);
    */
   static registerVectorProvider(name: string, providerClass: new (config: VectorDBConfig) => BaseVectorProvider) {
     this.vectorProviders[name] = providerClass;
@@ -21,44 +26,62 @@ export class ProviderRegistry {
 
   /**
    * Creates a vector database provider based on the configuration.
+   * Built-in providers are dynamically imported to avoid bundling all SDKs.
    */
   static async createVectorProvider(config: VectorDBConfig): Promise<BaseVectorProvider> {
     const { provider } = config;
 
-    // Check if we have a registered provider (dynamic loading)
+    // Custom registered provider takes priority
     if (this.vectorProviders[provider]) {
       return new this.vectorProviders[provider](config);
     }
 
-    // Fallback to internal switch for built-in providers
+    // Built-in providers — lazy-loaded
     switch (provider) {
-      case 'pinecone':
+      case 'pinecone': {
         const { PineconeProvider } = await import('../providers/vectordb/PineconeProvider');
         return new PineconeProvider(config);
+      }
       case 'pgvector':
-      case 'postgresql':
+      case 'postgresql': {
         const { PostgreSQLProvider } = await import('../providers/vectordb/PostgreSQLProvider');
         return new PostgreSQLProvider(config);
-      case 'mongodb':
+      }
+      case 'mongodb': {
         const { MongoDBProvider } = await import('../providers/vectordb/MongoDBProvider');
         return new MongoDBProvider(config);
-      case 'milvus':
+      }
+      case 'milvus': {
         const { MilvusProvider } = await import('../providers/vectordb/MilvusProvider');
         return new MilvusProvider(config);
-      case 'qdrant':
+      }
+      case 'qdrant': {
         const { QdrantProvider } = await import('../providers/vectordb/QdrantProvider');
         return new QdrantProvider(config);
-      case 'chromadb':
+      }
+      case 'chromadb': {
         const { ChromaDBProvider } = await import('../providers/vectordb/ChromaDBProvider');
         return new ChromaDBProvider(config);
-      case 'redis':
+      }
+      case 'redis': {
         const { RedisProvider } = await import('../providers/vectordb/RedisProvider');
         return new RedisProvider(config);
-      case 'weaviate':
+      }
+      case 'weaviate': {
         const { WeaviateProvider } = await import('../providers/vectordb/WeaviateProvider');
         return new WeaviateProvider(config);
+      }
+      case 'universal_rest':
+      case 'rest': {
+        const { UniversalVectorProvider } = await import('../providers/vectordb/UniversalVectorProvider');
+        return new UniversalVectorProvider(config);
+      }
       default:
-        throw new Error(`[ProviderRegistry] Unsupported vector provider: ${provider}`);
+        throw new Error(
+          `[ProviderRegistry] Unsupported vector provider: "${provider}". ` +
+          `Built-in providers: pinecone | pgvector | postgresql | mongodb | milvus | qdrant | chromadb | redis | weaviate | universal_rest. ` +
+          `For custom providers, call ProviderRegistry.registerVectorProvider("${provider}", YourClass).`
+        );
     }
   }
 

package/src/providers/vectordb/ChromaDBProvider.ts

@@ -4,7 +4,10 @@ import { BaseVectorProvider } from './BaseVectorProvider';
 import { VectorMatch, UpsertDocument } from '../../types';
 
 /**
- * ChromaDBProvider — implementation for ChromaDB using its REST API.
+ * ChromaDBProvider — ChromaDB implementation using its REST API.
+ *
+ * Required options: { baseUrl: string } — e.g. "http://localhost:8000"
+ * Optional: { host?: string, port?: number } — used to construct baseUrl when baseUrl not set
  */
 export class ChromaDBProvider extends BaseVectorProvider {
   private http: AxiosInstance;
@@ -13,8 +16,13 @@ export class ChromaDBProvider extends BaseVectorProvider {
   constructor(config: VectorDBConfig) {
     super(config);
     const opts = config.options as Record<string, unknown>;
-    const baseUrl = opts.baseUrl as string;
-    if (!baseUrl) throw new Error('[ChromaDBProvider] baseUrl is required');
+
+    // Accept either explicit baseUrl or host+port
+    const baseUrl =
+      (opts.baseUrl as string) ||
+      (opts.host ? `http://${opts.host}:${(opts.port as number) || 8000}` : undefined);
+
+    if (!baseUrl) throw new Error('[ChromaDBProvider] options.baseUrl is required');
 
     this.http = axios.create({
       baseURL: baseUrl,
@@ -22,10 +30,27 @@ export class ChromaDBProvider extends BaseVectorProvider {
     });
   }
 
+  /**
+   * Get or create the ChromaDB collection.
+   */
   async initialize(): Promise<void> {
-    // Get or create collection
-    const { data } = await this.http.get(`/api/v1/collections/${this.indexName}`);
-    this.collectionId = data.id;
+    try {
+      const { data } = await this.http.get(`/api/v1/collections/${this.indexName}`);
+      this.collectionId = data.id;
+      console.log(`[ChromaDBProvider] ✅ Collection "${this.indexName}" found (id: ${this.collectionId})`);
+    } catch (err) {
+      if (axios.isAxiosError(err) && err.response?.status === 404) {
+        // Create the collection if it doesn't exist
+        console.log(`[ChromaDBProvider] ⏳ Collection "${this.indexName}" not found. Creating...`);
+        const { data } = await this.http.post('/api/v1/collections', {
+          name: this.indexName,
+        });
+        this.collectionId = data.id;
+        console.log(`[ChromaDBProvider] ✅ Created collection "${this.indexName}" (id: ${this.collectionId})`);
+      } else {
+        throw err;
+      }
+    }
   }
 
   async upsert(doc: UpsertDocument, namespace?: string): Promise<void> {
@@ -34,7 +59,7 @@ export class ChromaDBProvider extends BaseVectorProvider {
 
   async batchUpsert(docs: UpsertDocument[], namespace?: string): Promise<void> {
     const payload = {
-      ids: docs.map(d => d.id),
+      ids: docs.map(d => String(d.id)),
       embeddings: docs.map(d => d.vector),
       documents: docs.map(d => d.content),
       metadatas: docs.map(d => ({
@@ -53,7 +78,7 @@ export class ChromaDBProvider extends BaseVectorProvider {
       where: namespace ? { namespace: { $eq: namespace } } : undefined,
     };
     const { data } = await this.http.post(`/api/v1/collections/${this.collectionId}/query`, payload);
-    
+
     const matches: VectorMatch[] = [];
     if (data.ids && data.ids[0]) {
       for (let i = 0; i < data.ids[0].length; i++) {
@@ -68,16 +93,16 @@ export class ChromaDBProvider extends BaseVectorProvider {
     return matches;
   }
 
-  async delete(id: string, _namespace?: string): Promise<void> {
+  async delete(id: string, namespace?: string): Promise<void> {
     await this.http.post(`/api/v1/collections/${this.collectionId}/delete`, {
       ids: [id],
-      where: _namespace ? { namespace: { $eq: _namespace } } : undefined,
+      where: namespace ? { namespace: { $eq: namespace } } : undefined,
     });
   }
 
-  async deleteNamespace(_namespace: string): Promise<void> {
+  async deleteNamespace(namespace: string): Promise<void> {
     await this.http.post(`/api/v1/collections/${this.collectionId}/delete`, {
-      where: { namespace: { $eq: _namespace } },
+      where: { namespace: { $eq: namespace } },
     });
   }