@retrivora-ai/rag-engine 1.9.7 → 1.9.9

package/src/core/LicenseVerifier.ts

@@ -0,0 +1,260 @@
+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 {
+  // A simple in-memory cache to save CPU overhead of cryptographic signature checks.
+  private static readonly cache = new Map<
+    string,
+    { payload: LicensePayload; cachedAt: number }
+  >();
+  private static readonly CACHE_TTL_MS = 60 * 1000; // Cache verified license for 60 seconds
+
+  // 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';
+    const now = Date.now();
+
+    // Check verification cache to bypass crypto computations on hot paths
+    if (licenseKey) {
+      const cacheKey = `${licenseKey}:${currentProjectId}:${provider || ''}:${publicKeyOverride || ''}`;
+      const cached = this.cache.get(cacheKey);
+      if (cached && now - cached.cachedAt < this.CACHE_TTL_MS) {
+        // Enforce expiration check even for cached items
+        const currentTimestampSec = Math.floor(now / 1000);
+        if (cached.payload.expiresAt && currentTimestampSec > cached.payload.expiresAt) {
+          this.cache.delete(cacheKey);
+        } else {
+          return cached.payload;
+        }
+      }
+    }
+
+    // 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)
+      }
+
+      if (licenseKey) {
+        const cacheKey = `${licenseKey}:${currentProjectId}:${provider || ''}:${publicKeyOverride || ''}`;
+        this.cache.set(cacheKey, { payload, cachedAt: now });
+      }
+
+      return payload;
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      throw new ConfigurationException(
+        `[Retrivora SDK] License key validation failed: ${message}`
+      );
+    }
+  }
+
+  static async verifyIP(licenseKey: string | undefined): Promise<void> {
+    if (!licenseKey) return;
+    try {
+      const parts = licenseKey.split('.');
+      if (parts.length !== 3) return;
+      const [, payloadB64] = parts;
+      const payload: LicensePayload & { ipv4?: string; ipv6?: string } = JSON.parse(
+        Buffer.from(payloadB64, 'base64url').toString('utf8')
+      );
+
+      const tier = (payload.tier || '').toLowerCase();
+      // Skip check if Enterprise or Custom
+      if (tier === 'enterprise' || tier === 'custom') {
+        return;
+      }
+
+      // If Hobby (free_trial) or Developer Pro, perform the check
+      if (payload.ipv4 || payload.ipv6) {
+        let currentIpv4 = '';
+        let currentIpv6 = '';
+        
+        try {
+          const res = await fetch('https://api4.ipify.org?format=json', { signal: AbortSignal.timeout(3000) });
+          const data = await res.json();
+          currentIpv4 = data.ip;
+        } catch (e) {}
+        
+        try {
+          const res = await fetch('https://api6.ipify.org?format=json', { signal: AbortSignal.timeout(3000) });
+          const data = await res.json();
+          currentIpv6 = data.ip;
+        } catch (e) {}
+
+        const localIps: string[] = [];
+        try {
+          const osModule = require('os');
+          const interfaces = osModule.networkInterfaces();
+          for (const name of Object.keys(interfaces)) {
+            for (const iface of interfaces[name] || []) {
+              if (!iface.internal) {
+                localIps.push(iface.address);
+              }
+            }
+          }
+        } catch (e) {}
+        
+        const isIpv4Match = payload.ipv4 && (
+          currentIpv4 === payload.ipv4 || 
+          localIps.includes(payload.ipv4)
+        );
+        const isIpv6Match = payload.ipv6 && (
+          currentIpv6 === payload.ipv6 || 
+          localIps.includes(payload.ipv6)
+        );
+
+        if (payload.ipv4 && payload.ipv6) {
+          if (!isIpv4Match && !isIpv6Match) {
+            throw new Error(
+              `License key access restricted. This license key was created for a different system (authorized IPv4: ${payload.ipv4}, IPv6: ${payload.ipv6}) and cannot be used on this machine.`
+            );
+          }
+        } else if (payload.ipv4 && !isIpv4Match) {
+          throw new Error(
+            `License key access restricted. This license key was created for a different system (authorized IPv4: ${payload.ipv4}) and cannot be used on this machine.`
+          );
+        } else if (payload.ipv6 && !isIpv6Match) {
+          throw new Error(
+            `License key access restricted. This license key was created for a different system (authorized IPv6: ${payload.ipv6}) and cannot be used on this machine.`
+          );
+        }
+      }
+    } catch (err: any) {
+      if (err.message.includes('License key access restricted')) {
+        throw new ConfigurationException(`[Retrivora SDK] ${err.message}`);
+      }
+    }
+  }
+}

package/src/core/MultiAgentCoordinator.ts

@@ -0,0 +1,239 @@
+import { ILLMProvider } from '../llm/ILLMProvider';
+import { MCPRegistry } from './mcp';
+import { ChatMessage, ChatResponse } from '../types';
+
+interface CoordinatorOptions {
+  llmProvider: ILLMProvider;
+  mcpRegistry: MCPRegistry;
+  documentSearch: (query: string) => Promise<{ reply: string; sources: any[] }>;
+  maxIterations?: number;
+}
+
+export class MultiAgentCoordinator {
+  private llmProvider: ILLMProvider;
+  private mcpRegistry: MCPRegistry;
+  private documentSearch: (query: string) => Promise<{ reply: string; sources: any[] }>;
+  private maxIterations: number;
+
+  constructor(options: CoordinatorOptions) {
+    this.llmProvider = options.llmProvider;
+    this.mcpRegistry = options.mcpRegistry;
+    this.documentSearch = options.documentSearch;
+    this.maxIterations = options.maxIterations ?? 5;
+  }
+
+  /**
+   * Run the multi-agent coordination loop synchronously and return the final response.
+   */
+  async run(question: string, history: ChatMessage[] = []): Promise<ChatResponse> {
+    const generator = this.runStream(question, history);
+    let finalReply = '';
+    const sources: any[] = [];
+
+    for await (const chunk of generator) {
+      if (typeof chunk === 'string') {
+        finalReply += chunk;
+      } else if (chunk && typeof chunk === 'object') {
+        if ('reply' in chunk && chunk.reply) {
+          finalReply += chunk.reply;
+        }
+        if ('sources' in chunk && chunk.sources) {
+          sources.push(...chunk.sources);
+        }
+      }
+    }
+
+    return {
+      reply: finalReply,
+      sources,
+    };
+  }
+
+  /**
+   * Run the multi-agent coordination loop as an async stream, yielding intermediate thought blocks and tool execution events.
+   */
+  async *runStream(question: string, history: ChatMessage[] = []): AsyncGenerator<string | ChatResponse> {
+    // 1. Gather all tools
+    const mcpTools = await this.mcpRegistry.getAllTools();
+    
+    // Construct system instructions with tools
+    let systemPrompt = `You are a goal-driven supervisor agent coordinating a RAG search engine and external MCP servers to solve the user's task.
+You must think step-by-step before acting. Write your internal thinking process inside a \`<think>...</think>\` block first, then act or answer.
+
+Available Tools:
+- document_search(query: string): Search the vector database and knowledge base for documents.
+
+${mcpTools.length > 0 ? 'MCP Server Tools:' : ''}
+${mcpTools.map(mt => `- ${mt.tool.name}(${JSON.stringify(mt.tool.inputSchema?.properties || {})}): ${mt.tool.description || 'No description provided.'}`).join('\n')}
+
+Tool Calling Protocol:
+If you need to call a tool, write:
+TOOL_CALL: <toolName>({ "argName": "value" })
+And then STOP generating immediately.
+
+Once you have gathered enough information to fully satisfy the user's query, write your final response directly after the closing \`</think>\` block.
+`;
+
+    // 2. Setup message history for the supervisor
+    const supervisorHistory: ChatMessage[] = [
+      { role: 'system' as any, content: systemPrompt },
+      ...history,
+      { role: 'user', content: question }
+    ];
+
+    let iteration = 0;
+    const allSources: any[] = [];
+
+    while (iteration < this.maxIterations) {
+      iteration++;
+      
+      const context = ''; // Context is populated by tool call results directly in message history
+      let fullModelResponse = '';
+      let textBuffer = '';
+      let inThinking = false;
+      let toolCallDetected: { name: string; args: any } | null = null;
+
+      if (!this.llmProvider.chatStream) {
+        // Fallback for non-streaming providers
+        const reply = await this.llmProvider.chat(supervisorHistory, context);
+        fullModelResponse = reply;
+        
+        // Parse simulated thinking / tool call
+        const thinkIndex = reply.indexOf('<think>');
+        const endThinkIndex = reply.indexOf('</think>');
+        
+        let thinkingText = '';
+        let answerText = reply;
+
+        if (thinkIndex !== -1 && endThinkIndex !== -1) {
+          thinkingText = reply.substring(thinkIndex + 7, endThinkIndex);
+          answerText = reply.substring(endThinkIndex + 8);
+          yield { type: 'thinking', text: thinkingText } as any;
+        }
+
+        const toolMatch = answerText.match(/TOOL_CALL:\s*(\w+)\s*\(([\s\S]*)\)/);
+        if (toolMatch) {
+          const toolName = toolMatch[1];
+          const toolArgsStr = toolMatch[2];
+          try {
+            const toolArgs = JSON.parse(toolArgsStr.trim());
+            toolCallDetected = { name: toolName, args: toolArgs };
+          } catch (e) {
+            console.error('[MultiAgentCoordinator] Failed to parse tool args:', toolArgsStr, e);
+          }
+        } else {
+          yield answerText;
+        }
+      } else {
+        // Run with streaming completion
+        const stream = this.llmProvider.chatStream(supervisorHistory, context);
+
+        for await (const chunk of stream) {
+          fullModelResponse += chunk;
+          textBuffer += chunk;
+
+          // Check for <think> tag
+          if (!inThinking && textBuffer.includes('<think>')) {
+            inThinking = true;
+            const idx = textBuffer.indexOf('<think>');
+            textBuffer = textBuffer.slice(idx + 7);
+          }
+
+          // Check for </think> tag
+          if (inThinking && textBuffer.includes('</think>')) {
+            inThinking = false;
+            const idx = textBuffer.indexOf('</think>');
+            const thinkingDelta = textBuffer.slice(0, idx);
+            yield { type: 'thinking', text: thinkingDelta } as any;
+            textBuffer = textBuffer.slice(idx + 8);
+          }
+
+          // If we are in thinking mode, yield thinking delta
+          if (inThinking && textBuffer.length > 0) {
+            yield { type: 'thinking', text: textBuffer } as any;
+            textBuffer = '';
+          }
+
+          // Check for TOOL_CALL during stream
+          if (!inThinking && textBuffer.includes('TOOL_CALL:')) {
+            const idx = textBuffer.indexOf('TOOL_CALL:');
+            const precedingText = textBuffer.slice(0, idx);
+            if (precedingText.trim()) {
+              yield precedingText;
+            }
+            // Await remainder of stream to get the full tool call signature
+            textBuffer = textBuffer.slice(idx);
+          }
+
+          // If not in thinking and no tool call prefix is blocking, yield final answer text
+          if (!inThinking && !textBuffer.includes('TOOL_CALL:') && textBuffer.length > 0) {
+            yield textBuffer;
+            textBuffer = '';
+          }
+        }
+
+        // Flush any remaining buffer text
+        if (textBuffer.trim()) {
+          const toolMatch = textBuffer.match(/TOOL_CALL:\s*(\w+)\s*\(([\s\S]*)\)/);
+          if (toolMatch) {
+            const toolName = toolMatch[1];
+            const toolArgsStr = toolMatch[2];
+            try {
+              const toolArgs = JSON.parse(toolArgsStr.trim());
+              toolCallDetected = { name: toolName, args: toolArgs };
+            } catch (e) {
+              console.error('[MultiAgentCoordinator] Failed to parse tool args:', toolArgsStr, e);
+            }
+          } else {
+            yield textBuffer;
+          }
+        }
+      }
+
+      // 3. Append the model response to the supervisor history
+      supervisorHistory.push({ role: 'assistant', content: fullModelResponse });
+
+      // 4. Handle tool execution if requested, else break the loop
+      if (toolCallDetected) {
+        const { name: toolName, args: toolArgs } = toolCallDetected;
+        
+        yield { 
+          type: 'thinking', 
+          text: `\n[Agent Call] Executing tool "${toolName}" with parameters ${JSON.stringify(toolArgs)}...\n` 
+        } as any;
+
+        let toolResultText = '';
+        try {
+          if (toolName === 'document_search') {
+            const searchRes = await this.documentSearch(toolArgs.query || '');
+            toolResultText = `document_search returned:\n${searchRes.reply}\n\nSources count: ${searchRes.sources.length}`;
+            if (searchRes.sources && searchRes.sources.length > 0) {
+              allSources.push(...searchRes.sources);
+              yield { reply: '', sources: searchRes.sources };
+            }
+          } else {
+            // Call MCP tool
+            const mcpRes = await this.mcpRegistry.callTool(toolName, toolArgs);
+            toolResultText = `MCP Tool "${toolName}" returned:\n${JSON.stringify(mcpRes.content)}`;
+          }
+        } catch (err: any) {
+          toolResultText = `Error calling tool "${toolName}": ${err.message}`;
+        }
+
+        // Append tool output to history
+        supervisorHistory.push({
+          role: 'user',
+          content: `TOOL_RESULT for ${toolName}:\n${toolResultText}`
+        });
+
+        yield {
+          type: 'thinking',
+          text: `[Agent Output] Tool "${toolName}" executed. Continuing reasoning...\n`
+        } as any;
+      } else {
+        // No tool call requested — response is complete
+        break;
+      }
+    }
+  }
+}