@retrivora-ai/rag-engine 0.4.3 → 0.4.4

package/src/config/constants.ts

@@ -39,6 +39,13 @@ export const EMBEDDING_PROVIDERS = [
   'custom',
 ] as const;
 
+export const GRAPH_DB_PROVIDERS = [
+  'neo4j',
+  'memgraph',
+  'simple',
+  'custom',
+] as const;
+
 export const PROVIDERS_WITH_EMBEDDINGS = [
   'openai',
   'ollama',

package/src/config/serverConfig.ts

@@ -169,6 +169,7 @@ export function getRagConfig(env: Record<string, string | undefined> = process.e
         "Hello! I'm your AI assistant. Ask me anything about your documents.",
       visualStyle: (readString(env, 'NEXT_PUBLIC_UI_VISUAL_STYLE') ?? readString(env, 'UI_VISUAL_STYLE') ?? 'glass') as 'glass' | 'solid',
       borderRadius: (readString(env, 'NEXT_PUBLIC_UI_BORDER_RADIUS') ?? readString(env, 'UI_BORDER_RADIUS') ?? 'xl') as 'none' | 'sm' | 'md' | 'lg' | 'xl' | 'full',
+      allowUpload: (readString(env, 'NEXT_PUBLIC_ALLOW_UPLOAD') ?? readString(env, 'UI_ALLOW_UPLOAD') ?? 'false') === 'true',
     },
     rag: {
       topK: readNumber(env, 'RAG_TOP_K', 5),

package/src/core/Pipeline.ts

@@ -1,11 +1,13 @@
 import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
+import { BaseGraphProvider } from '../providers/graphdb/BaseGraphProvider';
 import { ILLMProvider, ChatMessage } from '../llm/ILLMProvider';
 import { DocumentChunker } from '../rag/DocumentChunker';
+import { EntityExtractor } from '../rag/EntityExtractor';
 import { RagConfig } from '../config/RagConfig';
 import { ProviderRegistry } from './ProviderRegistry';
 import { BatchProcessor, BatchOptions } from './BatchProcessor';
 import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
-import { IngestDocument, ChatResponse } from '../types';
+import { IngestDocument, ChatResponse, GraphSearchResult } from '../types';
 
 interface QueryFieldHint {
   field?: string;
@@ -171,9 +173,11 @@ function buildQueryFilter(question: string, hints: QueryFieldHint[]): QueryFilte
  */
 export class Pipeline {
   private vectorDB!: BaseVectorProvider;
+  private graphDB?: BaseGraphProvider;
   private llmProvider!: ILLMProvider;
   private embeddingProvider!: ILLMProvider;
   private chunker: DocumentChunker;
+  private entityExtractor?: EntityExtractor;
   private config: RagConfig;
   private initialised = false;
 
@@ -202,6 +206,12 @@ export class Pipeline {
     this.llmProvider = llmProvider;
     this.embeddingProvider = embeddingProvider;
 
+    if (this.config.graphDb) {
+      this.graphDB = await ProviderRegistry.createGraphProvider(this.config.graphDb);
+      await this.graphDB.initialize();
+      this.entityExtractor = new EntityExtractor(this.llmProvider);
+    }
+
     await this.vectorDB.initialize();
     this.initialised = true;
   }
@@ -270,6 +280,16 @@ export class Pipeline {
           docId: doc.docId,
           chunksIngested: upsertResult.totalProcessed,
         });
+
+        // Graph ingestion
+        if (this.graphDB && this.entityExtractor) {
+          console.log(`[Pipeline] Extracting entities for doc ${doc.docId}...`);
+          for (const chunk of chunks) {
+            const { nodes, edges } = await this.entityExtractor.extract(chunk.content);
+            if (nodes.length > 0) await this.graphDB.addNodes(nodes);
+            if (edges.length > 0) await this.graphDB.addEdges(edges);
+          }
+        }
       } catch (error) {
         console.error(`[Pipeline] Failed to ingest document ${doc.docId}:`, error);
         results.push({ docId: doc.docId, chunksIngested: 0 });
@@ -286,29 +306,70 @@ export class Pipeline {
     const scoreThreshold = this.config.rag?.scoreThreshold ?? 0.0;
 
     try {
-      const queryVector = await this.embeddingProvider.embed(question, { taskType: 'query' });
+      let searchQuery = question;
+
+      // 1. Query Transformation
+      if (this.config.rag?.useQueryTransformation) {
+        searchQuery = await this.rewriteQuery(question, history);
+      }
+
+      // 2. Vector Retrieval
+      const queryVector = await this.embeddingProvider.embed(searchQuery, { taskType: 'query' });
       const fieldHints = extractQueryFieldHints(question);
-      // Build a provider-agnostic filter from hints + original query text. Providers
-      // may choose to use `queryText` for hybrid keyword matching and `metadata`
-      // for exact/filtered metadata queries.
       const filter = buildQueryFilter(question, fieldHints) as Record<string, unknown>;
-
-      // Keep the raw hints available as well for providers that want richer parsing
       filter.__entityHints = fieldHints;
 
       const rawMatches = await this.vectorDB.query(queryVector, topK, ns, filter);
-
       const sources = rawMatches.filter((m) => m.score >= scoreThreshold);
-      const context = sources.length
+
+      // 3. Graph Retrieval
+      let graphData: GraphSearchResult | undefined;
+      if (this.graphDB && this.config.rag?.useGraphRetrieval) {
+        graphData = await this.graphDB.query(searchQuery);
+      }
+
+      // 4. Context Augmentation
+      let context = sources.length
         ? sources.map((m, i) => `[Source ${i + 1}]\n${m.content}`).join('\n\n---\n\n')
         : 'No relevant context found.';
 
+      if (graphData && graphData.nodes.length > 0) {
+        const graphContext = graphData.nodes.map(n => 
+          `Entity: ${n.label} (${n.id})${n.properties ? ' - ' + JSON.stringify(n.properties) : ''}`
+        ).join('\n');
+        context = `GRAPH KNOWLEDGE:\n${graphContext}\n\nVECTOR CONTEXT:\n${context}`;
+      }
+
+      // 5. Generation
       const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
       const reply = await this.llmProvider.chat(messages, context);
 
-      return { reply, sources };
+      return { reply, sources, graphData };
     } catch (error) {
       throw new Error(`[Pipeline] Chat failed: ${error instanceof Error ? error.message : String(error)}`);
     }
   }
+
+  /**
+   * Rewrite the user query for better retrieval performance.
+   */
+  private async rewriteQuery(question: string, history: ChatMessage[]): Promise<string> {
+    const prompt = `
+Given the following conversation history and a new question, rewrite the question to be a better search query for a vector database. 
+Focus on extracting the core intent and entities. Do not answer the question, just rewrite it.
+
+History:
+${history.map(m => `${m.role}: ${m.content}`).join('\n')}
+
+New Question: ${question}
+
+Optimized Search Query:`;
+
+    const rewrite = await this.llmProvider.chat([
+      { role: 'system', content: 'You are an assistant that optimizes search queries for RAG systems.' },
+      { role: 'user', content: prompt }
+    ], '');
+
+    return rewrite.trim() || question;
+  }
 }

package/src/core/ProviderRegistry.ts

@@ -1,5 +1,6 @@
 import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
-import { VectorDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig';
+import { BaseGraphProvider } from '@/providers/graphdb/BaseGraphProvider';
+import { VectorDBConfig, GraphDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig';
 import { LLMFactory } from '../llm/LLMFactory';
 import { ILLMProvider } from '../llm/ILLMProvider';
 
@@ -12,18 +13,19 @@ import { ILLMProvider } from '../llm/ILLMProvider';
  */
 export class ProviderRegistry {
   private static vectorProviders: Record<string, new (config: VectorDBConfig) => BaseVectorProvider> = {};
+  private static graphProviders: Record<string, new (config: GraphDBConfig) => BaseGraphProvider> = {};
 
-  /**
-   * 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;
   }
 
+  /**
+   * Register a custom graph provider class by name.
+   */
+  static registerGraphProvider(name: string, providerClass: new (config: GraphDBConfig) => BaseGraphProvider) {
+    this.graphProviders[name] = providerClass;
+  }
+
   /**
    * Creates a vector database provider based on the configuration.
    * Built-in providers are dynamically imported to avoid bundling all SDKs.
@@ -85,6 +87,36 @@ export class ProviderRegistry {
     }
   }
 
+  /**
+   * Creates a graph database provider based on the configuration.
+   */
+  static async createGraphProvider(config: GraphDBConfig): Promise<BaseGraphProvider> {
+    const { provider } = config;
+
+    // Custom registered provider takes priority
+    if (this.graphProviders[provider]) {
+      return new this.graphProviders[provider](config);
+    }
+
+    // Built-in providers — lazy-loaded
+    switch (provider) {
+      case 'neo4j': {
+        // Placeholder for real Neo4j provider
+        throw new Error('[ProviderRegistry] Neo4j provider not implemented yet.');
+      }
+      case 'simple': {
+        const { SimpleGraphProvider } = await import('@/providers/graphdb/SimpleGraphProvider');
+        return new SimpleGraphProvider(config);
+      }
+      default:
+        throw new Error(
+          `[ProviderRegistry] Unsupported graph provider: "${provider}". ` +
+          `Built-in providers: simple. ` +
+          `For custom providers, call ProviderRegistry.registerGraphProvider("${provider}", YourClass).`
+        );
+    }
+  }
+
   /**
    * Creates an LLM provider based on the configuration.
    */

package/src/index.ts

@@ -6,6 +6,7 @@
 // ── React Components ─────────────────────────────────────────
 export { ChatWidget } from './components/ChatWidget';
 export { ChatWindow } from './components/ChatWindow';
+export { DocumentUpload } from './components/DocumentUpload';
 export { MessageBubble } from './components/MessageBubble';
 export { SourceCard } from './components/SourceCard';
 export { ConfigProvider, useConfig } from './components/ConfigProvider';

package/src/providers/graphdb/BaseGraphProvider.ts

@@ -0,0 +1,43 @@
+import { GraphDBConfig } from '../../config/RagConfig';
+import { GraphNode, Edge, GraphSearchResult } from '../../types';
+
+/**
+ * BaseGraphProvider — Abstract base class for Graph DB providers.
+ */
+export abstract class BaseGraphProvider {
+  protected config: GraphDBConfig;
+
+  constructor(config: GraphDBConfig) {
+    this.config = config;
+  }
+
+  /**
+   * Initialise connection to the graph database.
+   */
+  abstract initialize(): Promise<void>;
+
+  /**
+   * Add nodes to the graph.
+   */
+  abstract addNodes(nodes: GraphNode[]): Promise<void>;
+
+  /**
+   * Add edges (relationships) to the graph.
+   */
+  abstract addEdges(edges: Edge[]): Promise<void>;
+
+  /**
+   * Query the graph for relevant nodes and edges.
+   */
+  abstract query(queryText: string, limit?: number): Promise<GraphSearchResult>;
+
+  /**
+   * Check if the database is reachable.
+   */
+  abstract ping(): Promise<boolean>;
+
+  /**
+   * Close connection.
+   */
+  abstract disconnect(): Promise<void>;
+}

package/src/providers/graphdb/SimpleGraphProvider.ts

@@ -0,0 +1,66 @@
+import { BaseGraphProvider } from './BaseGraphProvider';
+import { GraphNode, Edge, GraphSearchResult } from '../../types';
+
+/**
+ * SimpleGraphProvider — In-memory graph storage for demonstration and local dev.
+ */
+export class SimpleGraphProvider extends BaseGraphProvider {
+  private nodes: Map<string, GraphNode> = new Map();
+  private edges: Edge[] = [];
+
+  async initialize(): Promise<void> {
+    console.log('[SimpleGraphProvider] Initialised in-memory graph store.');
+  }
+
+  async addNodes(nodes: GraphNode[]): Promise<void> {
+    for (const node of nodes) {
+      this.nodes.set(node.id, node);
+    }
+  }
+
+  async addEdges(edges: Edge[]): Promise<void> {
+    this.edges.push(...edges);
+  }
+
+  async query(queryText: string, limit = 5): Promise<GraphSearchResult> {
+    const q = queryText.toLowerCase();
+    
+    // Simple keyword matching for demonstration
+    const matchedNodes = Array.from(this.nodes.values()).filter(node => 
+      node.id.toLowerCase().includes(q) || 
+      node.label.toLowerCase().includes(q) ||
+      JSON.stringify(node.properties).toLowerCase().includes(q)
+    ).slice(0, limit);
+
+    const matchedNodeIds = new Set(matchedNodes.map(n => n.id));
+    const matchedEdges = this.edges.filter(edge => 
+      matchedNodeIds.has(edge.source) || matchedNodeIds.has(edge.target)
+    );
+
+    // Also include connected nodes
+    for (const edge of matchedEdges) {
+      if (!matchedNodeIds.has(edge.source)) {
+        const source = this.nodes.get(edge.source);
+        if (source) matchedNodes.push(source);
+      }
+      if (!matchedNodeIds.has(edge.target)) {
+        const target = this.nodes.get(edge.target);
+        if (target) matchedNodes.push(target);
+      }
+    }
+
+    return {
+      nodes: Array.from(new Set(matchedNodes)),
+      edges: matchedEdges
+    };
+  }
+
+  async ping(): Promise<boolean> {
+    return true;
+  }
+
+  async disconnect(): Promise<void> {
+    this.nodes.clear();
+    this.edges = [];
+  }
+}

package/src/rag/DocumentChunker.ts

@@ -23,19 +23,23 @@ export interface ChunkOptions {
   docId?: string | number;
   /** Extra metadata to attach to every chunk */
   metadata?: Record<string, unknown>;
+  /** Characters used to split text, in order of priority */
+  separators?: string[];
 }
 
 export class DocumentChunker {
   private readonly chunkSize: number;
   private readonly chunkOverlap: number;
+  private readonly separators: string[];
 
-  constructor(chunkSize = 1000, chunkOverlap = 200) {
+  constructor(chunkSize = 1000, chunkOverlap = 200, separators = ['\n\n', '\n', ' ', '']) {
     this.chunkSize = chunkSize;
     this.chunkOverlap = chunkOverlap;
+    this.separators = separators;
   }
 
   /**
-   * Split a single text string into overlapping chunks.
+   * Split a single text string into overlapping chunks using a recursive strategy.
    */
   chunk(text: string, options: ChunkOptions = {}): Chunk[] {
     const {
@@ -43,53 +47,92 @@ export class DocumentChunker {
       chunkOverlap = this.chunkOverlap,
       docId = `doc_${Date.now()}`,
       metadata = {},
+      separators = this.separators,
     } = options;
 
-    // Normalise whitespace
-    const cleaned = text.replace(/\r\n/g, '\n').trim();
-    if (!cleaned) return [];
-
-    // Split into sentences (handles . ! ? followed by space/newline)
-    const sentences = cleaned
-      .split(/(?<=[.!?])\s+|\n{2,}/)
-      .map((s) => s.trim())
-      .filter(Boolean);
-
-    const chunks: Chunk[] = [];
-    let current = '';
+    const finalChunks: Chunk[] = [];
+    const splits = this.recursiveSplit(text, separators, chunkSize);
+    
+    let currentChunk: string[] = [];
+    let currentLength = 0;
     let chunkIndex = 0;
 
-    for (const sentence of sentences) {
-      if ((current + ' ' + sentence).trim().length <= chunkSize) {
-        current = current ? `${current} ${sentence}` : sentence;
-      } else {
-        if (current) {
-          chunks.push({
-            id: `${docId}_chunk_${chunkIndex++}`,
-            content: current.trim(),
-            metadata: { ...metadata, docId, chunkIndex: chunkIndex - 1 },
-          });
-        }
+    for (const split of splits) {
+      if (currentLength + split.length > chunkSize && currentChunk.length > 0) {
+        // Current chunk is full, push it
+        finalChunks.push({
+          id: `${docId}_chunk_${chunkIndex++}`,
+          content: currentChunk.join('').trim(),
+          metadata: { ...metadata, docId, chunkIndex: chunkIndex - 1 },
+        });
 
-        // Start new chunk with overlap from the previous one
-        if (chunkOverlap > 0 && current.length > chunkOverlap) {
-          current = current.slice(-chunkOverlap) + ' ' + sentence;
-        } else {
-          current = sentence;
+        // Start new chunk with overlap
+        // This is a simplified overlap: we keep items from the end of the previous chunk
+        // until we reach the overlap size.
+        const overlapItems = [];
+        let overlapLen = 0;
+        for (let i = currentChunk.length - 1; i >= 0; i--) {
+          if (overlapLen + currentChunk[i].length <= chunkOverlap) {
+            overlapItems.unshift(currentChunk[i]);
+            overlapLen += currentChunk[i].length;
+          } else {
+            break;
+          }
         }
+        currentChunk = overlapItems;
+        currentLength = overlapLen;
       }
+      
+      currentChunk.push(split);
+      currentLength += split.length;
     }
 
     // Push the last remaining chunk
-    if (current.trim()) {
-      chunks.push({
+    if (currentChunk.length > 0) {
+      finalChunks.push({
         id: `${docId}_chunk_${chunkIndex}`,
-        content: current.trim(),
+        content: currentChunk.join('').trim(),
         metadata: { ...metadata, docId, chunkIndex },
       });
     }
 
-    return chunks;
+    return finalChunks;
+  }
+
+  /**
+   * Recursively split text based on separators.
+   */
+  private recursiveSplit(text: string, separators: string[], chunkSize: number): string[] {
+    const finalSplits: string[] = [];
+    
+    // Find the first separator that exists in the text
+    let separator = separators[separators.length - 1]; // default to last (usually empty string)
+    let nextSeparators: string[] = [];
+
+    for (let i = 0; i < separators.length; i++) {
+      if (text.includes(separators[i])) {
+        separator = separators[i];
+        nextSeparators = separators.slice(i + 1);
+        break;
+      }
+    }
+
+    // Split the text
+    const parts = text.split(separator);
+    
+    for (const part of parts) {
+      if (part.length <= chunkSize) {
+        finalSplits.push(part + separator);
+      } else if (nextSeparators.length > 0) {
+        // Part is too long, try splitting with next level separators
+        finalSplits.push(...this.recursiveSplit(part, nextSeparators, chunkSize));
+      } else {
+        // No more separators, just hard split (should not happen if empty string is the last separator)
+        finalSplits.push(part);
+      }
+    }
+
+    return finalSplits;
   }
 
   /**

package/src/rag/EntityExtractor.ts

@@ -0,0 +1,43 @@
+import { ILLMProvider } from '../llm/ILLMProvider';
+import { GraphNode, Edge } from '../types';
+
+/**
+ * EntityExtractor — Uses an LLM to extract entities and relations from text.
+ */
+export class EntityExtractor {
+  private llm: ILLMProvider;
+
+  constructor(llm: ILLMProvider) {
+    this.llm = llm;
+  }
+
+  /**
+   * Extract nodes and edges from a text chunk.
+   */
+  async extract(text: string): Promise<{ nodes: GraphNode[]; edges: Edge[] }> {
+    const prompt = `
+Extract entities and relationships from the following text.
+Format the output as JSON with two keys: "nodes" (array of {id, label, properties}) and "edges" (array of {source, target, type, properties}).
+Use the same ID for the same entity.
+
+Text:
+"${text}"
+
+Output JSON:
+`;
+
+    const response = await this.llm.chat([
+      { role: 'system', content: 'You are an expert at knowledge graph extraction. Respond only with valid JSON.' },
+      { role: 'user', content: prompt }
+    ], '');
+
+    try {
+      // Clean up LLM response (sometimes it adds markdown blocks)
+      const cleanJson = response.replace(/```json\n?|\n?```/g, '').trim();
+      return JSON.parse(cleanJson);
+    } catch (error) {
+      console.warn('[EntityExtractor] Failed to parse LLM response as JSON:', response);
+      return { nodes: [], edges: [] };
+    }
+  }
+}

package/src/types/index.ts

@@ -21,4 +21,23 @@ export interface IngestDocument {
 export interface ChatResponse {
   reply: string;
   sources: VectorMatch[];
+  graphData?: GraphSearchResult;
+}
+
+export interface GraphNode {
+  id: string;
+  label: string;
+  properties?: Record<string, unknown>;
+}
+
+export interface Edge {
+  source: string;
+  target: string;
+  type: string;
+  properties?: Record<string, unknown>;
+}
+
+export interface GraphSearchResult {
+  nodes: GraphNode[];
+  edges: Edge[];
 }