Update utils/database.py

utils/database.py

@@ -401,13 +401,6 @@ def add_query(conn: sqlite3.Connection, query: str, response: str, document_id:
         st.error(f"Error adding query: {e}")
         return False
 
-# Add to utils/database.py
-
-import sqlite3
-from typing import List, Dict, Optional
-from datetime import datetime
-from langchain_core.messages import HumanMessage, AIMessage
-import streamlit as st
 
 def create_chat_tables(conn: sqlite3.Connection) -> None:
     """Create necessary tables for chat management."""
@@ -1656,15 +1649,25 @@ def get_document_content(conn: sqlite3.Connection, document_id: int) -> Optional
         st.error(f"Error retrieving document content: {e}")
         return None
 
-def initialize_qa_system(vector_store):
-    """
-    Initialize QA system with optimized retrieval.
+def get_context_with_sources(retriever, query):
+    """Get context with source documents."""
+    docs = retriever.get_relevant_documents(query)
+    formatted_docs = []
+    
+    for doc in docs:
+        source = doc.metadata.get('source', 'Unknown source')
+        formatted_docs.append(f"\nFrom {source}:\n{doc.page_content}")
+    
+    return "\n".join(formatted_docs)
 
-    Args:
-        vector_store (FAISS): FAISS vector store instance.
-    Returns:
-        dict: QA system chain or None if initialization fails.
-    """
+def format_chat_history(chat_history):
+    """Format chat history for the prompt."""
+    if not chat_history or not isinstance(chat_history, list):
+        return []
+    return [msg for msg in chat_history if isinstance(msg, (HumanMessage, AIMessage))]
+
+def initialize_qa_system(vector_store):
+    """Initialize QA system with optimized retrieval."""
     try:
         llm = ChatOpenAI(
             temperature=0.5,
@@ -1673,60 +1676,35 @@ def initialize_qa_system(vector_store):
             api_key=os.environ.get("OPENAI_API_KEY")
         )
 
-        # Optimize retriever settings and add source tracking
+        # Optimize retriever settings
         retriever = vector_store.as_retriever(
             search_kwargs={
-                "k": 3,  # Retrieve fewer, more relevant chunks
-                "fetch_k": 5,  # Consider more candidates before selecting top k
-                "include_metadata": True  # Enable source tracking
+                "k": 3,
+                "fetch_k": 5,
+                "include_metadata": True
             }
         )
 
-        # Create a template that enforces clean formatting
+        # Create system prompt template
         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:
+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:
 
 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, such as headings, bullet points, bold, italics, and tables where appropriate.
-
-Use Markdown Syntax: Ensure the response is fully formatted using Markdown for optimal readability in the chat.
-
-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.
-"""),
+Use clear section headers to organize information logically.
+Utilize bullet points for lists or complex information.
+Cite specific sections or page numbers from the RFP document when referencing information.
+Maintain professional formatting using Markdown.
+Keep responses focused and directly related to the query.
+Acknowledge when information falls outside the provided context.
+Use formal and professional language.
+Ensure accuracy and completeness in responses.
+            """),
             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)
 
+        # Create the chain
         chain = (
             {
                 "context": lambda x: get_context_with_sources(retriever, x["input"]),