Update main.py

main.py

@@ -19,7 +19,7 @@ from content_analyzer.document_parser import DocumentProcessor
 from search_engine.indexer import RetrieverBuilder
 from intelligence.orchestrator import AgentWorkflow
 from configuration import definitions, parameters
-                                 
+                             
 
 # Rate limiting configuration - 3 requests per hour per IP
 WINDOW_S = 3600
@@ -31,17 +31,17 @@ def rate_limit(request):
     """Thread-safe rate limiting per IP address."""
     ip = getattr(request.client, "host", "unknown")
     now = time.time()
-    
+
     with _calls_lock:
         q = _calls[ip]
         # Remove expired entries
         while q and (now - q[0]) > WINDOW_S:
             q.popleft()
-        
+    
         if len(q) >= MAX_CALLS:
             import gradio as gr
             raise gr.Error(f"Rate limit: {MAX_CALLS} requests per {WINDOW_S//60} minutes. Please wait.")
-        
+    
         q.append(now)
 
 
@@ -66,14 +66,14 @@ def format_chat_history(history: List[Dict]) -> str:
     """Format chat history as markdown for display."""
     if not history:
         return "*No conversation history yet. Ask a question to get started!*"
-    
+
     formatted = []
     for i, entry in enumerate(history, 1):
         timestamp = entry.get("timestamp", "")
         question = entry.get("question", "")
         answer = entry.get("answer", "")
         confidence = entry.get("confidence", "N/A")
-        
+    
         formatted.append(f"""
 ---
 ### 💬 Q{i} ({timestamp})
@@ -83,7 +83,7 @@ def format_chat_history(history: List[Dict]) -> str:
 
 *Confidence: {confidence}*
 """)
-    
+
     return "\n".join(formatted)
 
 
@@ -91,19 +91,19 @@ def format_document_context(documents: List, question: str = "") -> str:
     """Format retrieved documents with annotation highlighting."""
     if not documents:
         return "*No documents retrieved yet.*"
-    
+
     formatted = [f"### 📚 Retrieved Context ({len(documents)} chunks)\n"]
-    
+
     # Extract key terms from question for highlighting
     key_terms = []
     if question:
         stopwords = {'the', 'a', 'an', 'is', 'are', 'was', 'were', 'in', 'on', 'at', 'to', 'for', 'of', 'and', 'or', 'what', 'how', 'why', 'when', 'where', 'which'}
         key_terms = [word.lower() for word in question.split() if word.lower() not in stopwords and len(word) > 2]
-    
+
     for i, doc in enumerate(documents[:5], 1):
         content = doc.page_content if hasattr(doc, 'page_content') else str(doc)
         source = doc.metadata.get('source', 'Unknown') if hasattr(doc, 'metadata') else 'Unknown'
-        
+    
         # Truncate long content
         if len(content) > 500:
             content = content[:500] + "..."
@@ -114,7 +114,7 @@ def format_document_context(documents: List, question: str = "") -> str:
             import re
             pattern = re.compile(re.escape(term), re.IGNORECASE)
             highlighted_content = pattern.sub(f"**{term}**", highlighted_content)
-        
+    
         formatted.append(f"""
 <details>
 <summary>📄 Chunk {i} - {os.path.basename(source)}</summary>
@@ -123,10 +123,10 @@ def format_document_context(documents: List, question: str = "") -> str:
 
 </details>
 """)
-    
+
     if len(documents) > 5:
         formatted.append(f"\n*... and {len(documents) - 5} more chunks*")
-    
+
     return "\n".join(formatted)
 
 
@@ -186,16 +186,16 @@ def main():
     _ensure_hfhub_hffolder_compat()  # must run before importing gradio
     import gradio as gr
     _setup_gradio_shim()
-    
+
     logger.info("=" * 60)
     logger.info("Starting SmartDoc AI application...")
     logger.info("=" * 60)
-    
+
     # Initialize components
     processor = DocumentProcessor()
     retriever_indexer = RetrieverBuilder()
     orchestrator = AgentWorkflow()
-    
+
     logger.info("All components initialized successfully")
 
     # CSS styling - Clean, accessible light theme with professional colors
@@ -205,7 +205,7 @@ def main():
         background: linear-gradient(180deg, #f8fafc 0%, #e2e8f0 100%) !important;
         font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
     }
-    
+
     /* Title styles - Dark text for readability */
     .app-title { 
         font-size: 2.2em !important; 
@@ -226,7 +226,7 @@ def main():
         font-size: 0.95em !important;
         line-height: 1.6 !important;
     }
-    
+
     /* Section headers */
     .section-header {
         color: #1e293b !important;
@@ -235,7 +235,7 @@ def main():
         padding-bottom: 8px !important;
         margin-bottom: 16px !important;
     }
-    
+
     /* Chat history panel - Clean white card with more height */
     .chat-history { 
         min-height: 500px;
@@ -259,7 +259,7 @@ def main():
     .chat-history strong {
         color: #1e293b !important;
     }
-    
+
     /* Document context panel */
     .doc-context { 
         max-height: 380px; 
@@ -286,7 +286,7 @@ def main():
     .doc-context p, .doc-context span {
         color: #475569 !important;
     }
-    
+
     /* Answer box - Success green accent, auto-height */
     .answer-box > div:nth-child(2) { 
         border-left: 4px solid #10b981 !important; 
@@ -317,7 +317,7 @@ def main():
         border-radius: 6px !important;
         overflow-x: auto !important;
     }
-    
+
     /* Verification box - Blue accent */
     .verification-box >  div:nth-child(2) {
         border-left: 4px solid #0ea5e9 !important;
@@ -333,7 +333,7 @@ def main():
     .verification-box strong {
         color: #075985 !important;
     }
-    
+
     /* Stats panel - Professional blue gradient */
     .stats-panel { 
         background: linear-gradient(135deg, #0369a1 0%, #0284c7 50%, #0ea5e9 100%) !important; 
@@ -346,7 +346,7 @@ def main():
     .stats-panel strong {
         color: #ffffff !important;
     }
-    
+
     /* Info panel */
     .info-panel {
         background: #eff6ff !important;
@@ -355,7 +355,7 @@ def main():
         padding: 12px !important;
         color: #1e40af !important;
     }
-    
+
     /* Form elements */
     .gr-input, .gr-textbox textarea {
         background: #ffffff !important;
@@ -367,13 +367,13 @@ def main():
         border-color: #0ea5e9 !important;
         box-shadow: 0 0 0 3px rgba(14, 165, 233, 0.1) !important;
     }
-    
+
     /* Labels */
     label {
         color: #374151 !important;
         font-weight: 500 !important;
     }   
-    
+
     /* Dropdown - High contrast with darker background for visibility */
     .gr-dropdown, 
     [data-testid="dropdown"],
@@ -398,7 +398,7 @@ def main():
         background: transparent !important;
         font-weight: 500 !important;
     }
-    
+
     /* Dropdown container and options */
     [data-testid="dropdown"] span, 
     .dropdown-container span,
@@ -406,7 +406,7 @@ def main():
         color: #1e293b !important;
         font-weight: 500 !important;
     }
-    
+
     /* Dropdown list options */
     .gr-dropdown ul, 
     .dropdown-options,
@@ -427,14 +427,14 @@ def main():
         background: #c7d2fe !important;
         color: #1e40af !important;
     }
-    
+
     /* Dropdown label */
     .gr-dropdown label,
     [data-testid="dropdown"] label {
         color: #1e40af !important;
         font-weight: 600 !important;
     }
-    
+
     /* Tabs - Clean styling */
     .tab-nav {
         border-bottom: 2px solid #e2e8f0 !important;
@@ -451,7 +451,7 @@ def main():
         border-bottom: 3px solid #0369a1 !important;
         font-weight: 600 !important;
     }
-    
+
     /* Markdown text */
     .prose, .markdown-text {
         color: #334155 !important;
@@ -463,7 +463,7 @@ def main():
     .prose strong, .markdown-text strong {
         color: #0f172a !important;
     }
-    
+
     /* Scrollbar styling */
     ::-webkit-scrollbar {
         width: 8px;
@@ -495,7 +495,7 @@ def main():
         background: #1d4ed8 !important;
         box-shadow: 0 4px 10px rgba(30, 64, 175, 0.4) !important;
     }         
-    
+
     /* Left side input boxes with borders */
     .left-panel-box {
         background: #fafafa !important;
@@ -508,7 +508,7 @@ def main():
         border-color: #64748b !important;
         box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1) !important;
     }
-    
+
     /* File upload box with border */
     .file-upload-box {
         background: #f8fafc !important;
@@ -521,7 +521,7 @@ def main():
         border-style: solid !important;
         background: #f0f9ff !important;
     }
-    
+
     /* Question input box with border */
     .question-box {
         background: #fffbeb !important;
@@ -631,7 +631,7 @@ setInterval(tick, 500);
     # Launch server - Compatible with both local and Hugging Face Spaces
     # HF Spaces sets SPACE_ID environment variable
     is_hf_space = os.environ.get("SPACE_ID") is not None
-    
+
     with gr.Blocks(title="SmartDoc AI") as demo:
         gr.Markdown("### SmartDoc AI - Document Q&A", elem_classes="app-title")
         gr.Markdown("Upload your documents and ask questions. Answers will appear below, just like a chat.", elem_classes="app-description")
@@ -683,7 +683,7 @@ setInterval(tick, 500);
             )
             try:
                 if not question_text.strip():
-                   
+               
                     chat_history.append({"role": "user", "content": question_text})
                     chat_history.append({"role": "assistant", "content": "Please enter a question."})
                     yield (
@@ -698,7 +698,7 @@ setInterval(tick, 500);
                     )
                     return
                 if not uploaded_files:                         
-                  
+              
                     chat_history.append({"role": "user", "content": question_text})
                     chat_history.append({"role": "assistant", "content": "Please upload at least one document."})
                     yield (
@@ -824,7 +824,7 @@ setInterval(tick, 500);
                 verification = result.get("verification_report", "No verification details available.")
                 logger.info(f"Verification (internal):\n{verification}")
                 # Do not display verification to user, only use internally
-               
+           
                 chat_history.append({"role": "user", "content": question_text})
                 chat_history.append({"role": "assistant", "content": f"**Answer:**\n{answer}"})
                 session_state.value["last_documents"] = retriever.invoke(question_text)
@@ -853,8 +853,8 @@ setInterval(tick, 500);
                 )
             except Exception as e:
                 logger.error(f"Processing error: {e}", exc_info=True)
-               
-               
+           
+           
                 chat_history.append({"role": "user", "content": question_text})
                 chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
                 yield (
@@ -897,37 +897,94 @@ setInterval(tick, 500);
             ex_data = EXAMPLES[example_key]
             question_text = ex_data["question"]
             file_names = ex_data["file_paths"]
-            
+        
             # Try to download from HF dataset if on Spaces
             if is_hf_space:
                 try:
-                    from huggingface_hub import hf_hub_download
+                    from datasets import load_dataset
+                    import tempfile
+                    
                     copied_files = []
                     file_info_text = f"✅ Loaded: {example_key}\n\n"
                     
-                    # Get HF token for private dataset access (optional)
+                    # Get HF token (optional for public datasets)
                     hf_token = os.environ.get("HF_TOKEN", None)
                     
-                    for file_path in file_names:
-                        filename = os.path.basename(file_path)
-                        try:
-                            local_path = hf_hub_download(
-                                repo_id="TilanB/smartdoc-samples",  # Correct dataset repo
-                                repo_type="dataset",
-                                filename=filename,
-                                token=hf_token,  # Pass token for private repos
-                            )
-                            copied_files.append(local_path)
-                            file_size_mb = os.path.getsize(local_path) / (1024 * 1024)
-                            file_info_text += f"📄 {filename} ({file_size_mb:.2f} MB)\n"
-                        except Exception as e:
-                            logger.error(f"Failed to download {filename}: {e}")
-                            file_info_text += f"❌ {filename} - Download failed\n"
-                    if not copied_files:
-                        return [], "", "❌ Could not load example files from dataset. Make sure the dataset is public or HF_TOKEN is set."
-                    return copied_files, question_text, file_info_text
-                except ImportError:
-                    return [], "", "❌ huggingface_hub not installed"
+                    try:
+                        # Load dataset - uses row-based structure
+                        logger.info(f"Loading dataset from HuggingFace: TilanB/smartdoc-samples")
+                        ds = load_dataset(
+                            "TilanB/smartdoc-samples",
+                            split="train",
+                            token=hf_token
+                        )
+                        logger.info(f"Dataset loaded with {len(ds)} rows")
+                        
+                        # Create temp directory for files
+                        temp_dir = tempfile.mkdtemp(prefix='hf_examples_')
+                        
+                        # Extract requested files from dataset rows
+                        for file_path in file_names:
+                            filename = os.path.basename(file_path)
+                            file_found = False
+                            
+                            # Search through dataset rows
+                            for row in ds:
+                                # Check if this row contains our file
+                                # Adjust field names based on your dataset structure
+                                row_filename = row.get('filename') or row.get('name') or row.get('path', '')
+                                
+                                if os.path.basename(row_filename) == filename:
+                                    temp_file_path = os.path.join(temp_dir, filename)
+                                    
+                                    # Handle different dataset column formats
+                                    if 'content' in row and row['content']:
+                                        # Binary content stored directly
+                                        with open(temp_file_path, 'wb') as f:
+                                            f.write(row['content'])
+                                    elif 'file' in row and row['file']:
+                                        # File object with bytes
+                                        file_obj = row['file']
+                                        if isinstance(file_obj, dict) and 'bytes' in file_obj:
+                                            with open(temp_file_path, 'wb') as f:
+                                                f.write(file_obj['bytes'])
+                                        elif isinstance(file_obj, bytes):
+                                            with open(temp_file_path, 'wb') as f:
+                                                f.write(file_obj)
+                                    elif 'data' in row and row['data']:
+                                        # Raw data field
+                                        with open(temp_file_path, 'wb') as f:
+                                            f.write(row['data'])
+                                    else:
+                                        logger.warning(f"Unknown dataset format for {filename}, available fields: {list(row.keys())}")
+                                        continue
+                                    
+                                    copied_files.append(temp_file_path)
+                                    file_size_mb = os.path.getsize(temp_file_path) / (1024 * 1024)
+                                    file_info_text += f"📄 {filename} ({file_size_mb:.2f} MB)\n"
+                                    file_found = True
+                                    logger.info(f"Successfully extracted {filename} from dataset")
+                                    break
+                            
+                            if not file_found:
+                                logger.warning(f"File {filename} not found in dataset rows")
+                                file_info_text += f"⚠️ {filename} - Not found in dataset\n"
+                        
+                        if not copied_files:
+                            # Log dataset structure for debugging
+                            if len(ds) > 0:
+                                logger.error(f"Dataset structure: {list(ds[0].keys())}")
+                            return [], "", f"❌ Could not find example files in dataset.\n\nDataset has {len(ds)} rows. Please check dataset structure or upload files manually."
+                        
+                        return copied_files, question_text, file_info_text
+                        
+                    except Exception as e:
+                        logger.error(f"Failed to load dataset: {e}", exc_info=True)
+                        return [], "", f"❌ Failed to load dataset: {str(e)}\n\nPlease upload files manually."
+                
+                except ImportError as e:
+                    logger.error(f"datasets package not installed: {e}")
+                    return [], "", "❌ 'datasets' package not installed"
             else:
                 # Local mode - use files from samples directory
                 import tempfile
@@ -948,17 +1005,17 @@ setInterval(tick, 500);
                 if not copied_files:
                     return [], "", "Could not load example files"
                 return copied_files, question_text, file_info_text
-        
+    
         example_dropdown.change(
             fn=load_example,
             inputs=[example_dropdown],
             outputs=[files, question, loaded_file_info]
         )
-        
+    
         # Show loaded_file_info when example is selected
         def show_info(example_key):
             return gr.update(visible=bool(example_key))
-        
+    
         example_dropdown.change(
             fn=show_info,
             inputs=[example_dropdown],
@@ -967,7 +1024,7 @@ setInterval(tick, 500);
     # Launch server - Compatible with both local and Hugging Face Spaces
     # HF Spaces sets SPACE_ID environment variable
     is_hf_space = os.environ.get("SPACE_ID") is not None
-    
+
     demo.queue()
     if is_hf_space:
         # Hugging Face Spaces configuration