# utils/database.py from langchain_community.chat_models import ChatOpenAI from langchain_core.messages import ( HumanMessage, AIMessage, SystemMessage, BaseMessage ) from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.runnables import RunnablePassthrough from langchain.chains import ConversationalRetrievalChain from langchain.chat_models import ChatOpenAI from langchain.agents import AgentExecutor, Tool, create_openai_tools_agent from langchain.agents.format_scratchpad.tools import format_to_tool_messages from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser from utils.document_chunker import DocumentChunker import os import streamlit as st import sqlite3 import traceback import time import io import tempfile from langchain_community.document_loaders import PyPDFLoader from sqlite3 import Error def create_connection(db_file): """Create a database connection to the SQLite database.""" conn = None try: conn = sqlite3.connect(db_file) return conn except Error as e: st.error("Failed to connect to database. Please try again or contact support.") return None # Add this function to your database.py file def get_db_connection(): """Get a thread-safe database connection.""" try: data_dir = Path("data") data_dir.mkdir(exist_ok=True) db_path = data_dir / 'rfp_analysis.db' # Create new connection conn = sqlite3.connect(str(db_path)) # Create tables if they don't exist create_tables(conn) return conn except Exception as e: st.error(f"Database connection error: {str(e)}") return None def create_tables(conn): """Create necessary tables in the database.""" try: sql_create_documents_table = ''' CREATE TABLE IF NOT EXISTS documents ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, content TEXT NOT NULL, upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); ''' sql_create_queries_table = ''' CREATE TABLE IF NOT EXISTS queries ( id INTEGER PRIMARY KEY AUTOINCREMENT, query TEXT NOT NULL, response TEXT NOT NULL, document_id INTEGER, query_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (document_id) REFERENCES documents (id) ); ''' sql_create_annotations_table = ''' CREATE TABLE IF NOT EXISTS annotations ( id INTEGER PRIMARY KEY AUTOINCREMENT, document_id INTEGER NOT NULL, annotation TEXT NOT NULL, page_number INTEGER, annotation_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (document_id) REFERENCES documents (id) ); ''' conn.execute(sql_create_documents_table) conn.execute(sql_create_queries_table) conn.execute(sql_create_annotations_table) except Error as e: st.error(f"Error: {e}") def insert_document(name, content): """Insert a document with thread-safe connection.""" try: conn = get_db_connection() if conn is None: return None cursor = conn.cursor() cursor.execute( "INSERT INTO documents (name, content) VALUES (?, ?)", (name, content) ) conn.commit() doc_id = cursor.lastrowid conn.close() return doc_id except Exception as e: st.error(f"Error inserting document: {str(e)}") if conn: conn.rollback() conn.close() return None def get_documents(conn): """Retrieve all documents from the database. Args: conn: SQLite database connection Returns: tuple: (list of document contents, list of document names) """ try: cursor = conn.cursor() cursor.execute("SELECT content, name FROM documents") results = cursor.fetchall() if not results: return [], [] # Separate contents and names document_contents = [row[0] for row in results] document_names = [row[1] for row in results] return document_contents, document_names except Error as e: st.error(f"Error retrieving documents: {e}") return [], [] def insert_document(conn, name, content): """Insert a new document into the database. Args: conn: SQLite database connection name (str): Name of the document content (str): Content of the document Returns: int: ID of the inserted document, or None if insertion failed """ try: cursor = conn.cursor() sql = '''INSERT INTO documents (name, content) VALUES (?, ?)''' cursor.execute(sql, (name, content)) conn.commit() return cursor.lastrowid except Error as e: st.error(f"Error inserting document: {e}") return None def verify_vector_store(vector_store): """Verify that the vector store has documents loaded. Args: vector_store: FAISS vector store instance Returns: bool: True if vector store is properly initialized with documents """ try: # Try to perform a simple similarity search test_results = vector_store.similarity_search("test", k=1) return len(test_results) > 0 except Exception as e: st.error(f"Vector store verification failed: {e}") return False def handle_document_upload(uploaded_files): """Handle document upload with improved chunking and progress tracking.""" # Initialize containers first - before any processing progress_container = st.empty() status_container = st.empty() details_container = st.empty() progress_bar = progress_container.progress(0) try: # Initialize session state variables if 'qa_system' not in st.session_state: st.session_state.qa_system = None if 'vector_store' not in st.session_state: st.session_state.vector_store = None # Initialize persistence manager persistence = PersistenceManager() # Generate a session ID based on timestamp and files session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}" # Initialize embeddings (10% progress) status_container.info("🔄 Initializing embeddings model...") embeddings = get_embeddings_model() if not embeddings: status_container.error("❌ Failed to initialize embeddings model") return progress_bar.progress(10) # Initialize document chunker chunker = DocumentChunker( chunk_size=1000, chunk_overlap=200, max_tokens_per_chunk=2000 ) # Process documents document_pairs = [] # List to store (content, filename) pairs progress_per_file = 70 / len(uploaded_files) current_progress = 10 for idx, uploaded_file in enumerate(uploaded_files): file_name = uploaded_file.name status_container.info(f"🔄 Processing document {idx + 1}/{len(uploaded_files)}: {file_name}") details_container.text(f"📄 Current file: {file_name}") # Create temporary file for PDF processing with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file: tmp_file.write(uploaded_file.getvalue()) tmp_file.flush() try: # Load PDF content loader = PyPDFLoader(tmp_file.name) pdf_documents = loader.load() content = "\n".join(doc.page_content for doc in pdf_documents) # Store original content in database doc_id = insert_document(st.session_state.db_conn, file_name, content) if not doc_id: status_container.error(f"❌ Failed to store document: {file_name}") continue document_pairs.append((content, file_name)) finally: # Ensure temporary file is cleaned up try: os.unlink(tmp_file.name) except Exception as e: st.warning(f"Could not delete temporary file: {e}") current_progress += progress_per_file progress_bar.progress(int(current_progress)) if not document_pairs: status_container.error("❌ No documents were successfully processed") return # Chunk documents (80% progress) status_container.info("🔄 Chunking documents...") details_container.text("📑 Splitting documents into manageable chunks...") chunks, chunk_metadatas = chunker.process_documents(document_pairs) if not chunks: status_container.error("❌ Failed to chunk documents") return progress_bar.progress(80) # Save chunks for persistence persistence.save_chunks(chunks, chunk_metadatas, session_id) # Initialize vector store (90% progress) status_container.info("🔄 Initializing vector store...") details_container.text("🔍 Creating vector embeddings...") vector_store = initialize_faiss(embeddings, chunks, chunk_metadatas) if not vector_store: status_container.error("❌ Failed to initialize vector store") return # Save vector store and update session state persistence.save_vector_store(vector_store, session_id) st.session_state.vector_store = vector_store st.session_state.current_session_id = session_id progress_bar.progress(90) # Initialize QA system (100% progress) status_container.info("🔄 Setting up QA system...") qa_system = initialize_qa_system(vector_store) if not qa_system: status_container.error("❌ Failed to initialize QA system") return st.session_state.qa_system = qa_system progress_bar.progress(100) # Success message status_container.success("✅ Documents processed successfully!") details_container.markdown(f""" 🎉 **Ready to chat!** - Documents processed: {len(document_pairs)} - Total chunks created: {len(chunks)} - Average chunk size: {sum(len(chunk) for chunk in chunks) / len(chunks):.0f} characters - Vector store initialized and saved - QA system ready - Session ID: {session_id} You can now start asking questions about your documents! """) st.balloons() st.session_state.chat_ready = True except Exception as e: status_container.error(f"❌ Error processing documents: {str(e)}") details_container.error(traceback.format_exc()) st.session_state.vector_store = None st.session_state.qa_system = None st.session_state.chat_ready = False finally: # Clean up progress display after successful processing if st.session_state.get('qa_system') is not None: time.sleep(5) progress_container.empty() def display_vector_store_info(): """Display information about the current vector store state.""" if 'vector_store' not in st.session_state: st.info("ℹī¸ No documents loaded yet.") return try: # Get the vector store from session state vector_store = st.session_state.vector_store # Get basic stats test_query = vector_store.similarity_search("test", k=1) doc_count = len(test_query) # Create an expander for detailed info with st.expander("📊 Knowledge Base Status"): col1, col2 = st.columns(2) with col1: st.metric( label="Documents Loaded", value=doc_count ) with col2: st.metric( label="System Status", value="Ready" if verify_vector_store(vector_store) else "Not Ready" ) # Display sample queries if verify_vector_store(vector_store): st.markdown("### 🔍 Sample Document Snippets") sample_docs = vector_store.similarity_search("", k=3) for i, doc in enumerate(sample_docs, 1): with st.container(): st.markdown(f"**Snippet {i}:**") st.text(doc.page_content[:200] + "...") except Exception as e: st.error(f"Error displaying vector store info: {e}") st.error(traceback.format_exc()) def initialize_qa_system(vector_store): """Initialize QA system with proper chat handling.""" try: llm = ChatOpenAI( temperature=0.5, model_name="gpt-4", api_key=os.environ.get("OPENAI_API_KEY") ) # Create retriever function retriever = vector_store.as_retriever(search_kwargs={"k": 2}) # Create a template that enforces clean formatting prompt = ChatPromptTemplate.from_messages([ ("system", """You are an expert consultant specializing in analyzing Request for Proposal (RFP) documents. Your goal is to assist users by providing clear, concise, and professional insights based on the content provided. Please adhere to the following guidelines when crafting your responses: Begin with a summary that highlights the key findings or answers the main query. Structured Format: Use clear and descriptive section headers to organize the information logically. Bullet Points: Utilize bullet points for lists or complex information to enhance readability. Source Attribution: Cite specific sections or page numbers from the RFP document when referencing information. Professional Formatting: Maintain a clean and professional layout using Markdown formatting where appropriate (e.g., headings, bold, italics). Focused Content: Keep your responses concise and directly related to the user's query, avoiding unnecessary information. Scope Awareness: If a query falls outside the provided information or context, politely acknowledge this and suggest consulting the relevant sections or additional sources. Confidentiality: Respect the confidentiality of the information provided and avoid sharing any sensitive data beyond the scope of the query. Tone and Language: Use formal and professional language, ensuring clarity and precision in your responses. Accuracy: Double-check all information for accuracy and completeness before providing it to the user. """), MessagesPlaceholder(variable_name="chat_history"), ("human", "{input}\n\nContext: {context}") ]) def get_chat_history(inputs): chat_history = inputs.get("chat_history", []) if not isinstance(chat_history, list): return [] return [msg for msg in chat_history if isinstance(msg, BaseMessage)] def get_context(inputs): docs = retriever.get_relevant_documents(inputs["input"]) context_parts = [] for doc in docs: source = doc.metadata.get('source', 'Unknown source') context_parts.append(f"\nFrom {source}:\n{doc.page_content}") return "\n".join(context_parts) chain = ( { "context": get_context, "chat_history": get_chat_history, "input": lambda x: x["input"] } | prompt | llm ) return chain except Exception as e: st.error(f"Error initializing QA system: {e}") return None # FAISS vector store initialization def initialize_faiss(embeddings, documents, document_names): """Initialize FAISS vector store.""" try: from langchain.vectorstores import FAISS vector_store = FAISS.from_texts( documents, embeddings, metadatas=[{"source": name} for name in document_names], ) return vector_store except Exception as e: st.error(f"Error initializing FAISS: {e}") return None # Embeddings model retrieval @st.cache_resource def get_embeddings_model(): """Get the embeddings model.""" try: from langchain.embeddings import HuggingFaceEmbeddings model_name = "sentence-transformers/all-MiniLM-L6-v2" embeddings = HuggingFaceEmbeddings(model_name=model_name) return embeddings except Exception as e: st.error(f"Error loading embeddings model: {e}") return None