import os print(os.system("apt-get update && apt-get install -y poppler-utils")) os.system("apt-get update && apt-get install -y tesseract-ocr") import tempfile import streamlit as st from PIL import Image import pytesseract # pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract" from pdf2image import convert_from_path import pypdf from dotenv import load_dotenv import time from langchain_core.messages import HumanMessage, AIMessage, SystemMessage from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.output_parsers import StrOutputParser from langchain_together import Together from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_community.vectorstores import FAISS from langchain_community.embeddings import HuggingFaceEmbeddings import langgraph from langgraph.graph import END from typing import List, Dict, Any, TypedDict, Optional # Load environment variables load_dotenv() # Set page configuration st.set_page_config( page_title="Document Q&A", page_icon="📚", layout="wide", initial_sidebar_state="expanded" ) # Custom CSS for better UI st.markdown(""" """, unsafe_allow_html=True) # Example questions EXAMPLE_QUESTIONS = [ "How do the different topics in these documents relate to each other?", "What is the structure of this document?", "Can you analyze the writing style of this text?", "Extract all dates and events mentioned in the document", "What are the main arguments presented in this document?" ] # Initialize the LLM @st.cache_resource def get_llm(): return Together( model="deepseek-ai/DeepSeek-V3", temperature=0.7, max_tokens=1024 ) # Initialize embeddings @st.cache_resource def get_embeddings(): return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # Initialize text splitter @st.cache_resource def get_text_splitter(): return RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200 ) # Function to extract text from PDF def extract_text_from_pdf(pdf_file): pdf_reader = pypdf.PdfReader(pdf_file) text = "" for page in pdf_reader.pages: text += page.extract_text() or "" return text # Function to extract text from image using OCR def extract_text_from_image(image_file): image = Image.open(image_file) text = pytesseract.image_to_string(image) return text # Function to process PDF with OCR if needed def process_pdf_with_ocr(pdf_file): # First try normal text extraction text = extract_text_from_pdf(pdf_file) # If little or no text was extracted, try OCR if len(text.strip()) < 100: images = convert_from_path(pdf_file) text = "" for image in images: text += pytesseract.image_to_string(image) return text # Function to process uploaded files def process_uploaded_files(uploaded_files): all_texts = [] file_info = [] for file in uploaded_files: # Create a temporary file with tempfile.NamedTemporaryFile(delete=False) as temp_file: temp_file.write(file.getvalue()) temp_file_path = temp_file.name # Process based on file type if file.name.lower().endswith('.pdf'): text = process_pdf_with_ocr(temp_file_path) file_type = "PDF" elif file.name.lower().endswith(('.png', '.jpg', '.jpeg')): text = extract_text_from_image(temp_file_path) file_type = "Image" elif file.name.lower().endswith(('.txt', '.md')): text = file.getvalue().decode('utf-8') file_type = "Text" else: text = f"Unsupported file format: {file.name}" file_type = "Unknown" all_texts.append(f"--- Content from {file.name} ---\n{text}") file_info.append({"name": file.name, "type": file_type}) # Clean up the temporary file os.unlink(temp_file_path) return "\n\n".join(all_texts), file_info # Function to create vector store from text def create_vectorstore(text): text_splitter = get_text_splitter() chunks = text_splitter.split_text(text) # Use FAISS instead of Chroma to avoid SQLite dependency return FAISS.from_texts( texts=chunks, embedding=get_embeddings() ) # Define the state schema for the graph using TypedDict class GraphState(TypedDict): messages: List documents: List thinking: str # Define the RAG agent using LangGraph def create_rag_agent(vectorstore): # Define the retrieval component def retrieve(state: GraphState) -> GraphState: query = state["messages"][-1].content docs = vectorstore.similarity_search(query, k=5) return {"documents": docs, "messages": state["messages"], "thinking": state.get("thinking", "")} # Define the generation component with thinking step def generate(state: GraphState) -> GraphState: messages = state["messages"] documents = state["documents"] # Extract relevant context from documents context = "\n\n".join([f"Document {i+1}:\n{doc.page_content}" for i, doc in enumerate(documents)]) # First, have the model think about the query thinking_prompt = ChatPromptTemplate.from_messages([ SystemMessage(content="You are an assistant that thinks step by step before answering."), MessagesPlaceholder(variable_name="messages"), SystemMessage(content=f"Here is relevant context from the knowledge base:\n{context}\n\nThink step by step about how to answer the query using this context.") ]) thinking = thinking_prompt | get_llm() | StrOutputParser() thinking_result = thinking.invoke({"messages": messages}) # Then generate the final answer answer_prompt = ChatPromptTemplate.from_messages([ SystemMessage(content="You are a helpful assistant that provides accurate information based on the given context."), MessagesPlaceholder(variable_name="messages"), SystemMessage(content=f"Here is relevant context from the knowledge base:\n{context}\n\nHere is your thinking process:\n{thinking_result}\n\nNow provide a clear and helpful answer based on this context and thinking.") ]) answer = answer_prompt | get_llm() | StrOutputParser() response = answer.invoke({"messages": messages}) return { "messages": messages + [AIMessage(content=response)], "thinking": thinking_result, "documents": documents } # Create the graph from langgraph.graph import StateGraph workflow = StateGraph(GraphState) workflow.add_node("retrieve", retrieve) workflow.add_node("generate", generate) workflow.set_entry_point("retrieve") workflow.add_edge("retrieve", "generate") workflow.add_edge("generate", END) # Compile the graph app = workflow.compile() return app # Function to clear all session state def clear_session_state(): for key in list(st.session_state.keys()): del st.session_state[key] # Main app layout def main(): # Initialize session state for showing examples if "show_examples" not in st.session_state: st.session_state.show_examples = True # Initialize messages if not exists if "messages" not in st.session_state: st.session_state.messages = [] # Initialize thinking history if not exists if "thinking_history" not in st.session_state: st.session_state.thinking_history = [] # Sidebar for document upload and controls with st.sidebar: st.markdown('
📚 Document Q&A
', unsafe_allow_html=True) st.markdown("""
  1. Upload your documents using the form below
  2. Process the documents
  3. Ask questions about your documents
  4. View the AI's answers and thinking process
""", unsafe_allow_html=True) # Document upload section st.markdown('
📄 Upload Documents
', unsafe_allow_html=True) st.markdown('
Select files to upload:
', unsafe_allow_html=True) # File uploader uploaded_files = st.file_uploader("Upload documents", type=["pdf", "txt", "png", "jpg", "jpeg"], accept_multiple_files=True, label_visibility="collapsed") # Process button if uploaded_files: if st.button("Process Documents"): with st.spinner("Processing documents..."): # Process progress bar progress_bar = st.progress(0) for i in range(100): time.sleep(0.01) progress_bar.progress(i + 1) # Process the files text, file_info = process_uploaded_files(uploaded_files) st.session_state.vectorstore = create_vectorstore(text) st.session_state.documents_processed = True st.session_state.file_info = file_info # Display success message st.success(f"✅ Processed {len(uploaded_files)} documents successfully!") # Document info section if "file_info" in st.session_state and st.session_state.file_info: st.markdown('
', unsafe_allow_html=True) st.markdown('
📋 Document Information
', unsafe_allow_html=True) # Display file list for i, file in enumerate(st.session_state.file_info): st.markdown(f"""
{file['name']}
{file['type']} file
""", unsafe_allow_html=True) # Remove documents button if st.button("Remove All Documents"): if "vectorstore" in st.session_state: del st.session_state.vectorstore if "file_info" in st.session_state: del st.session_state.file_info if "documents_processed" in st.session_state: del st.session_state.documents_processed st.success("All documents removed!") st.rerun() # Controls section st.markdown('
', unsafe_allow_html=True) st.markdown('
⚙️ Controls
', unsafe_allow_html=True) # Clear chat button if st.button("Clear Chat"): if "messages" in st.session_state: st.session_state.messages = [] if "thinking_history" in st.session_state: st.session_state.thinking_history = [] st.rerun() # Reset all button if st.button("Reset All"): clear_session_state() st.rerun() # Hide/Show examples button if st.button("Hide Examples" if st.session_state.show_examples else "Show Examples"): st.session_state.show_examples = not st.session_state.show_examples st.rerun() # Main content area st.title("Document Q&A Assistant") # Example questions section - only show if flag is True if st.session_state.show_examples: st.markdown("### Example Questions") cols = st.columns(len(EXAMPLE_QUESTIONS)) for i, question in enumerate(EXAMPLE_QUESTIONS): with cols[i]: if st.button(question, key=f"example_{hash(question)}"): st.session_state.messages.append(HumanMessage(content=question)) # Generate response if vectorstore exists if "vectorstore" in st.session_state: with st.spinner("Thinking..."): # Create RAG agent rag_agent = create_rag_agent(st.session_state.vectorstore) # Run the agent result = rag_agent.invoke({ "messages": [HumanMessage(content=question)], "documents": [], "thinking": "" }) # Store thinking process st.session_state.thinking_history.append(result["thinking"]) # Add AI message to chat history st.session_state.messages.append(result["messages"][-1]) else: # Add AI message to chat history st.session_state.messages.append(AIMessage(content="Please upload and process documents first.")) st.rerun() # Chat container st.markdown("### 💬 Chat") chat_container = st.container() with chat_container: # Display chat messages if st.session_state.messages: for i, message in enumerate(st.session_state.messages): if isinstance(message, HumanMessage): st.markdown(f"""
User: {message.content}
""", unsafe_allow_html=True) else: st.markdown(f"""
Assistant: {message.content}
""", unsafe_allow_html=True) # Show thinking process if available if "thinking_history" in st.session_state and i//2 < len(st.session_state.thinking_history): thinking = st.session_state.thinking_history[i//2] # Create a unique key for this thinking process thinking_key = f"thinking_{i//2}" # Store the visibility state in session_state if not already there if thinking_key not in st.session_state: st.session_state[thinking_key] = False # Toggle button for thinking process toggle_text = "Show thinking" if not st.session_state[thinking_key] else "Hide thinking" # Create the toggle button if st.button(toggle_text, key=f"toggle_{thinking_key}"): st.session_state[thinking_key] = not st.session_state[thinking_key] st.rerun() # Show thinking process if toggled on if st.session_state[thinking_key]: with st.expander("Thinking Process", expanded=True): st.write(thinking) else: st.info("Upload documents and start asking questions!") # Chat input st.markdown("### Ask a question about your documents") with st.form(key="chat_form", clear_on_submit=True): user_input = st.text_input("Type your question here...", key="user_question", label_visibility="collapsed") cols = st.columns([6, 1]) with cols[0]: submit_button = st.form_submit_button("Ask", use_container_width=True) if submit_button and user_input: # Add user message to chat history st.session_state.messages.append(HumanMessage(content=user_input)) # Generate response if vectorstore exists if "vectorstore" in st.session_state: with st.spinner("Thinking..."): # Create RAG agent rag_agent = create_rag_agent(st.session_state.vectorstore) # Run the agent result = rag_agent.invoke({ "messages": [HumanMessage(content=user_input)], "documents": [], "thinking": "" }) # Store thinking process st.session_state.thinking_history.append(result["thinking"]) # Add AI message to chat history st.session_state.messages.append(result["messages"][-1]) else: # Add AI message to chat history st.session_state.messages.append(AIMessage(content="Please upload and process documents first.")) # Rerun to update the UI st.rerun() if __name__ == "__main__": main()