codeNJ/Instinctive-Studio-assessment
0
1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from typing import List, Optional, Dict, Any4import zipfile5import json6import os7import sqlite38import PyPDF29import re10from sentence_transformers import SentenceTransformer11import faiss12import numpy as np13from sklearn.feature_extraction.text import TfidfVectorizer14from sklearn.metrics.pairwise import cosine_similarity15import nltk16from nltk.tokenize import sent_tokenize17import time18import traceback19import uvicorn20import gradio as gr21from fastapi.middleware.cors import CORSMiddleware22from fastapi.responses import HTMLResponse, JSONResponse23import asyncio24from contextlib import asynccontextmanager25import requests26import urllib.parse27from pathlib import Path28 29# Download NLTK data for sentence tokenization30nltk.download('punkt', quiet=True)31 32# Load sources.json if it exists33try:34 with open('sources.json', 'r') as f:35 sources = json.load(f)36except FileNotFoundError:37 print("sources.json not found, using empty sources list")38 sources = []39 40# Pydantic models for request/response41class AskRequest(BaseModel):42 q: str43 k: int = 544 mode: str = "answer" # "answer" or "search"45 46class Context(BaseModel):47 text: str48 score: float49 title: str50 url: str51 vector_score: float52 keyword_score: float53 title_score: float54 hybrid_score: float55 cosine_score: float56 exact_matches: int57 58class AskResponse(BaseModel):59 answer: Optional[str] = None60 contexts: List[Context]61 reranker_used: bool = True62 63# Global system instance64safety_system = None65 66@asynccontextmanager67async def lifespan(app: FastAPI):68 # Startup69 global safety_system70 print("Setting up industrial safety retrieval system...")71 safety_system = IndustrialSafetyRetrievalSystem()72 print("Downloading and processing documents...")73 74 # Run processing in a separate thread to avoid blocking75 import threading76 def process_docs():77 safety_system.download_and_process_documents()78 print("System ready!")79 80 thread = threading.Thread(target=process_docs)81 thread.daemon = True82 thread.start()83 84 yield # App runs here85 86 # Shutdown87 # Clean up resources if needed88 print("Shutting down...")89 90app = FastAPI(91 title="Industrial Safety Retrieval API",92 description="API for retrieving industrial safety information from PDF documents",93 lifespan=lifespan94)95 96# Add CORS middleware97app.add_middleware(98 CORSMiddleware,99 allow_origins=["*"],100 allow_credentials=True,101 allow_methods=["*"],102 allow_headers=["*"],103)104 105class IndustrialSafetyRetrievalSystem:106 def __init__(self):107 print("Loading sentence transformer model...")108 self.model = SentenceTransformer('all-MiniLM-L6-v2')109 print("Model loaded successfully!")110 self.vectorizer = TfidfVectorizer(stop_words='english')111 self.db_path = 'safety_docs.db'112 self.chunk_id_mapping = [] # Maps FAISS index to database chunk ID113 self.setup_database()114 115 def get_db_connection(self):116 """Get a new database connection with proper settings"""117 conn = sqlite3.connect(self.db_path)118 conn.execute("PRAGMA busy_timeout = 30000")119 return conn120 121 def setup_database(self):122 """Create database tables for chunks and metadata"""123 print("Setting up database...")124 conn = self.get_db_connection()125 cursor = conn.cursor()126 127 cursor.execute('''128 CREATE TABLE IF NOT EXISTS sources (129 id INTEGER PRIMARY KEY,130 title TEXT,131 url TEXT,132 file_path TEXT133 )134 ''')135 136 cursor.execute('''137 CREATE TABLE IF NOT EXISTS chunks (138 id INTEGER PRIMARY KEY,139 source_id INTEGER,140 chunk_text TEXT,141 page_number INTEGER,142 chunk_index INTEGER,143 FOREIGN KEY (source_id) REFERENCES sources (id)144 )145 ''')146 147 cursor.execute('''148 CREATE TABLE IF NOT EXISTS embeddings (149 chunk_id INTEGER PRIMARY KEY,150 embedding BLOB,151 FOREIGN KEY (chunk_id) REFERENCES chunks (id)152 )153 ''')154 155 conn.commit()156 conn.close()157 print("Database setup complete!")158 159 def clear_database(self):160 """Clear existing data from database"""161 max_retries = 5162 for attempt in range(max_retries):163 try:164 conn = self.get_db_connection()165 cursor = conn.cursor()166 167 # Disable foreign keys temporarily for faster deletion168 cursor.execute("PRAGMA foreign_keys = OFF")169 170 # Clear tables in reverse order to avoid foreign key constraints171 cursor.execute("DELETE FROM embeddings")172 cursor.execute("DELETE FROM chunks")173 cursor.execute("DELETE FROM sources")174 175 # Re-enable foreign keys176 cursor.execute("PRAGMA foreign_keys = ON")177 178 conn.commit()179 conn.close()180 print("Database cleared successfully")181 182 # Clear the mapping as well183 self.chunk_id_mapping = []184 return True185 186 except sqlite3.OperationalError as e:187 print(f"Database locked (attempt {attempt + 1}/{max_retries}), retrying...")188 time.sleep(2) # Wait before retrying189 if attempt == max_retries - 1:190 print(f"Failed to clear database after {max_retries} attempts: {e}")191 return False192 193 def download_pdfs_from_sources(self):194 """Download PDFs from sources.json and save to pdf_docs directory"""195 pdf_dir = 'pdf_docs'196 os.makedirs(pdf_dir, exist_ok=True)197 198 downloaded_files = []199 200 for source in sources:201 url = source.get('url', '')202 title = source.get('title', 'Unknown')203 204 if not url.lower().endswith('.pdf'):205 print(f"Skipping non-PDF source: {title} ({url})")206 continue207 208 try:209 # Create a safe filename from the title210 safe_title = "".join(c for c in title if c.isalnum() or c in (' ', '-', '_')).rstrip()211 safe_title = safe_title.replace(' ', '_')[:100] # Limit length212 filename = f"{safe_title}.pdf"213 file_path = os.path.join(pdf_dir, filename)214 215 # Check if file already exists216 if os.path.exists(file_path):217 print(f"File already exists: {filename}")218 downloaded_files.append({219 'file_path': file_path,220 'file_name': filename,221 'title': title,222 'url': url223 })224 continue225 226 print(f"Downloading: {title} from {url}")227 228 # Download the PDF229 response = requests.get(url, timeout=30)230 response.raise_for_status()231 232 # Save the PDF233 with open(file_path, 'wb') as f:234 f.write(response.content)235 236 print(f"Downloaded: {filename} ({len(response.content)} bytes)")237 238 downloaded_files.append({239 'file_path': file_path,240 'file_name': filename,241 'title': title,242 'url': url243 })244 245 except Exception as e:246 print(f"Error downloading {title} from {url}: {e}")247 248 return downloaded_files249 250 def extract_text_from_pdf(self, file_path):251 """Extract text from PDF file"""252 text = ""253 try:254 print(f" Opening PDF: {os.path.basename(file_path)}")255 with open(file_path, 'rb') as file:256 pdf_reader = PyPDF2.PdfReader(file)257 258 if pdf_reader.is_encrypted:259 print(" PDF is encrypted, trying to decrypt...")260 try:261 pdf_reader.decrypt('')262 print(" PDF decrypted successfully!")263 except:264 print(" Could not decrypt PDF")265 return text266 267 total_pages = len(pdf_reader.pages)268 print(f" Total pages: {total_pages}")269 270 successful_pages = 0271 for page_num, page in enumerate(pdf_reader.pages):272 try:273 page_text = page.extract_text()274 if page_text and page_text.strip():275 text += f"\n--- PAGE {page_num + 1} ---\n"276 text += page_text277 successful_pages += 1278 else:279 text += f"\n--- PAGE {page_num + 1} (No text content) ---\n"280 except Exception as page_error:281 print(f" Error on page {page_num + 1}: {page_error}")282 text += f"\n--- PAGE {page_num + 1} (Error reading page) ---\n"283 284 print(f" Successfully extracted text from {successful_pages}/{total_pages} pages")285 print(f" Total text length: {len(text)} characters")286 287 except Exception as e:288 print(f" Error reading PDF: {e}")289 print(f" Traceback: {traceback.format_exc()}")290 return ""291 292 return text293 294 def chunk_text(self, text, max_chunk_size=500):295 """Split text into sensible chunks"""296 chunks = []297 298 if not text or not text.strip():299 print(" No text to chunk")300 return chunks301 302 # Clean the text303 text = re.sub(r'--- PAGE \d+ ---', '', text)304 text = re.sub(r'\(No text content\)', '', text)305 text = re.sub(r'\(Error reading page\)', '', text)306 text = text.strip()307 308 if not text:309 print(" Text is empty after cleaning")310 return chunks311 312 paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]313 print(f" Found {len(paragraphs)} paragraphs")314 315 for paragraph in paragraphs:316 if len(paragraph) <= max_chunk_size:317 chunks.append(paragraph)318 else:319 try:320 sentences = sent_tokenize(paragraph)321 current_chunk = ""322 for sentence in sentences:323 if len(current_chunk) + len(sentence) <= max_chunk_size:324 current_chunk += " " + sentence if current_chunk else sentence325 else:326 if current_chunk:327 chunks.append(current_chunk)328 current_chunk = sentence329 if current_chunk:330 chunks.append(current_chunk)331 except:332 # Fallback: split by words333 words = paragraph.split()334 current_chunk = []335 current_length = 0336 337 for word in words:338 if current_length + len(word) + 1 <= max_chunk_size:339 current_chunk.append(word)340 current_length += len(word) + 1341 else:342 if current_chunk:343 chunks.append(" ".join(current_chunk))344 current_chunk = [word]345 current_length = len(word)346 347 if current_chunk:348 chunks.append(" ".join(current_chunk))349 350 print(f" Created {len(chunks)} chunks")351 if chunks:352 print(f" Sample chunk: {chunks[0][:100]}...")353 return chunks354 355 def get_pdf_files(self):356 """Get all PDF files from the pdf_docs directory"""357 pdf_files = []358 pdf_dir = 'pdf_docs'359 360 if not os.path.exists(pdf_dir):361 print(f"Error: Directory '{pdf_dir}' not found!")362 return pdf_files363 364 for file_name in os.listdir(pdf_dir):365 if file_name.endswith('.pdf'):366 file_path = os.path.join(pdf_dir, file_name)367 368 # Get source info from sources.json369 title, url = self.get_source_info(file_name)370 371 pdf_files.append({372 'file_path': file_path,373 'file_name': file_name,374 'title': title,375 'url': url376 })377 378 return pdf_files379 380 def debug_database_state(self):381 """Debug function to check database state"""382 conn = self.get_db_connection()383 cursor = conn.cursor()384 385 print("\n=== DATABASE DEBUG INFO ===")386 387 # Check sources388 cursor.execute("SELECT COUNT(*) FROM sources")389 source_count = cursor.fetchone()[0]390 print(f"Sources in database: {source_count}")391 392 # Check chunks393 cursor.execute("SELECT COUNT(*) FROM chunks")394 chunk_count = cursor.fetchone()[0]395 print(f"Chunks in database: {chunk_count}")396 397 # Check embeddings398 cursor.execute("SELECT COUNT(*) FROM embeddings")399 embedding_count = cursor.fetchone()[0]400 print(f"Embeddings in database: {embedding_count}")401 402 # Check mapping403 print(f"Chunk ID mapping length: {len(self.chunk_id_mapping)}")404 405 # Show sample data406 if chunk_count > 0:407 cursor.execute("SELECT chunk_text FROM chunks LIMIT 1")408 sample_chunk = cursor.fetchone()409 if sample_chunk:410 print(f"Sample chunk text: {sample_chunk[0][:100]}...")411 412 conn.close()413 print("===========================\n")414 415 def get_source_info(self, file_name):416 """Get proper source info from sources.json"""417 # Try to match the filename with sources.json entries418 file_key = file_name.replace('.pdf', '').replace('_', ' ').replace('-', ' ').lower()419 420 for source in sources:421 source_title = source.get('title', '').lower()422 source_url = source.get('url', 'unknown')423 424 # Check if filename matches title or URL425 if (file_key in source_title or 426 any(word in source_title for word in file_key.split() if len(word) > 3) or427 file_name in source_url):428 return source.get('title', file_name), source.get('url', 'unknown')429 430 # Fallback to filename431 return file_name.replace('.pdf', ''), 'unknown'432 433 def download_and_process_documents(self):434 """Download PDFs from sources.json and process them"""435 print("Starting document download and processing...")436 437 # Clear database first438 if not self.clear_database():439 print("Warning: Could not clear database, proceeding anyway...")440 441 # Download PDFs from sources.json442 pdf_files = self.download_pdfs_from_sources()443 444 # Also check for existing PDFs in the directory445 existing_pdfs = self.get_pdf_files()446 447 # Combine downloaded and existing files448 all_pdf_files = pdf_files + existing_pdfs449 450 # Remove duplicates451 seen_paths = set()452 unique_pdf_files = []453 for pdf in all_pdf_files:454 if pdf['file_path'] not in seen_paths:455 seen_paths.add(pdf['file_path'])456 unique_pdf_files.append(pdf)457 458 if not unique_pdf_files:459 print("No PDF files found to process!")460 return461 462 print(f"Found {len(unique_pdf_files)} PDF files to process")463 464 all_embeddings = []465 processed_files = 0466 467 # Process each PDF file468 for pdf_info in unique_pdf_files:469 file_path = pdf_info['file_path']470 file_name = pdf_info['file_name']471 title = pdf_info.get('title', '')472 url = pdf_info.get('url', 'unknown')473 474 if not os.path.exists(file_path):475 print(f"File not found: {file_path}")476 continue477 478 print(f"\n=== Processing: {file_name} ===")479 480 # Get a new database connection for this file481 conn = self.get_db_connection()482 cursor = conn.cursor()483 484 try:485 # If title wasn't provided, try to get it from sources.json486 if not title:487 title, url = self.get_source_info(file_name)488 489 # Insert source into database490 cursor.execute(491 "INSERT INTO sources (title, url, file_path) VALUES (?, ?, ?)",492 (title, url, file_path)493 )494 source_id = cursor.lastrowid495 496 # Extract and chunk text497 text = self.extract_text_from_pdf(file_path)498 if not text or not text.strip():499 print(f" No text extracted from {file_name}")500 conn.close()501 continue502 503 chunks = self.chunk_text(text)504 if not chunks:505 print(f" No chunks created from {file_name}")506 conn.close()507 continue508 509 # Store chunks and generate embeddings510 for chunk_idx, chunk_text in enumerate(chunks):511 cursor.execute(512 "INSERT INTO chunks (source_id, chunk_text, page_number, chunk_index) VALUES (?, ?, ?, ?)",513 (source_id, chunk_text, 1, chunk_idx)514 )515 chunk_id = cursor.lastrowid516 517 # Generate embedding518 embedding = self.model.encode([chunk_text])[0]519 520 # Store embedding in database521 cursor.execute(522 "INSERT INTO embeddings (chunk_id, embedding) VALUES (?, ?)",523 (chunk_id, embedding.tobytes())524 )525 526 # Add to our collections for FAISS527 all_embeddings.append(embedding)528 self.chunk_id_mapping.append(chunk_id) # Map FAISS index to chunk ID529 530 conn.commit()531 processed_files += 1532 print(f" Successfully processed {file_name}")533 534 except Exception as e:535 print(f" Error processing {file_name}:")536 print(f" Error: {e}")537 print(f" Traceback: {traceback.format_exc()}")538 conn.rollback()539 finally:540 conn.close()541 542 # Debug database state after processing543 self.debug_database_state()544 545 # Create FAISS index with all embeddings546 if all_embeddings:547 try:548 print(f"\nCreating FAISS index with {len(all_embeddings)} embeddings...")549 embeddings_array = np.array(all_embeddings).astype('float32')550 print(f"Embeddings array shape: {embeddings_array.shape}")551 print(f"Chunk ID mapping length: {len(self.chunk_id_mapping)}")552 553 # Normalize embeddings for better cosine similarity554 faiss.normalize_L2(embeddings_array)555 556 # Use Inner Product index for normalized vectors (equivalent to cosine similarity)557 self.faiss_index = faiss.IndexFlatIP(embeddings_array.shape[1])558 self.faiss_index.add(embeddings_array)559 print(f"FAISS index created with {self.faiss_index.ntotal} vectors (using cosine similarity)")560 561 except Exception as e:562 print(f"Error creating FAISS index: {e}")563 print(f"FAISS traceback: {traceback.format_exc()}")564 else:565 print("No embeddings were generated!")566 567 print(f"\nProcessing complete: {processed_files}/{len(unique_pdf_files)} files with {len(all_embeddings)} chunks")568 569 def search(self, query, top_k=10):570 """Search for relevant chunks with improved scoring"""571 print(f"\nSearching for: '{query}'")572 573 if not hasattr(self, 'faiss_index') or self.faiss_index.ntotal == 0:574 print("No FAISS index available!")575 return []576 577 print(f"FAISS index has {self.faiss_index.ntotal} vectors")578 print(f"Chunk ID mapping has {len(self.chunk_id_mapping)} entries")579 580 # Vector similarity search with normalized query581 query_embedding = self.model.encode([query])[0].astype('float32').reshape(1, -1)582 faiss.normalize_L2(query_embedding) # Normalize query for cosine similarity583 scores, indices = self.faiss_index.search(query_embedding, min(top_k * 3, self.faiss_index.ntotal))584 585 print(f"FAISS search returned {len(indices[0])} results")586 587 # Get chunk details using the mapping588 conn = self.get_db_connection()589 cursor = conn.cursor()590 results = []591 592 for faiss_idx, score in zip(indices[0], scores[0]):593 # Skip invalid indices594 if faiss_idx >= len(self.chunk_id_mapping) or faiss_idx < 0:595 print(f"Skipping invalid FAISS index: {faiss_idx}")596 continue597 598 # Get the actual chunk ID from our mapping599 chunk_id = self.chunk_id_mapping[faiss_idx]600 601 cursor.execute('''602 SELECT c.id, c.chunk_text, s.title, s.url603 FROM chunks c604 JOIN sources s ON c.source_id = s.id605 WHERE c.id = ?606 ''', (chunk_id,))607 608 row = cursor.fetchone()609 if row:610 chunk_id, chunk_text, title, url = row611 612 # Vector score is now cosine similarity (0 to 1, higher is better)613 vector_score = max(0, float(score)) # Ensure non-negative614 615 # Enhanced keyword scoring616 query_terms = [term.lower().strip() for term in query.split() if len(term.strip()) > 2]617 chunk_text_lower = chunk_text.lower()618 619 # Exact matches620 exact_matches = sum(1 for term in query_terms if term in chunk_text_lower)621 622 # Partial matches (stemming-like)623 partial_matches = 0624 for term in query_terms:625 if len(term) > 4: # Only for longer terms626 stem = term[:4] # Simple stemming627 if stem in chunk_text_lower and term not in chunk_text_lower:628 partial_matches += 0.5629 630 # Calculate keyword score with bonuses631 if query_terms:632 keyword_score = (exact_matches + partial_matches) / len(query_terms)633 # Bonus for multiple matches in close proximity634 if exact_matches > 1:635 keyword_score *= 1.2636 else:637 keyword_score = 0638 639 # Context relevance score (title matching)640 title_score = 0641 if title and query_terms:642 title_lower = title.lower()643 title_matches = sum(1 for term in query_terms if term in title_lower)644 title_score = title_matches / len(query_terms)645 646 # Combined hybrid score with better weighting647 # Increase vector weight since embeddings are generally more reliable648 hybrid_score = (0.6 * vector_score +649 0.3 * keyword_score +650 0.1 * title_score)651 652 # Boost score for very relevant chunks653 if keyword_score > 0.5: # High keyword match654 hybrid_score *= 1.15655 if vector_score > 0.8: # High semantic similarity656 hybrid_score *= 1.1657 658 # Ensure score doesn't exceed 1.0659 hybrid_score = min(hybrid_score, 1.0)660 661 results.append({662 'chunk_id': chunk_id,663 'chunk_text': chunk_text,664 'title': title,665 'url': url,666 'vector_score': vector_score,667 'keyword_score': keyword_score,668 'title_score': title_score,669 'hybrid_score': hybrid_score,670 'cosine_score': score,671 'exact_matches': exact_matches672 })673 else:674 print(f"No chunk found for ID {chunk_id}")675 676 conn.close()677 678 print(f"Found {len(results)} matching chunks")679 680 results.sort(key=lambda x: x['hybrid_score'], reverse=True)681 return results[:top_k]682 683 def answer_question(self, question, top_k=5, confidence_threshold=0.4):684 """Generate answer based on retrieved chunks with abstention logic"""685 print(f"\nAnswering question: '{question}'")686 relevant_chunks = self.search(question, top_k)687 688 if not relevant_chunks:689 return "I couldn't find any relevant information in the industrial safety documents to answer this question. Please try rephrasing your question or asking about a different safety topic."690 691 # Check if we should abstain based on confidence692 best_score = relevant_chunks[0]['hybrid_score']693 avg_score = sum(chunk['hybrid_score'] for chunk in relevant_chunks[:3]) / min(3, len(relevant_chunks))694 695 print(f"Best score: {best_score:.3f}, Average top-3 score: {avg_score:.3f}")696 697 if best_score < confidence_threshold:698 return f"I found some potentially related content (confidence: {best_score:.3f}), but I'm not confident it adequately addresses your question about '{question}'. The retrieved information may not be specific enough. Please try a more specific question or different keywords related to industrial safety."699 700 print(f"Found {len(relevant_chunks)} relevant chunks with good confidence")701 702 # Create grounded answer703 answer_parts = []704 citations = []705 seen_sources = set()706 707 # Use top 3 chunks for answer708 for i, chunk in enumerate(relevant_chunks[:3]):709 # Add chunk content710 chunk_text = chunk['chunk_text'].strip()711 if chunk_text:712 answer_parts.append(f"According to the safety documentation: {chunk_text}")713 714 # Add unique citations715 source_key = f"{chunk['title']}|{chunk['url']}"716 if source_key not in seen_sources:717 citations.append(f"[{len(citations) + 1}] {chunk['title']} - {chunk['url']}")718 seen_sources.add(source_key)719 720 if not answer_parts:721 return "I found potentially relevant documents but couldn't extract clear information to answer your question. Please try rephrasing with more specific safety-related terms."722 723 # Construct final answer724 answer = f"Based on the industrial safety documentation:\n\n"725 answer += "\n\n".join(answer_parts)726 727 if citations:728 answer += f"\n\nSources:\n" + "\n".join(citations)729 730 # Add confidence note for borderline cases731 if best_score < 0.6:732 answer += f"\n\nNote: Answer confidence is moderate ({best_score:.2f}). Please verify with additional sources if this is for critical safety decisions."733 734 return answer735 736@app.post("/ask", response_model=AskResponse)737async def ask_question(request: AskRequest):738 """Endpoint to ask questions about industrial safety"""739 global safety_system740 741 if safety_system is None:742 raise HTTPException(status_code=503, detail="System not initialized yet")743 744 try:745 if request.mode == "search":746 # Return search results only747 results = safety_system.search(request.q, request.k)748 contexts = [749 Context(750 text=result['chunk_text'],751 score=result['hybrid_score'],752 title=result['title'],753 url=result['url'],754 vector_score=result['vector_score'],755 keyword_score=result['keyword_score'],756 title_score=result['title_score'],757 hybrid_score=result['hybrid_score'],758 cosine_score=result['cosine_score'],759 exact_matches=result['exact_matches']760 )761 for result in results762 ]763 return AskResponse(answer=None, contexts=contexts, reranker_used=True)764 765 else: # answer mode766 # Generate answer with contexts767 results = safety_system.search(request.q, request.k)768 answer = safety_system.answer_question(request.q, request.k)769 770 contexts = [771 Context(772 text=result['chunk_text'],773 score=result['hybrid_score'],774 title=result['title'],775 url=result['url'],776 vector_score=result['vector_score'],777 keyword_score=result['keyword_score'],778 title_score=result['title_score'],779 hybrid_score=result['hybrid_score'],780 cosine_score=result['cosine_score'],781 exact_matches=result['exact_matches']782 )783 for result in results784 ]785 786 return AskResponse(answer=answer, contexts=contexts, reranker_used=True)787 788 except Exception as e:789 raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")790 791@app.get("/health")792async def health_check():793 """Health check endpoint"""794 global safety_system795 if safety_system is None:796 return {"status": "initializing"}797 798 if not hasattr(safety_system, 'faiss_index'):799 return {"status": "processing_documents"}800 801 return {802 "status": "ready",803 "documents_processed": len(safety_system.chunk_id_mapping) if safety_system else 0,804 "faiss_index_size": safety_system.faiss_index.ntotal if hasattr(safety_system, 'faiss_index') else 0805 }806 807@app.get("/")808async def root():809 """Root endpoint with API information"""810 return {811 "message": "Industrial Safety Retrieval API",812 "endpoints": {813 "POST /ask": "Ask questions about industrial safety",814 "GET /health": "Check system health status",815 "GET /": "This information page"816 },817 "parameters": {818 "q": "Question to ask",819 "k": "Number of results to return (default: 5)",820 "mode": "'answer' for full answer or 'search' for contexts only (default: 'answer')"821 }822 }823 824# Gradio Interface825def create_gradio_interface():826 """Create Gradio interface for the safety retrieval system"""827 828 def gradio_ask_question(question, num_results=5, mode="answer"):829 """Wrapper function for Gradio interface"""830 if safety_system is None:831 return "System not initialized yet. Please wait...", []832 833 if not hasattr(safety_system, 'faiss_index'):834 return "System is still processing documents. Please wait...", []835 836 try:837 if mode == "answer":838 answer = safety_system.answer_question(question, num_results)839 results = safety_system.search(question, num_results)840 contexts = [841 {842 "text": result['chunk_text'][:200] + "..." if len(result['chunk_text']) > 200 else result['chunk_text'],843 "score": round(result['hybrid_score'], 3),844 "title": result['title'],845 "url": result['url'],846 "confidence": f"{result['hybrid_score']:.3f}"847 }848 for result in results849 ]850 return answer, contexts851 else:852 results = safety_system.search(question, num_results)853 contexts = [854 {855 "text": result['chunk_text'][:200] + "..." if len(result['chunk_text']) > 200 else result['chunk_text'],856 "score": round(result['hybrid_score'], 3),857 "title": result['title'],858 "url": result['url'],859 "confidence": f"{result['hybrid_score']:.3f}"860 }861 for result in results862 ]863 return "Search results (showing top contexts):", contexts864 865 except Exception as e:866 return f"Error: {str(e)}", []867 868 with gr.Blocks(title="Industrial Safety Q&A", theme=gr.themes.Soft()) as demo:869 gr.Markdown("# 🛡️ Industrial Safety Document Q&A")870 gr.Markdown("Ask questions about industrial safety standards and procedures")871 872 with gr.Row():873 with gr.Column(scale=2):874 question = gr.Textbox(875 label="Question",876 placeholder="e.g., What are the electrical safety requirements?",877 lines=2,878 max_lines=4879 )880 881 with gr.Row():882 num_results = gr.Slider(883 minimum=1,884 maximum=10,885 value=5,886 step=1,887 label="Number of results"888 )889 mode = gr.Radio(890 choices=["answer", "search"],891 value="answer",892 label="Response Mode"893 )894 895 submit_btn = gr.Button("Ask Question", variant="primary")896 897 with gr.Column(scale=3):898 answer_output = gr.Textbox(899 label="Answer",900 lines=6,901 interactive=False902 )903 904 contexts_output = gr.DataFrame(905 label="Supporting Contexts",906 headers=["text", "score", "title", "url", "confidence"]907 )908 909 examples = gr.Examples(910 examples=[911 ["What are the electrical safety requirements?", 5, "answer"],912 ["How should lockout tagout procedures be implemented?", 3, "answer"],913 ["What PPE is required for chemical handling?", 4, "search"],914 ["What are the emergency procedures for fire incidents?", 5, "answer"]915 ],916 inputs=[question, num_results, mode],917 outputs=[answer_output, contexts_output],918 fn=gradio_ask_question,919 cache_examples=False,920 label="Example Queries"921 )922 923 submit_btn.click(924 fn=gradio_ask_question,925 inputs=[question, num_results, mode],926 outputs=[answer_output, contexts_output],927 api_name="ask_question"928 )929 930 question.submit(931 fn=gradio_ask_question,932 inputs=[question, num_results, mode],933 outputs=[answer_output, contexts_output]934 )935 936 return demo937 938gradio_app = create_gradio_interface()939app = gr.mount_gradio_app(app, gradio_app, path="/")940 941@app.get("/gradio")942async def gradio_redirect():943 return HTMLResponse("""944 <!DOCTYPE html>945 <html>946 <head>947 <meta http-equiv="refresh" content="0; url=/" />948 </head>949 <body>950 <p>Redirecting to <a href="/">Gradio Interface</a></p>951 </body>952 </html>953 """)954 955if __name__ == "__main__":956 uvicorn.run(app, host="0.0.0.0", port=7860)