feat: Enhance PageIndex chat and search with OCR content and link RAG documents to PageIndex IDs.

Author: amethystani

src/components/DashboardViews/DocumentsView.tsx

@@ -195,6 +195,7 @@ export default function DocumentsView({ isDark = true }: { isDark?: boolean }) {
     const [pageIndexDocIdByDocId, setPageIndexDocIdByDocId] = useState<Record<string, string>>({});
     const [pageIndexLoadingByDocId, setPageIndexLoadingByDocId] = useState<Record<string, boolean>>({});
     const [pageIndexErrorByDocId, setPageIndexErrorByDocId] = useState<Record<string, string>>({});
+    const [currentPageIndexDocId, setCurrentPageIndexDocId] = useState<string | null>(null);
     const [chatMessages, setChatMessages] = useState<{ role: string; content: string }[]>([]);
     const [chatInput, setChatInput] = useState('');
     const [isChatting, setIsChatting] = useState(false);
@@ -299,6 +300,7 @@ export default function DocumentsView({ isDark = true }: { isDark?: boolean }) {
                 );
 
                 setPageIndexTree(result);
+                setCurrentPageIndexDocId(docId);
                 setIntelligenceTab('tree');
 
                 // Initial AI greeting
@@ -321,7 +323,8 @@ export default function DocumentsView({ isDark = true }: { isDark?: boolean }) {
                     selectedFile,
                     kbId,
                     user.id,
-                    () => { } // Silent progress for standard RAG
+                    () => { }, // Silent progress for standard RAG
+                    docId
                 ).then(async (ragDoc) => {
                     if (ragDoc) {
                         await ragService.setPageIndexDocId(ragDoc.id, docId);
@@ -356,7 +359,12 @@ export default function DocumentsView({ isDark = true }: { isDark?: boolean }) {
 
         try {
             const history = chatMessages.map(m => ({ role: m.role, content: m.content }));
-            const response = await pageIndexService.chatWithDocument(userMsg, pageIndexTree, history);
+            const response = await pageIndexService.chatWithDocument(
+                userMsg,
+                pageIndexTree,
+                history,
+                { docId: currentPageIndexDocId || undefined }
+            );
             setChatMessages(prev => [...prev, { role: 'assistant', content: response }]);
         } catch (err) {
             setChatMessages(prev => [...prev, { role: 'assistant', content: "Sorry, I encountered an error while processing your request." }]);

src/components/PageIndex/PageIndexView.tsx

@@ -166,7 +166,7 @@ export default function PageIndexView() {
 
         try {
             const history = chatMessages.map(m => ({ role: m.role, content: m.content }));
-            const response = await pageIndexService.chatWithDocument(userMsg, tree, history);
+            const response = await pageIndexService.chatWithDocument(userMsg, tree, history, { docId: docId || undefined });
             setChatMessages(prev => [...prev, { role: 'assistant', content: response }]);
         } catch (err) {
             setChatMessages(prev => [...prev, { role: 'assistant', content: "Sorry, I encountered an error while processing your request." }]);

src/services/localLLMService.ts

@@ -692,6 +692,21 @@ class GroqLLMService {
               });
             }
 
+            const ocrMatches = await pageIndexService.searchOcrPages(pageIndexDoc.pageindex_doc_id, query, 5);
+            if (ocrMatches.length > 0) {
+              return JSON.stringify({
+                source: 'pageindex_ocr',
+                document: {
+                  id: pageIndexDoc.id,
+                  name: pageIndexDoc.file_name,
+                  pageindex_doc_id: pageIndexDoc.pageindex_doc_id,
+                },
+                results: ocrMatches.map((m, i) => ({ rank: i + 1, ...m })),
+                total: ocrMatches.length,
+                note: 'Tree had no matches; using OCR content.',
+              });
+            }
+
             return JSON.stringify({
               source: 'pageindex',
               document: {
@@ -701,7 +716,7 @@ class GroqLLMService {
               },
               results: [],
               total: 0,
-              note: 'No matching nodes found in PageIndex.',
+              note: 'No matching nodes found in PageIndex or OCR.',
             });
           } else if (documentName) {
             pageIndexNote = 'No PageIndex document found for the requested name.';

src/services/pageIndexService.ts

@@ -23,8 +23,15 @@ export interface TreeResponse {
     retrieval_ready?: boolean;
 }
 
+interface OcrPage {
+    page_index: number;
+    markdown: string;
+    images?: string[];
+}
+
 class PageIndexService {
     private treeCache = new Map<string, PageIndexNode[]>();
+    private ocrCache = new Map<string, OcrPage[] | string>();
 
     private getPageIndexHeaders() {
         if (!PAGEINDEX_API_KEY) {
@@ -141,18 +148,62 @@ class PageIndexService {
         throw new Error(`PageIndex status: ${treeResponse.status}`);
     }
 
+    /**
+     * Fetch OCR results (page/node/raw).
+     */
+    async getOcr(docId: string, format: 'page' | 'node' | 'raw' = 'page'): Promise<OcrPage[] | string> {
+        const cacheKey = `${docId}:${format}`;
+        const cached = this.ocrCache.get(cacheKey);
+        if (cached) return cached;
+
+        const headers = this.getPageIndexHeaders();
+        const response = await axios.get(`${PAGEINDEX_API_URL}/doc/${docId}/`, {
+            headers,
+            params: {
+                type: 'ocr',
+                format,
+            },
+        });
+
+        const data = response.data;
+        const result = (Array.isArray(data) ? data : data?.result ?? data) as OcrPage[] | string;
+        this.ocrCache.set(cacheKey, result);
+        return result;
+    }
+
     /**
      * Chat with Groq using the flattened tree as context
      */
     async chatWithDocument(
         question: string, 
         treeNodes: PageIndexNode[], 
-        history: { role: string; content: string }[] = []
+        history: { role: string; content: string }[] = [],
+        options?: { docId?: string }
     ): Promise<string> {
         if (!GROQ_API_KEY) {
             throw new Error('Missing Groq API key');
         }
         const flattenedTree = this.flattenTree(treeNodes);
+        const treeTextChars = this.countTextChars(treeNodes);
+        let ocrContext = '';
+
+        if (options?.docId && treeTextChars < 200) {
+            try {
+                const matches = await this.searchOcrPages(options.docId, question, 4);
+                if (matches.length > 0) {
+                    ocrContext = matches
+                        .map(m => `[Page ${m.page_index}] ${m.excerpt}`)
+                        .join('\n');
+                } else {
+                    const raw = await this.getOcr(options.docId, 'raw');
+                    if (typeof raw === 'string' && raw.trim().length > 0) {
+                        ocrContext = raw.substring(0, 4000);
+                    }
+                }
+            } catch {
+                // Ignore OCR fallback errors, continue with tree only
+            }
+        }
 
         const systemPrompt = `You are a helpful document assistant.
 You have been given a structured index tree extracted from a PDF document.
@@ -168,7 +219,7 @@ Rules:
             { role: 'system', content: systemPrompt },
             { 
                 role: 'user', 
-                content: `Here is the indexed structure of the document:\n\n${flattenedTree}\n\nI will now ask questions about this document.` 
+                content: `Here is the indexed structure of the document:\n\n${flattenedTree}\n\n${ocrContext ? `Additional OCR content (page-based):\n${ocrContext}\n\n` : ''}I will now ask questions about this document.` 
             },
             { role: 'assistant', content: "Got it! I've read the full index. Ask me anything." },
             ...history,
@@ -220,6 +271,21 @@ Rules:
         return result;
     }
 
+    /**
+     * Count total text characters in the tree (summary/text).
+     */
+    private countTextChars(nodes: PageIndexNode[]): number {
+        let count = 0;
+        for (const node of nodes) {
+            const text = node.summary || node.text || '';
+            count += text.length;
+            if (node.nodes && node.nodes.length > 0) {
+                count += this.countTextChars(node.nodes);
+            }
+        }
+        return count;
+    }
+
     /**
      * Search the tree for a query (simple keyword match).
      */
@@ -280,6 +346,39 @@ Rules:
             });
     }
 
+    /**
+     * Search OCR page results for a query.
+     */
+    async searchOcrPages(docId: string, query: string, limit: number = 5): Promise<Array<{
+        page_index: number;
+        excerpt: string;
+        score: number;
+    }>> {
+        const terms = query.toLowerCase().split(/\s+/).filter(Boolean);
+        if (terms.length === 0) return [];
+
+        const ocr = await this.getOcr(docId, 'page');
+        if (!Array.isArray(ocr)) return [];
+
+        const scored = ocr.map((page) => {
+            const hay = (page.markdown || '').toLowerCase();
+            let score = 0;
+            for (const term of terms) {
+                if (hay.includes(term)) score += 1;
+            }
+            return { page, score };
+        }).filter(item => item.score > 0);
+
+        return scored
+            .sort((a, b) => b.score - a.score)
+            .slice(0, limit)
+            .map(({ page, score }) => ({
+                page_index: page.page_index,
+                excerpt: page.markdown.length > 400 ? `${page.markdown.substring(0, 400)}...` : page.markdown,
+                score,
+            }));
+    }
+
     /**
      * Count total nodes in the tree
      */

src/services/ragService.ts

@@ -202,6 +202,7 @@ class RAGService {
         kbId: string,
         userId: string,
         onProgress?: (progress: ProcessingProgress) => void,
+        pageIndexDocId?: string,
     ): Promise<UploadedDocument | null> {
         const fileType = file.name.split('.').pop()?.toLowerCase() || 'txt';
         let documentRecord: UploadedDocument | null = null;
@@ -233,6 +234,7 @@ class RAGService {
                     file_size: file.size,
                     storage_path: storagePath,
                     status: 'processing',
+                    ...(pageIndexDocId ? { pageindex_doc_id: pageIndexDocId } : {}),
                 })
                 .select()
                 .single();