app.py

# 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>
{context}
</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)