@retrivora-ai/rag-engine 1.7.9 → 1.8.1

package/src/llm/providers/GeminiProvider.ts

@@ -14,53 +14,103 @@
  *   - model:  string  – e.g. "text-embedding-004"
  */
 
-import { GoogleGenAI } from '@google/genai';
+import { GoogleGenerativeAI } from '@google/generative-ai';
 import { ILLMProvider, EmbedOptions } from '../ILLMProvider';
 import { ChatMessage, ChatOptions } from '../../types';
 import { LLMConfig, EmbeddingConfig } from '../../config/RagConfig';
 import { IProviderValidator, IProviderHealthChecker, HealthCheckResult } from '../../core/ProviderInterfaces';
 import { ValidationError } from '../../core/ConfigValidator';
 
+const DEFAULT_EMBEDDING_MODEL = 'text-embedding-004';
+const DEFAULT_TEMPERATURE = 0.7;
+const DEFAULT_MAX_TOKENS = 1024;
+
+// ---------------------------------------------------------------------------
+// Internal helpers
+// ---------------------------------------------------------------------------
+
+/** Strip suffixes like `:latest` that are valid in Ollama but not in Google SDK. */
+function sanitizeModel(model: string): string {
+  return model ? model.split(':')[0] : model;
+}
+
+/** Build a GoogleGenerativeAI client for the given API key. */
+function buildClient(apiKey: string): GoogleGenerativeAI {
+  return new GoogleGenerativeAI(apiKey);
+}
+
+/**
+ * Apply a task-specific prefix to a text string when configured to do so.
+ * Returns the original string unchanged when no prefix applies.
+ */
+function applyPrefix(
+  text: string,
+  taskType: EmbedOptions['taskType'] | undefined,
+  queryPrefix: string | undefined,
+  docPrefix: string | undefined,
+): string {
+  if (taskType === 'query' && queryPrefix && !text.startsWith(queryPrefix)) {
+    return `${queryPrefix}${text}`;
+  }
+  if (taskType === 'document' && docPrefix && !text.startsWith(docPrefix)) {
+    return `${docPrefix}${text}`;
+  }
+  return text;
+}
+
+// ---------------------------------------------------------------------------
+// Provider
+// ---------------------------------------------------------------------------
+
 export class GeminiProvider implements ILLMProvider {
-  private readonly client: GoogleGenAI;
+  private readonly client: GoogleGenerativeAI;
   private readonly llmConfig: LLMConfig;
   private readonly embeddingConfig?: EmbeddingConfig;
 
   constructor(llmConfig: LLMConfig, embeddingConfig?: EmbeddingConfig) {
-    if (!llmConfig.apiKey) throw new Error('[GeminiProvider] llmConfig.apiKey is required');
-    this.client = new GoogleGenAI({ apiKey: llmConfig.apiKey });
-    
-    // Sanitize model names (e.g. remove ':latest' suffix used by Ollama)
-    this.llmConfig = {
-      ...llmConfig,
-      model: this.sanitizeModel(llmConfig.model)
-    };
-    
+    if (!llmConfig.apiKey) {
+      throw new Error('[GeminiProvider] llmConfig.apiKey is required');
+    }
+
+    this.llmConfig = { ...llmConfig, model: sanitizeModel(llmConfig.model) };
+    this.client = buildClient(this.llmConfig.apiKey!);
+
     if (embeddingConfig) {
       this.embeddingConfig = {
         ...embeddingConfig,
-        model: this.sanitizeModel(embeddingConfig.model)
+        model: sanitizeModel(embeddingConfig.model),
       };
     }
   }
 
+  // -------------------------------------------------------------------------
+  // Static factory helpers
+  // -------------------------------------------------------------------------
+
   static getValidator(): IProviderValidator {
     return {
       validate(config: Record<string, unknown>): ValidationError[] {
         const errors: ValidationError[] = [];
+
         if (!config.apiKey && !process.env.GOOGLE_GENAI_API_KEY) {
           errors.push({
             field: 'llm.apiKey',
             message: 'Gemini API key is required',
             suggestion: 'Set GOOGLE_GENAI_API_KEY environment variable or provide in config',
-            severity: 'error'
+            severity: 'error',
           });
         }
+
         if (!config.model) {
-          errors.push({ field: 'llm.model', message: 'Gemini model name is required', severity: 'error' });
+          errors.push({
+            field: 'llm.model',
+            message: 'Gemini model name is required',
+            severity: 'error',
+          });
         }
+
         return errors;
-      }
+      },
     };
   }
 
@@ -68,13 +118,17 @@ export class GeminiProvider implements ILLMProvider {
     return {
       async check(config: Record<string, unknown>): Promise<HealthCheckResult> {
         const timestamp = Date.now();
-        const apiKey = (config.apiKey as string) || process.env.GOOGLE_GENAI_API_KEY || '';
-        const modelName = config.model as string;
+        const apiKey = (config.apiKey as string | undefined) ?? process.env.GOOGLE_GENAI_API_KEY ?? '';
+        const modelName = sanitizeModel(config.model as string);
+
         try {
-          const { GoogleGenAI } = await import('@google/genai');
-          const genAI = new GoogleGenAI({ apiKey });
-          // Verify model availability
-          await genAI.models.get({ model: modelName });
+          const { GoogleGenerativeAI } = await import('@google/generative-ai');
+          const ai = new GoogleGenerativeAI(apiKey);
+
+          // Lightweight verification: attempt a minimal embed call
+          const model = ai.getGenerativeModel({ model: DEFAULT_EMBEDDING_MODEL });
+          await model.embedContent('health-check');
+
           return {
             healthy: true,
             provider: 'gemini',
@@ -89,179 +143,170 @@ export class GeminiProvider implements ILLMProvider {
             timestamp,
           };
         }
-      }
+      },
     };
   }
 
-  private sanitizeModel(model: string): string {
-    if (!model) return model;
-    // Strip :latest suffix common in Ollama but invalid in Google SDK
-    return model.split(':')[0];
+  // -------------------------------------------------------------------------
+  // Private utilities
+  // -------------------------------------------------------------------------
+
+  /** Resolve the embedding client — uses a separate client when the embedding
+   *  API key differs from the LLM API key. */
+  private get embeddingClient(): GoogleGenerativeAI {
+    if (this.embeddingConfig?.apiKey && this.embeddingConfig.apiKey !== this.llmConfig.apiKey) {
+      return buildClient(this.embeddingConfig.apiKey);
+    }
+    return this.client;
   }
 
+  /** Resolve the embedding model to use, in order of specificity. */
+  private resolveEmbeddingModel(optionsModel?: string): string {
+    return sanitizeModel(optionsModel ?? this.embeddingConfig?.model ?? DEFAULT_EMBEDDING_MODEL);
+  }
 
-  async chat(messages: ChatMessage[], context: string, options?: ChatOptions): Promise<string> {
-    const systemPrompt =
+  /**
+   * Build the system instruction string, inserting the RAG context either via
+   * the `{{context}}` placeholder or by appending it.
+   */
+  private buildSystemInstruction(context: string): string {
+    const base =
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Answer questions based on the provided context.\n\nContext:\n${context}`;
 
-    const systemInstruction = systemPrompt.includes('{{context}}')
-      ? systemPrompt.replace('{{context}}', context)
-      : `${systemPrompt}\n\nContext:\n${context}`;
+    return base.includes('{{context}}') ? base.replace('{{context}}', context) : `${base}\n\nContext:\n${context}`;
+  }
+
+  /**
+   * Convert ChatMessage[] to the Gemini contents format.
+   *
+   * Rules enforced by the Gemini API:
+   *   - Only `user` and `model` roles are allowed in `contents`.
+   *   - System messages must be passed via `systemInstruction`, not here.
+   *   - Messages must strictly alternate between `user` and `model`.
+   */
+  private buildGeminiContents(messages: ChatMessage[]): Array<{ role: 'user' | 'model'; parts: [{ text: string }] }> {
+    return messages
+      .filter((m) => m.role !== 'system')
+      .map((m) => ({
+        role: m.role === 'assistant' ? ('model' as const) : ('user' as const),
+        parts: [{ text: m.content }] as [{ text: string }],
+      }));
+  }
 
-    // Gemini requires alternating user/model messages, starting with user.
-    // If the first message is not user, or there are consecutive same-role messages, we need to adapt.
-    // To simplify and comply with typical usage, we'll map 'assistant' to 'model'.
-    const geminiMessages = messages.map((m) => ({
-      role: m.role === 'assistant' ? 'model' : 'user',
-      parts: [{ text: m.content }],
-    }));
+  // -------------------------------------------------------------------------
+  // ILLMProvider — chat
+  // -------------------------------------------------------------------------
 
-    const response = await this.client.models.generateContent({
+  async chat(messages: ChatMessage[], context: string, options?: ChatOptions): Promise<string> {
+    const model = this.client.getGenerativeModel({
       model: this.llmConfig.model,
-      contents: geminiMessages,
-      config: {
-        systemInstruction,
-        temperature: options?.temperature ?? this.llmConfig.temperature ?? 0.7,
-        maxOutputTokens: options?.maxTokens ?? this.llmConfig.maxTokens ?? 1024,
+      systemInstruction: this.buildSystemInstruction(context),
+    });
+
+    const result = await model.generateContent({
+      contents: this.buildGeminiContents(messages),
+      generationConfig: {
+        temperature: options?.temperature ?? this.llmConfig.temperature ?? DEFAULT_TEMPERATURE,
+        maxOutputTokens: options?.maxTokens ?? this.llmConfig.maxTokens ?? DEFAULT_MAX_TOKENS,
         stopSequences: options?.stop,
       },
     });
 
-    const text = typeof (response as { text: unknown }).text === 'function'
-      ? (response as { text: () => string }).text()
-      : (response.text ?? '');
-    return text;
+    return result.response.text();
   }
 
   async *chatStream(messages: ChatMessage[], context: string, options?: ChatOptions): AsyncIterable<string> {
-    const systemPrompt =
-      this.llmConfig.systemPrompt ??
-      `You are a helpful assistant. Answer questions based on the provided context.\n\nContext:\n${context}`;
-
-    const systemInstruction = systemPrompt.includes('{{context}}')
-      ? systemPrompt.replace('{{context}}', context)
-      : `${systemPrompt}\n\nContext:\n${context}`;
-
-    const geminiMessages = messages.map((m) => ({
-      role: m.role === 'assistant' ? 'model' : 'user',
-      parts: [{ text: m.content }],
-    }));
-
-    const result = await this.client.models.generateContentStream({
+    const model = this.client.getGenerativeModel({
       model: this.llmConfig.model,
-      contents: geminiMessages,
-      config: {
-        systemInstruction,
-        temperature: options?.temperature ?? this.llmConfig.temperature ?? 0.7,
-        maxOutputTokens: options?.maxTokens ?? this.llmConfig.maxTokens ?? 1024,
+      systemInstruction: this.buildSystemInstruction(context),
+    });
+
+    const result = await model.generateContentStream({
+      contents: this.buildGeminiContents(messages),
+      generationConfig: {
+        temperature: options?.temperature ?? this.llmConfig.temperature ?? DEFAULT_TEMPERATURE,
+        maxOutputTokens: options?.maxTokens ?? this.llmConfig.maxTokens ?? DEFAULT_MAX_TOKENS,
         stopSequences: options?.stop,
       },
     });
 
-    const stream = ((result as unknown as Record<string, unknown>)?.stream ?? result) as AsyncIterable<{ text?: string }>;
-    if (!stream) {
-      throw new Error('[GeminiProvider] generateContentStream returned undefined');
-    }
-
-    for await (const chunk of stream) {
-      const text = typeof (chunk as { text: unknown }).text === 'function'
-        ? (chunk as { text: () => string }).text()
-        : (chunk as { text?: string }).text;
+    for await (const chunk of result.stream) {
+      const text = chunk.text();
       if (text) yield text;
     }
   }
 
+  // -------------------------------------------------------------------------
+  // ILLMProvider — embeddings
+  // -------------------------------------------------------------------------
+
   async embed(text: string, options?: EmbedOptions): Promise<number[]> {
-    const model = this.sanitizeModel(
-      options?.model ??
-      this.embeddingConfig?.model ??
-      'text-embedding-004'
+    const content = applyPrefix(
+      text,
+      options?.taskType,
+      this.embeddingConfig?.queryPrefix,
+      this.embeddingConfig?.documentPrefix,
     );
 
-    const apiKey = this.embeddingConfig?.apiKey ?? this.llmConfig.apiKey;
-    const client = apiKey !== this.llmConfig.apiKey
-      ? new GoogleGenAI({ apiKey })
-      : this.client;
-
-    let content = text;
-    
-    // Dynamically apply prefixes from configuration if they exist
-    const queryPrefix = this.embeddingConfig?.queryPrefix;
-    const docPrefix = this.embeddingConfig?.documentPrefix;
-    
-    if (options?.taskType === 'query' && queryPrefix) {
-      if (!content.startsWith(queryPrefix)) {
-        content = `${queryPrefix}${text}`;
-      }
-    } else if (options?.taskType === 'document' && docPrefix) {
-      if (!content.startsWith(docPrefix)) {
-        content = `${docPrefix}${text}`;
-      }
-    }
+    const modelName = this.resolveEmbeddingModel(options?.model);
+    const model = this.embeddingClient.getGenerativeModel({ model: modelName });
 
-    const response = await client.models.embedContent({
-      model,
-      contents: content,
-      config: this.embeddingConfig?.dimensions ? { outputDimensionality: this.embeddingConfig.dimensions } : undefined
-    }).catch(err => {
-      console.error(`[GeminiProvider] Embedding failed for model "${model}":`, err.message);
+    const response = await model.embedContent({
+      content: { role: 'user', parts: [{ text: content }] },
+    }).catch((err: Error) => {
+      console.error(`[GeminiProvider] Embedding failed for model "${modelName}":`, err.message);
       throw err;
     });
 
-    return response.embeddings?.[0]?.values ?? [];
+    return response.embedding.values;
   }
 
   async batchEmbed(texts: string[], options?: EmbedOptions): Promise<number[][]> {
-    // Sequential fallback for Gemini (API doesn't have a direct batch endpoint in the same way, though you can pass multiple contents.
-    // For simplicity and to match the other providers' structure, we'll iterate or pass array.
-    // The new SDK supports an array of contents.
-    const model = this.sanitizeModel(
-      options?.model ??
-      this.embeddingConfig?.model ??
-      'text-embedding-004'
-    );
+    const modelName = this.resolveEmbeddingModel(options?.model);
+    const model = this.embeddingClient.getGenerativeModel({ model: modelName });
 
-    const apiKey = this.embeddingConfig?.apiKey ?? this.llmConfig.apiKey;
-    const client = apiKey !== this.llmConfig.apiKey
-      ? new GoogleGenAI({ apiKey })
-      : this.client;
-
-    let contents = texts;
-    
-    // Apply prefixes if needed
-    const queryPrefix = this.embeddingConfig?.queryPrefix;
-    const docPrefix = this.embeddingConfig?.documentPrefix;
-    
-    if (options?.taskType === 'query' && queryPrefix) {
-      contents = texts.map(text => text.startsWith(queryPrefix) ? text : `${queryPrefix}${text}`);
-    } else if (options?.taskType === 'document' && docPrefix) {
-      contents = texts.map(text => text.startsWith(docPrefix) ? text : `${docPrefix}${text}`);
-    }
+    const requests = texts.map((text) => ({
+      content: {
+        role: 'user',
+        parts: [{
+          text: applyPrefix(
+            text,
+            options?.taskType,
+            this.embeddingConfig?.queryPrefix,
+            this.embeddingConfig?.documentPrefix,
+          )
+        }]
+      },
+    }));
 
-    const response = await client.models.embedContent({
-      model,
-      contents,
-      config: this.embeddingConfig?.dimensions ? { outputDimensionality: this.embeddingConfig.dimensions } : undefined
-    }).catch(err => {
-      console.error(`[GeminiProvider] Batch embedding failed for model "${model}":`, err.message);
+    const response = await model.batchEmbedContents({
+      requests: requests.map(r => ({
+        model: `models/${modelName}`,
+        content: r.content,
+      }))
+    }).catch((err: Error) => {
+      console.error(`[GeminiProvider] Batch embedding failed for model "${modelName}":`, err.message);
       throw err;
     });
 
-    return response.embeddings?.map(e => e.values ?? []) ?? [];
+    return response.embeddings.map((e) => e.values ?? []);
   }
 
+  // -------------------------------------------------------------------------
+  // ILLMProvider — health
+  // -------------------------------------------------------------------------
+
   async ping(): Promise<boolean> {
     try {
-      // Perform a lightweight embed call to verify connectivity
-      await this.client.models.embedContent({
-        model: 'text-embedding-004',
-        contents: 'ping',
+      const model = this.embeddingClient.getGenerativeModel({
+        model: this.resolveEmbeddingModel(),
       });
+      await model.embedContent('ping');
       return true;
     } catch (err) {
       console.error('[GeminiProvider] Ping failed:', err);
       return false;
     }
   }
-}
+  }

package/src/llm/providers/OllamaProvider.ts

@@ -85,7 +85,7 @@ export class OllamaProvider implements ILLMProvider {
   }
 
   async chat(messages: ChatMessage[], context: string, options?: ChatOptions): Promise<string> {
-    const system = this.buildSystemPrompt(context);
+    const system = this.buildSystemPrompt(context, options?.systemPrompt);
 
     const { data } = await this.http.post<OllamaChatResponse>('/api/chat', {
       model: this.llmConfig.model,
@@ -104,7 +104,7 @@ export class OllamaProvider implements ILLMProvider {
   }
 
   async *chatStream(messages: ChatMessage[], context: string, options?: ChatOptions): AsyncIterable<string> {
-    const system = this.buildSystemPrompt(context);
+    const system = this.buildSystemPrompt(context, options?.systemPrompt);
 
     const response = await this.http.post('/api/chat', {
       model: this.llmConfig.model,
@@ -159,7 +159,11 @@ export class OllamaProvider implements ILLMProvider {
     }
   }
 
-  private buildSystemPrompt(context: string): string {
+  private buildSystemPrompt(context: string, overridePrompt?: string): string {
+    if (overridePrompt) {
+      // Per-call override: use as-is (context is already embedded in user message)
+      return overridePrompt;
+    }
     const systemPrompt =
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Use the provided context to answer the user's question.\n\nContext:\n${context}`;

package/src/llm/providers/OpenAIProvider.ts

@@ -80,15 +80,17 @@ export class OpenAIProvider implements ILLMProvider {
   }
 
   async chat(messages: ChatMessage[], context: string, options?: ChatOptions): Promise<string> {
-    const systemContent =
+    const resolvedSystemPrompt = options?.systemPrompt ??
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Answer questions based on the provided context.\n\nContext:\n${context}`;
 
     const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = {
       role: 'system',
-      content: systemContent.includes('{{context}}')
-        ? systemContent.replace('{{context}}', context)
-        : `${systemContent}\n\nContext:\n${context}`,
+      content: resolvedSystemPrompt.includes('{{context}}')
+        ? resolvedSystemPrompt.replace('{{context}}', context)
+        : options?.systemPrompt
+          ? resolvedSystemPrompt // per-call override: use as-is, context already in user message
+          : `${resolvedSystemPrompt}\n\nContext:\n${context}`,
     };
 
     const formattedMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
@@ -111,15 +113,17 @@ export class OpenAIProvider implements ILLMProvider {
   }
   
   async *chatStream(messages: ChatMessage[], context: string, options?: ChatOptions): AsyncIterable<string> {
-    const systemContent =
+    const resolvedSystemPrompt = options?.systemPrompt ??
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Answer questions based on the provided context.\n\nContext:\n${context}`;
 
     const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = {
       role: 'system',
-      content: systemContent.includes('{{context}}')
-        ? systemContent.replace('{{context}}', context)
-        : `${systemContent}\n\nContext:\n${context}`,
+      content: resolvedSystemPrompt.includes('{{context}}')
+        ? resolvedSystemPrompt.replace('{{context}}', context)
+        : options?.systemPrompt
+          ? resolvedSystemPrompt
+          : `${resolvedSystemPrompt}\n\nContext:\n${context}`,
     };
 
     const formattedMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [

package/src/types/chat.ts

@@ -21,6 +21,12 @@ export interface ChatOptions {
   temperature?: number;
   /** Stop sequences */
   stop?: string[];
+  /**
+   * Per-call system prompt override. When provided, providers should use this
+   * instead of (or in addition to) their configured llmConfig.systemPrompt.
+   * Useful for one-off analytical calls (e.g. UITransformer.analyzeAndDecide).
+   */
+  systemPrompt?: string;
 }
 
 export interface UseRagChatOptions {

package/src/types/index.ts

@@ -22,6 +22,7 @@ export interface ChatResponse {
   reply: string;
   sources: VectorMatch[];
   graphData?: GraphSearchResult;
+  ui_transformation?: unknown;
 }
 
 export interface GraphNode {
@@ -65,5 +66,85 @@ export interface Product {
   description?: string;
 }
 
+/**
+ * Supported UI visualization types
+ */
+export type VisualizationType =
+  | 'pie_chart'
+  | 'bar_chart'
+  | 'line_chart'
+  | 'radar_chart'
+  | 'table'
+  | 'product_carousel'
+  | 'carousel'
+  | 'text';
+
+/**
+ * Base structure for all UI transformation responses
+ */
+export interface UITransformationResponse {
+  type: VisualizationType;
+  title: string;
+  description?: string;
+  data: unknown;
+}
+
+/**
+ * Pie chart data structure
+ */
+export interface PieChartData {
+  label: string;
+  value: number;
+  inStockCount?: number;
+  outOfStockCount?: number;
+  [key: string]: unknown;
+}
+
+/**
+ * Bar chart data structure
+ */
+export interface BarChartData {
+  category: string;
+  value: number;
+  inStockCount?: number;
+  outOfStockCount?: number;
+  [key: string]: unknown;
+}
+
+/**
+ * Line chart data point
+ */
+export interface LineChartDataPoint {
+  timestamp: string | number;
+  value: number;
+  label?: string;
+  [key: string]: unknown;
+}
+
+/**
+ * Table representation
+ */
+export interface TableData {
+  columns: string[];
+  rows: (string | number | boolean)[][];
+}
+
+
+/**
+ * Product carousel item
+ */
+export interface CarouselProduct {
+  id: string | number;
+  name: string;
+  price?: number | string;
+  image?: string;
+  inStock?: boolean;
+  brand?: string;
+  description?: string;
+  [key: string]: unknown;
+}
+
+
+
 export * from './props';
 export * from './chat';