@retrivora-ai/rag-engine 1.9.6 → 1.9.8

package/src/core/VectorPlugin.ts

@@ -5,6 +5,8 @@ import { Pipeline } from './Pipeline';
 import { ProviderHealthCheck, HealthCheckResult } from './ProviderHealthCheck';
 import { IngestDocument, ChatResponse } from '../types';
 import { ChatMessage } from '../types';
+import { LicenseVerifier } from './LicenseVerifier';
+import { wrapError, RetrivoraErrorCode } from '../exceptions';
 
 /**
  * VectorPlugin — main orchestrator class.
@@ -23,12 +25,22 @@ export class VectorPlugin {
   private validationPromise: Promise<void>;
 
   constructor(hostConfig?: Partial<RagConfig>) {
-    this.config = ConfigResolver.resolve(hostConfig);
+    const resolvedConfig = ConfigResolver.resolve(hostConfig);
+    this.config = resolvedConfig;
     
     // Start validation early
-    this.validationPromise = ConfigValidator.validateAndThrow(this.config);
+    this.validationPromise = (async () => {
+      await ConfigValidator.validateAndThrow(resolvedConfig);
+      LicenseVerifier.verify(
+        resolvedConfig.licenseKey,
+        resolvedConfig.projectId,
+        resolvedConfig.vectorDb?.provider
+      );
+    })();
+    // Prevent unhandled promise rejection warning
+    this.validationPromise.catch(() => {});
     
-    this.pipeline = new Pipeline(this.config);
+    this.pipeline = new Pipeline(resolvedConfig);
   }
 
   /**
@@ -70,31 +82,82 @@ export class VectorPlugin {
    * Run a chat query.
    */
   async chat(message: string, history: ChatMessage[] = [], namespace?: string): Promise<ChatResponse> {
-    await this.validationPromise;
-    return this.pipeline.ask(message, history, namespace);
+    try {
+      await this.validationPromise;
+    } catch (err) {
+      throw wrapError(err, 'CONFIGURATION_ERROR');
+    }
+    try {
+      return await this.pipeline.ask(message, history, namespace);
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
   /**
    * Run a streaming chat query.
    */
   async *chatStream(message: string, history: ChatMessage[] = [], namespace?: string) {
-    await this.validationPromise;
-    yield* this.pipeline.askStream(message, history, namespace);
+    try {
+      await this.validationPromise;
+    } catch (err) {
+      throw wrapError(err, 'CONFIGURATION_ERROR');
+    }
+    try {
+      const stream = this.pipeline.askStream(message, history, namespace);
+      for await (const chunk of stream) {
+        yield chunk;
+      }
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
   /**
    * Ingest documents into the vector database.
    */
   async ingest(documents: IngestDocument[], namespace?: string): Promise<Array<{ docId: string | number; chunksIngested: number }>> {
-    await this.validationPromise;
-    return this.pipeline.ingest(documents, namespace);
+    try {
+      await this.validationPromise;
+    } catch (err) {
+      throw wrapError(err, 'CONFIGURATION_ERROR');
+    }
+    try {
+      return await this.pipeline.ingest(documents, namespace);
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
   /**
    * Get auto-suggestions based on a query prefix.
    */
   async getSuggestions(query: string, namespace?: string): Promise<string[]> {
-    await this.validationPromise;
-    return this.pipeline.getSuggestions(query, namespace);
+    try {
+      await this.validationPromise;
+    } catch (err) {
+      throw wrapError(err, 'CONFIGURATION_ERROR');
+    }
+    try {
+      return await this.pipeline.getSuggestions(query, namespace);
+    } catch (err) {
+      throw wrapError(err, 'RETRIEVAL_FAILED');
+    }
   }
 }
+

package/src/core/mcp.ts

@@ -0,0 +1,261 @@
+import { spawn, ChildProcess } from 'child_process';
+import { MCPServerConfig } from '../config/RagConfig';
+
+interface MCPTool {
+  name: string;
+  description?: string;
+  inputSchema: any;
+}
+
+interface MCPCallResult {
+  content: Array<{ type: string; text?: string; [key: string]: any }>;
+  isError?: boolean;
+}
+
+export class MCPClient {
+  private childProcess?: ChildProcess;
+  private sseUrl?: string;
+  private messageIdCounter = 0;
+  private pendingRequests = new Map<number, { resolve: (val: any) => void; reject: (err: any) => void }>();
+  private stdioBuffer = '';
+  private initialized = false;
+
+  constructor(private config: MCPServerConfig) {}
+
+  private getNextMessageId(): number {
+    return ++this.messageIdCounter;
+  }
+
+  /**
+   * Connect and initialize the MCP server connection.
+   */
+  async connect(): Promise<void> {
+    if (this.initialized) return;
+
+    if (this.config.transport === 'stdio') {
+      await this.connectStdio();
+    } else {
+      await this.connectSSE();
+    }
+
+    // Call initialize
+    try {
+      const response = await this.sendJsonRpc('initialize', {
+        protocolVersion: '2024-11-05',
+        capabilities: {},
+        clientInfo: {
+          name: 'retrivora-sdk',
+          version: '1.0.0',
+        },
+      });
+
+      // Send initialized notification
+      await this.sendNotification('notifications/initialized');
+      this.initialized = true;
+      console.log(`[MCPClient] Connected and initialized server "${this.config.name}" successfully.`);
+    } catch (err: any) {
+      this.disconnect();
+      throw new Error(`[MCPClient] Failed to initialize server "${this.config.name}": ${err.message}`);
+    }
+  }
+
+  private async connectStdio(): Promise<void> {
+    const cmd = this.config.command;
+    if (!cmd) {
+      throw new Error(`[MCPClient] Command option is required for stdio transport on "${this.config.name}"`);
+    }
+
+    const args = this.config.args || [];
+    const env = { ...process.env, ...(this.config.env || {}) };
+
+    this.childProcess = spawn(cmd, args, { env, shell: true });
+
+    if (!this.childProcess || !this.childProcess.stdout || !this.childProcess.stdin) {
+      throw new Error(`[MCPClient] Failed to spawn stdio child process for "${this.config.name}"`);
+    }
+
+    this.childProcess.stdout.on('data', (data: Buffer) => {
+      this.stdioBuffer += data.toString('utf-8');
+      this.processStdioLines();
+    });
+
+    this.childProcess.stderr?.on('data', (data: Buffer) => {
+      console.warn(`[MCPClient] [stderr] [${this.config.name}]:`, data.toString('utf-8'));
+    });
+
+    this.childProcess.on('close', (code) => {
+      console.log(`[MCPClient] Connection closed for "${this.config.name}" with exit code: ${code}`);
+      this.disconnect();
+    });
+
+    // Wait a brief tick to ensure stdout is set up
+    await new Promise((resolve) => setTimeout(resolve, 100));
+  }
+
+  private processStdioLines() {
+    let newlineIndex: number;
+    while ((newlineIndex = this.stdioBuffer.indexOf('\n')) !== -1) {
+      const line = this.stdioBuffer.substring(0, newlineIndex).trim();
+      this.stdioBuffer = this.stdioBuffer.substring(newlineIndex + 1);
+
+      if (!line) continue;
+
+      try {
+        const payload = JSON.parse(line);
+        if (payload.id !== undefined) {
+          const req = this.pendingRequests.get(Number(payload.id));
+          if (req) {
+            this.pendingRequests.delete(Number(payload.id));
+            if (payload.error) {
+              req.reject(new Error(payload.error.message || JSON.stringify(payload.error)));
+            } else {
+              req.resolve(payload.result);
+            }
+          }
+        }
+      } catch (e) {
+        console.warn(`[MCPClient] Error parsing stdio JSON-RPC line from "${this.config.name}":`, e, 'Line content:', line);
+      }
+    }
+  }
+
+  private async connectSSE(): Promise<void> {
+    if (!this.config.url) {
+      throw new Error(`[MCPClient] URL is required for sse transport on "${this.config.name}"`);
+    }
+
+    // Set up standard SSE handshake / SSE connection
+    // In serverless contexts, we pull initial SSE transport configuration,
+    // and route HTTP POSTs to the client URL.
+    this.sseUrl = this.config.url;
+  }
+
+  private async sendJsonRpc(method: string, params: any): Promise<any> {
+    const id = this.getNextMessageId();
+    const request = {
+      jsonrpc: '2.0',
+      id,
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (!this.childProcess || !this.childProcess.stdin) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" is not connected (stdio process missing)`);
+      }
+      return new Promise((resolve, reject) => {
+        this.pendingRequests.set(id, { resolve, reject });
+        this.childProcess!.stdin!.write(JSON.stringify(request) + '\n', 'utf-8');
+      });
+    } else {
+      // SSE / HTTP POST fallback for cloud/serverless setups
+      if (!this.sseUrl) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" SSE URL is not initialized`);
+      }
+      const response = await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(request),
+      });
+      if (!response.ok) {
+        throw new Error(`SSE JSON-RPC POST failed with status ${response.status}`);
+      }
+      const payload = await response.json();
+      if (payload.error) {
+        throw new Error(payload.error.message || JSON.stringify(payload.error));
+      }
+      return payload.result;
+    }
+  }
+
+  private async sendNotification(method: string, params?: any): Promise<void> {
+    const notification = {
+      jsonrpc: '2.0',
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (this.childProcess?.stdin) {
+        this.childProcess.stdin.write(JSON.stringify(notification) + '\n', 'utf-8');
+      }
+    } else if (this.sseUrl) {
+      await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(notification),
+      }).catch(err => console.warn('[MCPClient] Failed to send notification over SSE:', err));
+    }
+  }
+
+  /**
+   * List all tools offered by the MCP server.
+   */
+  async listTools(): Promise<MCPTool[]> {
+    await this.connect();
+    const result = await this.sendJsonRpc('tools/list', {});
+    return result?.tools || [];
+  }
+
+  /**
+   * Execute a tool on the MCP server.
+   */
+  async callTool(name: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    await this.connect();
+    return await this.sendJsonRpc('tools/call', { name, arguments: args });
+  }
+
+  disconnect() {
+    this.initialized = false;
+    this.pendingRequests.forEach(req => req.reject(new Error('MCPClient disconnected')));
+    this.pendingRequests.clear();
+    
+    if (this.childProcess) {
+      try {
+        this.childProcess.kill();
+      } catch { /* silent */ }
+      this.childProcess = undefined;
+    }
+  }
+}
+
+export class MCPRegistry {
+  private clients: MCPClient[] = [];
+
+  constructor(servers: MCPServerConfig[] = []) {
+    this.clients = servers.map(cfg => new MCPClient(cfg));
+  }
+
+  async getAllTools(): Promise<Array<{ serverName: string; tool: MCPTool }>> {
+    const list: Array<{ serverName: string; tool: MCPTool }> = [];
+    await Promise.all(
+      this.clients.map(async (client) => {
+        try {
+          const tools = await client.listTools();
+          tools.forEach(t => list.push({ serverName: client['config'].name, tool: t }));
+        } catch (e) {
+          console.warn(`[MCPRegistry] Failed to fetch tools from server "${client['config'].name}":`, e);
+        }
+      })
+    );
+    return list;
+  }
+
+  async callTool(toolName: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    // Find client hosting this tool. (In standard MCP, tool names are unique per server block)
+    // We will search across all connected clients.
+    for (const client of this.clients) {
+      try {
+        const tools = await client.listTools();
+        if (tools.some(t => t.name === toolName)) {
+          return await client.callTool(toolName, args);
+        }
+      } catch { /* silent fallback */ }
+    }
+    throw new Error(`[MCPRegistry] Tool "${toolName}" not found on any registered MCP server.`);
+  }
+
+  disconnectAll() {
+    this.clients.forEach(c => c.disconnect());
+  }
+}

package/src/exceptions/index.ts

@@ -57,3 +57,55 @@ export class AuthenticationException extends RetrivoraError {
     super(message, 'AUTHENTICATION_ERROR', details);
   }
 }
+
+/**
+ * Wraps any unknown error into an appropriate RetrivoraError subclass.
+ */
+export function wrapError(
+  err: unknown,
+  defaultCode: RetrivoraErrorCode,
+  defaultMessage?: string,
+): RetrivoraError {
+  if (err instanceof RetrivoraError) {
+    return err;
+  }
+
+  const error = err as any;
+  const message = error?.message || defaultMessage || String(err);
+  const status = error?.status || error?.statusCode || error?.response?.status;
+  const code = error?.code;
+
+  // 1. Rate Limit Recognition
+  if (status === 429 || /rate[- ]?limit/i.test(message) || code === 'RATE_LIMIT_EXCEEDED') {
+    return new RateLimitException(message, err);
+  }
+
+  // 2. Authentication Recognition
+  if (
+    status === 401 ||
+    status === 403 ||
+    /unauthorized|auth|api[- ]?key/i.test(message) ||
+    code === 'INVALID_API_KEY'
+  ) {
+    return new AuthenticationException(message, err);
+  }
+
+  // Fallback Mapping based on defaultCode
+  switch (defaultCode) {
+    case 'PROVIDER_NOT_FOUND':
+      return new ProviderNotFoundException('provider', message, err);
+    case 'EMBEDDING_FAILED':
+      return new EmbeddingFailedException(message, err);
+    case 'RETRIEVAL_FAILED':
+      return new RetrievalException(message, err);
+    case 'RATE_LIMITED':
+      return new RateLimitException(message, err);
+    case 'CONFIGURATION_ERROR':
+      return new ConfigurationException(message, err);
+    case 'AUTHENTICATION_ERROR':
+      return new AuthenticationException(message, err);
+    default:
+      return new RetrivoraError(message, defaultCode, err);
+  }
+}
+