@retrivora-ai/rag-engine 1.9.6 → 1.9.8

package/src/core/DatabaseStorage.ts

@@ -0,0 +1,469 @@
+import * as fs from 'fs';
+import * as path from 'path';
+import { RagConfig } from '../config/RagConfig';
+import { LicenseVerifier } from './LicenseVerifier';
+
+export interface HistoryMessage {
+  id: string;
+  role: string;
+  content: string;
+  sources?: any;
+  uiTransformation?: any;
+  trace?: any;
+  createdAt?: string;
+}
+
+export interface FeedbackItem {
+  id: string;
+  messageId: string;
+  sessionId: string;
+  rating: string;
+  comment?: string;
+  createdAt?: string;
+}
+
+export class DatabaseStorage {
+  private config: RagConfig;
+  private provider: string;
+  // Dynamic PG Pool
+  private pgPool: any = null;
+  // Dynamic MongoDB Client
+  private mongoClient: any = null;
+  private mongoDbName = '';
+
+  private historyTableName: string;
+  private feedbackTableName: string;
+
+  // Filesystem fallback configuration
+  private fallbackDir = path.join(process.cwd(), '.retrivora');
+  private historyFile = path.join(this.fallbackDir, 'history.json');
+  private feedbackFile = path.join(this.fallbackDir, 'feedback.json');
+
+  constructor(config: RagConfig) {
+    this.config = config;
+    this.provider = config.vectorDb?.provider || 'fs';
+
+    // Retrieve and sanitize table/collection names (alphanumeric + underscores only)
+    const rawHistoryTable = config.rag?.history?.tableName || 'retrivora_history';
+    const rawFeedbackTable = config.rag?.feedback?.tableName || 'retrivora_feedback';
+    this.historyTableName = rawHistoryTable.replace(/[^a-zA-Z0-9_]/g, '_');
+    this.feedbackTableName = rawFeedbackTable.replace(/[^a-zA-Z0-9_]/g, '_');
+  }
+
+  /**
+   * Asserts the user is running a valid Premium/Enterprise license in production.
+   * In non-production (development / test) environments this is a complete no-op —
+   * all History and Feedback features are freely accessible for local development.
+   */
+  private checkPremiumSubscription() {
+    // ── Development / test: fully open, no checks ────────────────────────────
+    if (process.env.NODE_ENV !== 'production') {
+      return;
+    }
+
+    // ── Production: license key required ────────────────────────────────────
+    if (!this.config.licenseKey) {
+      throw new Error(
+        '[Retrivora SDK] Chat History and Feedback features require a Premium or Enterprise license key in production.'
+      );
+    }
+
+    try {
+      const payload = LicenseVerifier.verify(this.config.licenseKey, this.config.projectId);
+      const tier = (payload.tier || '').toLowerCase();
+      if (
+        tier !== 'premium' &&
+        tier !== 'enterprise' &&
+        tier !== 'growth' &&
+        tier !== 'pro' &&
+        tier !== 'developer_pro'
+      ) {
+        throw new Error(
+          `[Retrivora SDK] Chat History and Feedback features are not available on the "${tier}" tier. ` +
+          `Please upgrade to a Developer Pro or Enterprise plan.`
+        );
+      }
+    } catch (err) {
+      throw new Error(
+        `[Retrivora SDK] Subscription check failed: ${err instanceof Error ? err.message : String(err)}`
+      );
+    }
+  }
+
+  private checkHistoryEnabled() {
+    this.checkPremiumSubscription();
+    if (this.config.rag?.history?.enabled === false) {
+      throw new Error('[Retrivora SDK] Chat History is disabled in RagConfig.');
+    }
+  }
+
+  private checkFeedbackEnabled() {
+    this.checkPremiumSubscription();
+    if (this.config.rag?.feedback?.enabled === false) {
+      throw new Error('[Retrivora SDK] User Feedback is disabled in RagConfig.');
+    }
+  }
+
+  private async getPGPool() {
+    const globalKey = `__retrivora_pg_pool_${this.config.projectId}`;
+    const _g = global as any;
+    if (_g[globalKey]) {
+      this.pgPool = _g[globalKey];
+      return this.pgPool;
+    }
+    const opts = this.config.vectorDb.options as Record<string, any>;
+    if (!opts.connectionString) {
+      throw new Error('[DatabaseStorage] pg connectionString option is missing.');
+    }
+    const { Pool } = await import('pg');
+    this.pgPool = new Pool({ connectionString: opts.connectionString });
+    _g[globalKey] = this.pgPool;
+
+    // Initialize tables using dynamic sanitized table names
+    const client = await this.pgPool.connect();
+    try {
+      await client.query(`
+        CREATE TABLE IF NOT EXISTS ${this.historyTableName} (
+          id                TEXT PRIMARY KEY,
+          session_id        TEXT NOT NULL,
+          role              TEXT NOT NULL,
+          content           TEXT NOT NULL,
+          sources           JSONB,
+          ui_transformation JSONB,
+          trace             JSONB,
+          created_at        TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+        )
+      `);
+      await client.query(`
+        CREATE TABLE IF NOT EXISTS ${this.feedbackTableName} (
+          id          TEXT PRIMARY KEY,
+          message_id  TEXT NOT NULL UNIQUE,
+          session_id  TEXT NOT NULL,
+          rating      TEXT NOT NULL,
+          comment     TEXT,
+          created_at  TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+        )
+      `);
+    } finally {
+      client.release();
+    }
+    return this.pgPool;
+  }
+
+  private async getMongoClient() {
+    const globalKey = `__retrivora_mongo_client_${this.config.projectId}`;
+    const _g = global as any;
+    if (_g[globalKey]) {
+      this.mongoClient = _g[globalKey];
+      this.mongoDbName = (this.config.vectorDb.options as any).database;
+      return this.mongoClient;
+    }
+    const opts = this.config.vectorDb.options as Record<string, any>;
+    if (!opts.uri || !opts.database) {
+      throw new Error('[DatabaseStorage] mongo uri and database options are missing.');
+    }
+    const { MongoClient } = await import('mongodb');
+    this.mongoClient = new MongoClient(opts.uri);
+    await this.mongoClient.connect();
+    _g[globalKey] = this.mongoClient;
+    this.mongoDbName = opts.database;
+    return this.mongoClient;
+  }
+
+  private getMongoDb() {
+    return this.mongoClient.db(this.mongoDbName);
+  }
+
+  // Helper for FS fallback reads
+  private readLocalFile(filePath: string): any[] {
+    if (!fs.existsSync(filePath)) return [];
+    try {
+      const raw = fs.readFileSync(filePath, 'utf-8');
+      return JSON.parse(raw) || [];
+    } catch {
+      return [];
+    }
+  }
+
+  // Helper for FS fallback writes
+  private writeLocalFile(filePath: string, data: any[]): void {
+    if (!fs.existsSync(this.fallbackDir)) {
+      fs.mkdirSync(this.fallbackDir, { recursive: true });
+    }
+    fs.writeFileSync(filePath, JSON.stringify(data, null, 2), 'utf-8');
+  }
+
+  // ─── Message History Operations ──────────────────────────────────────────
+
+  async saveMessage(
+    sessionId: string,
+    message: Omit<HistoryMessage, 'createdAt'>
+  ): Promise<void> {
+    this.checkHistoryEnabled();
+    const createdAt = new Date().toISOString();
+    const msgId = message.id || `msg_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`;
+
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      await pool.query(
+        `INSERT INTO ${this.historyTableName} (id, session_id, role, content, sources, ui_transformation, trace, created_at)
+         VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+         ON CONFLICT (id) DO UPDATE 
+         SET content = EXCLUDED.content, sources = EXCLUDED.sources, ui_transformation = EXCLUDED.ui_transformation, trace = EXCLUDED.trace`,
+        [
+          msgId,
+          sessionId,
+          message.role,
+          message.content,
+          message.sources ? JSON.stringify(message.sources) : null,
+          message.uiTransformation ? JSON.stringify(message.uiTransformation) : null,
+          message.trace ? JSON.stringify(message.trace) : null,
+          createdAt,
+        ]
+      );
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      const col = db.collection(this.historyTableName);
+      await col.updateOne(
+        { id: msgId },
+        {
+          $set: {
+            id: msgId,
+            session_id: sessionId,
+            role: message.role,
+            content: message.content,
+            sources: message.sources || null,
+            ui_transformation: message.uiTransformation || null,
+            trace: message.trace || null,
+            created_at: createdAt,
+          },
+        },
+        { upsert: true }
+      );
+    } else {
+      // Local FS fallback
+      const history = this.readLocalFile(this.historyFile);
+      const index = history.findIndex((h) => h.id === msgId);
+      const item = {
+        id: msgId,
+        session_id: sessionId,
+        role: message.role,
+        content: message.content,
+        sources: message.sources || null,
+        ui_transformation: message.uiTransformation || null,
+        trace: message.trace || null,
+        created_at: createdAt,
+      };
+
+      if (index !== -1) {
+        history[index] = item;
+      } else {
+        history.push(item);
+      }
+      this.writeLocalFile(this.historyFile, history);
+    }
+  }
+
+  async getHistory(sessionId: string): Promise<HistoryMessage[]> {
+    this.checkHistoryEnabled();
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      const res = await pool.query(
+        `SELECT id, role, content, sources, ui_transformation as "uiTransformation", trace, created_at as "createdAt"
+         FROM ${this.historyTableName} 
+         WHERE session_id = $1 
+         ORDER BY created_at ASC`,
+        [sessionId]
+      );
+      return res.rows;
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      const col = db.collection(this.historyTableName);
+      const items = await col
+        .find({ session_id: sessionId })
+        .sort({ created_at: 1 })
+        .toArray();
+      return items.map((item: any) => ({
+        id: item.id,
+        role: item.role,
+        content: item.content,
+        sources: item.sources,
+        uiTransformation: item.ui_transformation,
+        trace: item.trace,
+        createdAt: item.created_at,
+      }));
+    } else {
+      const history = this.readLocalFile(this.historyFile);
+      return history
+        .filter((h) => h.session_id === sessionId)
+        .sort((a, b) => new Date(a.created_at).getTime() - new Date(b.created_at).getTime());
+    }
+  }
+
+  async clearHistory(sessionId: string): Promise<void> {
+    this.checkHistoryEnabled();
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      await pool.query(`DELETE FROM ${this.historyTableName} WHERE session_id = $1`, [sessionId]);
+      await pool.query(`DELETE FROM ${this.feedbackTableName} WHERE session_id = $1`, [sessionId]);
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      await db.collection(this.historyTableName).deleteMany({ session_id: sessionId });
+      await db.collection(this.feedbackTableName).deleteMany({ session_id: sessionId });
+    } else {
+      const history = this.readLocalFile(this.historyFile);
+      const filteredHistory = history.filter((h) => h.session_id !== sessionId);
+      this.writeLocalFile(this.historyFile, filteredHistory);
+
+      const feedback = this.readLocalFile(this.feedbackFile);
+      const filteredFeedback = feedback.filter((f) => f.session_id !== sessionId);
+      this.writeLocalFile(this.feedbackFile, filteredFeedback);
+    }
+  }
+
+  async listSessions(): Promise<string[]> {
+    this.checkHistoryEnabled();
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      const res = await pool.query(`SELECT DISTINCT session_id FROM ${this.historyTableName}`);
+      return res.rows.map((r: any) => r.session_id);
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      return db.collection(this.historyTableName).distinct('session_id');
+    } else {
+      const history = this.readLocalFile(this.historyFile);
+      const sessions = new Set<string>(history.map((h) => h.session_id));
+      return Array.from(sessions);
+    }
+  }
+
+  // ─── Feedback Operations ──────────────────────────────────────────────────
+
+  async saveFeedback(feedback: Omit<FeedbackItem, 'id' | 'createdAt'>): Promise<void> {
+    this.checkFeedbackEnabled();
+    const id = `fb_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`;
+    const createdAt = new Date().toISOString();
+
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      await pool.query(
+        `INSERT INTO ${this.feedbackTableName} (id, message_id, session_id, rating, comment, created_at)
+         VALUES ($1, $2, $3, $4, $5, $6)
+         ON CONFLICT (message_id) DO UPDATE 
+         SET rating = EXCLUDED.rating, comment = EXCLUDED.comment`,
+        [id, feedback.messageId, feedback.sessionId, feedback.rating, feedback.comment || null, createdAt]
+      );
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      const col = db.collection(this.feedbackTableName);
+      await col.updateOne(
+        { message_id: feedback.messageId },
+        {
+          $set: {
+            id,
+            message_id: feedback.messageId,
+            session_id: feedback.sessionId,
+            rating: feedback.rating,
+            comment: feedback.comment || null,
+            created_at: createdAt,
+          },
+        },
+        { upsert: true }
+      );
+    } else {
+      const list = this.readLocalFile(this.feedbackFile);
+      const index = list.findIndex((f) => f.message_id === feedback.messageId);
+      const item = {
+        id,
+        message_id: feedback.messageId,
+        session_id: feedback.sessionId,
+        rating: feedback.rating,
+        comment: feedback.comment || null,
+        created_at: createdAt,
+      };
+
+      if (index !== -1) {
+        list[index] = item;
+      } else {
+        list.push(item);
+      }
+      this.writeLocalFile(this.feedbackFile, list);
+    }
+  }
+
+  async getFeedback(messageId: string): Promise<FeedbackItem | null> {
+    this.checkFeedbackEnabled();
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      const res = await pool.query(
+        `SELECT id, message_id as "messageId", session_id as "sessionId", rating, comment, created_at as "createdAt"
+         FROM ${this.feedbackTableName} 
+         WHERE message_id = $1`,
+        [messageId]
+      );
+      return res.rows[0] || null;
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      const col = db.collection(this.feedbackTableName);
+      const item = await col.findOne({ message_id: messageId });
+      if (!item) return null;
+      return {
+        id: item.id,
+        messageId: item.message_id,
+        sessionId: item.session_id,
+        rating: item.rating,
+        comment: item.comment,
+        createdAt: item.created_at,
+      };
+    } else {
+      const list = this.readLocalFile(this.feedbackFile);
+      return list.find((f) => f.message_id === messageId) || null;
+    }
+  }
+
+  async listFeedback(): Promise<FeedbackItem[]> {
+    this.checkFeedbackEnabled();
+    if (this.provider === 'postgresql' || this.provider === 'pgvector') {
+      const pool = await this.getPGPool();
+      const res = await pool.query(
+        `SELECT id, message_id as "messageId", session_id as "sessionId", rating, comment, created_at as "createdAt"
+         FROM ${this.feedbackTableName} 
+         ORDER BY created_at DESC`
+      );
+      return res.rows;
+    } else if (this.provider === 'mongodb') {
+      await this.getMongoClient();
+      const db = this.getMongoDb();
+      const col = db.collection(this.feedbackTableName);
+      const items = await col.find({}).sort({ created_at: -1 }).toArray();
+      return items.map((item: any) => ({
+        id: item.id,
+        messageId: item.message_id,
+        sessionId: item.session_id,
+        rating: item.rating,
+        comment: item.comment,
+        createdAt: item.created_at,
+      }));
+    } else {
+      const list = this.readLocalFile(this.feedbackFile);
+      return [...list].sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime());
+    }
+  }
+
+  async disconnect() {
+    if (this.pgPool) {
+      await this.pgPool.end();
+      this.pgPool = null;
+    }
+    if (this.mongoClient) {
+      await this.mongoClient.close();
+      this.mongoClient = null;
+    }
+  }
+}

package/src/core/LicenseVerifier.ts

@@ -0,0 +1,154 @@
+import * as crypto from 'crypto';
+import { ConfigurationException } from '../exceptions';
+
+export interface LicensePayload {
+  projectId: string;
+  expiresAt: number; // UNIX timestamp in seconds
+  tier: string;
+}
+
+/**
+ * LicenseVerifier — handles cryptographic validation of signed JWT license keys.
+ * Enables zero-latency local license validation without external network calls.
+ */
+export class LicenseVerifier {
+  // Retrivora's Public Key used to verify the license key signature.
+  // In production, this matches the private key held by Retrivora SaaS.
+  private static readonly PUBLIC_KEY = `-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArMieCkCBtTfByRNOserm
+XM4T0Am29JyNzB4XQPe+WJJ56CN73H1HLXx9n1ByFcxkEJ9po+SURj5DnvUXwEy+
+xstx33wDPpVmcMQe33j5CRYS0S4gICiNqIRN7XkTlJtrcXqrmh1/hdOAwfRbjO1T
+NVtORHwUwWfI/lBLyxUPGtKPed+0fQ7NtcW4o00JXxs3EjCB5BV9CbnXoTjQb/4h
+iwZof4l8rb7oaNzOsVg0RSyxsETX7Ex5sI0CjBz1p2mYp4oVhw/A/h8Ls7H0XzjC
++Eb6BLtsytt5JHoSp0jExLlCpDmpys2L+kETcBVx+IqiXtdN1n6KbkksvEDUVzLI
+MwIDAQAB
+-----END PUBLIC KEY-----`;
+
+  /**
+   * Decodes, verifies signature, and checks metadata of a license key.
+   *
+   * @param licenseKey - Base64url signed JWT license key.
+   * @param currentProjectId - Project namespace ID.
+   * @param publicKeyOverride - Optional override for unit/integration tests.
+   */
+  static verify(
+    licenseKey: string | undefined,
+    currentProjectId: string,
+    provider?: string,
+    publicKeyOverride?: string
+  ): LicensePayload {
+    const isProduction = process.env.NODE_ENV === 'production';
+
+    // 1. Development Fallback (Fail-open locally but log warning)
+    if (!licenseKey) {
+      if (isProduction) {
+        throw new ConfigurationException(
+          '[Retrivora SDK] License key (licenseKey) is required in production environments.'
+        );
+      }
+      console.warn(
+        `\x1b[33m[Retrivora SDK] WARNING: Running without a license key. This namespace (${currentProjectId}) is permitted for local development only.\x1b[0m`
+      );
+      return {
+        projectId: currentProjectId,
+        expiresAt: Math.floor(Date.now() / 1000) + 86400, // Valid for 24h
+        tier: 'hobby',
+      };
+    }
+
+    try {
+      const parts = licenseKey.split('.');
+      if (parts.length !== 3) {
+        throw new Error('Malformed token structure (expected 3 parts).');
+      }
+
+      const [headerB64, payloadB64, signatureB64] = parts;
+      const dataToVerify = `${headerB64}.${payloadB64}`;
+
+      const publicKey = publicKeyOverride ?? this.PUBLIC_KEY;
+      
+      // 2. Cryptographic signature check
+      const signature = Buffer.from(signatureB64, 'base64url');
+      const data = Buffer.from(dataToVerify);
+
+      const isValid = crypto.verify(
+        'sha256',
+        data,
+        publicKey,
+        signature
+      );
+
+      if (!isValid) {
+        throw new Error('Signature verification failed.');
+      }
+
+      // 3. Decode & validate payload checks
+      const payload: LicensePayload = JSON.parse(
+        Buffer.from(payloadB64, 'base64url').toString('utf8')
+      );
+
+      if (!payload.projectId) {
+        throw new Error('License payload is missing "projectId".');
+      }
+
+      if (payload.projectId !== currentProjectId) {
+        throw new Error(
+          `Project ID mismatch. License is bound to project namespace "${payload.projectId}" ` +
+            `but configuration has "${currentProjectId}".`
+        );
+      }
+
+      if (!payload.expiresAt) {
+        throw new Error('License payload is missing expiration ("expiresAt").');
+      }
+
+      const currentTimestampSec = Math.floor(Date.now() / 1000);
+      if (currentTimestampSec > payload.expiresAt) {
+        throw new Error(
+          `License key has expired. Expiration: ${new Date(
+            payload.expiresAt * 1000
+          ).toDateString()}`
+        );
+      }
+
+      // 4. Validate Vector DB Provider against Tier (Production only)
+      if (isProduction && provider) {
+        const tier = (payload.tier || '').toLowerCase();
+        const normalizedProvider = provider.toLowerCase();
+
+        if (tier === 'hobby') {
+          // Hobby tier only supports PostgreSQL/pgvector/Supabase
+          const allowedHobby = ['postgresql', 'pgvector', 'supabase'];
+          if (!allowedHobby.includes(normalizedProvider)) {
+            throw new Error(
+              `The database provider "${provider}" is not allowed on the Hobby tier. ` +
+              `Hobby tier is restricted to PostgreSQL and Supabase. ` +
+              `Please upgrade to a Developer Pro or Enterprise plan.`
+            );
+          }
+        } else if (tier === 'pro' || tier === 'developer_pro' || tier === 'premium' || tier === 'growth') {
+          // Developer Pro supports: PostgreSQL, Pinecone, Qdrant, MongoDB, Milvus, ChromaDB
+          const allowedPro = [
+            'postgresql', 'pgvector', 'supabase',
+            'pinecone', 'qdrant', 'mongodb', 'milvus',
+            'chroma', 'chromadb'
+          ];
+          if (!allowedPro.includes(normalizedProvider)) {
+            throw new Error(
+              `The database provider "${provider}" is not allowed on the Developer Pro tier. ` +
+              `Please upgrade to an Enterprise plan.`
+            );
+          }
+        }
+        // Enterprise has access to all databases (including redis, weaviate, custom ones)
+      }
+
+      return payload;
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      throw new ConfigurationException(
+        `[Retrivora SDK] License key validation failed: ${message}`
+      );
+    }
+  }
+}