# First import our patch before any gradio imports try: from patch_gradio import patched_get_type except ImportError: pass # Now import the rest of the packages import gradio as gr from pathlib import Path from tempfile import mkdtemp # LangChain & Embedding/LLM from langchain_community.vectorstores import FAISS from langchain_huggingface.embeddings import HuggingFaceEmbeddings from langchain_huggingface import HuggingFaceEndpoint from langchain.chains import create_retrieval_chain from langchain.chains.combine_documents import create_stuff_documents_chain from langchain_core.prompts import PromptTemplate from langchain_text_splitters import RecursiveCharacterTextSplitter # Custom Data Extractor from data_extraction import DataExtractor # Constants EMBED_MODEL_ID = "BAAI/bge-large-en-v1.5" GEN_MODEL_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1" TOP_K = 5 TEMP_DIR = Path(mkdtemp()) milvus_uri = str(TEMP_DIR / "docling.db") # Initialize extractor = DataExtractor() embedding = HuggingFaceEmbeddings(model_name=EMBED_MODEL_ID) llm = HuggingFaceEndpoint(repo_id=GEN_MODEL_ID, task="text-generation") template = """You are a helpful assistant. Based on the following context, answer the user's query. If you don't know the answer, just say that you don't know. Don't try to make up an answer. Use the given context strictly to answer the question. {context} Question: {input} Helpful Answer:""" rag_prompt = PromptTemplate.from_template(template) # App State vectorstore = None retriever = None # Store the retriever separately rag_chain = None # Function to process using direct extraction methods def process_direct(file_path): text = "" if file_path.suffix.lower() == ".pdf": text = extractor.extract_text_from_pdf(file_path) elif file_path.suffix.lower() == ".docx": text = extractor.extract_text_from_doc(file_path) elif file_path.suffix.lower() == ".txt": text = extractor.extract_text_from_txt(file_path) else: return None, "Unsupported file format. Please use PDF, DOCX, or TXT." # Clean text text = extractor.clean_text(text) # Split text using RecursiveCharacterTextSplitter text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, length_function=len ) from langchain_core.documents import Document doc = Document( page_content=text, metadata={ "source": str(file_path), "filename": file_path.name } ) return text_splitter.split_documents([doc]), "Direct extraction completed successfully." # Function to process using docling def process_docling(file_path): try: if file_path.suffix.lower() == ".pdf": # Convert PDF to markdown markdown_path = TEMP_DIR / "temp.md" extractor.pdf_to_markdown(file_path, markdown_path) # Load markdown markdown_raw = extractor.load_markdown_file(markdown_path) elif file_path.suffix.lower() == ".md": markdown_raw = extractor.load_markdown_file(file_path) else: return None, "Only PDF or Markdown files are supported for docling extraction." # Update metadata with filename for doc in markdown_raw: doc.metadata.update({ "filename": file_path.name }) # Split markdown splits = extractor.spit_markdown(markdown_raw) # Update metadata for all splits for split in splits: split.metadata.update({ "filename": file_path.name }) return splits, "Docling extraction completed successfully." except Exception as e: return None, f"Error during docling extraction: {str(e)}" # Step 1: File Upload & Index def process_file(file, extraction_method): global vectorstore, retriever, rag_chain if file is None: return "Please upload a file first." file_path = Path(file.name) # Process based on selected extraction method if extraction_method == "Direct Extraction": data_splits, message = process_direct(file_path) else: # Docling Extraction data_splits, message = process_docling(file_path) if data_splits is None: return message # Create vector store vectorstore = FAISS.from_documents(documents=data_splits, embedding=embedding) # Create retriever and chain retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K}) question_answer_chain = create_stuff_documents_chain(llm, prompt=rag_prompt) rag_chain = create_retrieval_chain(retriever, question_answer_chain) # Count chunks created chunk_count = len(data_splits) return f"File processed using {extraction_method}. {message} Created {chunk_count} document chunks. You can now ask questions." # Format source information def format_sources(docs): source_info = [] seen_sources = set() for i, doc in enumerate(docs, 1): # Extract source information source = doc.metadata.get("source", "Unknown") filename = doc.metadata.get("filename", Path(source).name if source != "Unknown" else "Unknown") # Get header information if available (from docling extraction) header_info = [] for level in range(1, 4): header = doc.metadata.get(f"Header_{level}", "") if header: header_info.append(header) # Create a unique identifier for this source source_id = f"{filename}{'_' + '_'.join(header_info) if header_info else ''}" if source_id not in seen_sources: seen_sources.add(source_id) # Format the source information source_text = f"**Source {i}: {filename}**" if header_info: source_text += f"\n - Section: {' > '.join(header_info)}" # Add the full text content without truncation source_text += f"\n - Content: \"{doc.page_content}\"" source_info.append(source_text) return "\n\n".join(source_info) # Step 2: Chat with RAG def chat_with_doc(query): if not rag_chain: return "Upload and process a document first.", "" # Use the stored retriever directly to get documents if not retriever: return "Retriever not initialized. Please process a document first.", "" # Get documents directly from retriever retrieved_docs = retriever.get_relevant_documents(query) # Then invoke the chain with the query result = rag_chain.invoke({"input": query}) answer = result.get("answer", "No answer found.") # Format source information source_text = format_sources(retrieved_docs) return answer, source_text # Graceful error handling wrapper for Gradio functions def safe_process_file(file, extraction_method): try: return process_file(file, extraction_method) except Exception as e: import traceback error_trace = traceback.format_exc() print(f"Error processing file: {str(e)}\n{error_trace}") return f"Error processing file: {str(e)}" def safe_chat_with_doc(query): try: return chat_with_doc(query) except Exception as e: import traceback error_trace = traceback.format_exc() print(f"Error processing query: {str(e)}\n{error_trace}") return f"Error processing query: {str(e)}", "" # Gradio Interface with gr.Blocks() as app: gr.Markdown("# Document Q&A System") with gr.Tab("Upload Document"): file_input = gr.File(label="Upload File", file_types=[".pdf", ".md", ".docx", ".txt"]) extraction_method = gr.Radio( ["Direct Extraction", "Docling Extraction"], label="Extraction Method", value="Direct Extraction", info="Direct extraction is faster but simpler. Docling extraction preserves document structure better." ) process_btn = gr.Button("Process Document") file_output = gr.Textbox(label="Processing Status", interactive=False) process_btn.click( fn=safe_process_file, # Use the wrapped safe function inputs=[file_input, extraction_method], outputs=file_output ) with gr.Tab("Chat with Document"): chat_input = gr.Textbox(label="Ask a Question", placeholder="Type your question...") ask_btn = gr.Button("Ask") chat_output = gr.Textbox(label="Answer", interactive=False) sources_output = gr.Textbox(label="Sources", interactive=False) ask_btn.click( fn=safe_chat_with_doc, # Use the wrapped safe function inputs=chat_input, outputs=[chat_output, sources_output] ) chat_input.submit( fn=safe_chat_with_doc, # Use the wrapped safe function inputs=chat_input, outputs=[chat_output, sources_output] ) if __name__ == "__main__": app.launch(debug=True)