Update app.py

app.py

@@ -164,9 +164,9 @@ def merge_markdown_contents(contents):
 
 def upload_to_firebase(user_id, file):
     """
-    Upload document to Firebase and extract content for querying.
+    Upload document to Firebase, extract content, and add it to the knowledge base.
     """
-    content = convert_file_to_md(file)  # Ensure this function is working as expected
+    content = convert_file_to_md(file)  # Ensure this function extracts content correctly
     if not content:
         return None, "Failed to extract content from the file."
 
@@ -176,15 +176,41 @@ def upload_to_firebase(user_id, file):
     # Save document to Firebase
     db.child("users").child(user_id).child("KnowledgeBase").child(doc_id).set(document_data)
 
-    # Update session state
-    if "documents" not in st.session_state:
-        st.session_state["documents"] = {}
-    st.session_state["documents"][doc_id] = document_data
+    # Add content to the knowledge base
+    if "knowledge_base" not in st.session_state:
+        st.session_state["knowledge_base"] = []
+    st.session_state["knowledge_base"].append({"doc_id": doc_id, "content": content})
+
+    # Index the document content for semantic search
+    index_document_content(content, doc_id)
 
-    st.sidebar.success(f"Document '{file.name}' uploaded successfully!")
+    st.sidebar.success(f"Document '{file.name}' uploaded successfully and added to the knowledge base!")
     return content, None
 
 
+def index_document_content(doc_content, doc_id):
+    """
+    Indexes the document content by splitting it into chunks and creating embeddings.
+    """
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    texts = text_splitter.split_text(doc_content)
+
+    # Create embeddings for each chunk
+    embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
+    doc_metadata = [{"doc_id": doc_id, "chunk_id": i} for i in range(len(texts))]
+    vector_store = FAISS.from_texts(texts, embeddings, metadatas=doc_metadata)
+
+    # Save the vector store in session state
+    if "vector_store" not in st.session_state:
+        st.session_state["vector_store"] = {}
+    st.session_state["vector_store"][doc_id] = vector_store
+
+
+
+
+
+
+
 
 def fetch_trustbuilders(user_id):
     """
@@ -1060,21 +1086,23 @@ def google_search(query):
 
 def rag_response(query):
     """
-    Handle RAG-based queries when uploaded document context is not mentioned.
+    Handle queries by searching both static and dynamically uploaded knowledge base.
     """
     try:
-        if "uploaded document" in query.lower():
-            # Handle document-specific queries
-            return handle_document_query(query)
-
-        # Proceed with the existing knowledge base logic
-        retrieved_docs = search_knowledge_base(query)  # Replace with actual KB search logic
-        if not retrieved_docs:
+        # Retrieve relevant chunks from the vector store
+        results = []
+        if "vector_store" in st.session_state:
+            for vector_store in st.session_state["vector_store"].values():
+                results.extend(vector_store.similarity_search(query, k=3))  # Adjust `k` for the number of results
+
+        # Combine results into a context
+        context = "\n".join([result.page_content for result in results])
+        if not context:
             return "No relevant information found in the knowledge base."
 
-        context = "\n".join(doc.page_content for doc in retrieved_docs)
+        # Generate AI response with the retrieved context
         prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
-        llm = ChatOpenAI(model="gpt-4o", temperature=0.3, api_key=openai_api_key)
+        llm = ChatOpenAI(model="gpt-4", temperature=0.3, api_key=openai_api_key)
         response = llm.invoke(prompt)
 
         return response.content
@@ -1083,7 +1111,6 @@ def rag_response(query):
         return "An error occurred during the RAG response generation process."
 
 
-
 # Define tools
 @tool
 def knowledge_base_tool(query: str):
@@ -1799,54 +1826,7 @@ def load_user_memory(user_id):
         st.session_state["documents"] = {}
         st.session_state["vector_store"] = {}
 
-def get_document_content(doc_name=None):
-    """
-    Retrieve content of an uploaded document from Streamlit session state.
-    """
-    documents = st.session_state.get("documents", {})
-    if not documents:
-        return None, "No documents have been uploaded."
-
-    if doc_name:
-        doc_name = doc_name.strip().lower()
-        for doc_id, doc_data in documents.items():
-            if doc_data.get("name", "").strip().lower() == doc_name:
-                return doc_data.get("content"), None
-        return None, f"Document '{doc_name}' not found."
-
-    # Default to the most recently uploaded document
-    last_doc = list(documents.values())[-1]
-    return last_doc.get("content"), None
-
-
-def handle_document_query(query):
-    """
-    Handle user queries related to uploaded documents.
-    """
-    # Extract document name from the query
-    doc_name_match = re.search(r"[\"']?([^\"']+\.(pdf|docx|doc|txt))[\"']?", query, re.IGNORECASE)
-    doc_name = doc_name_match.group(1) if doc_name_match else None
 
-    if not doc_name:
-        return "Please specify a document name in your query."
-
-
-    # Fetch document content
-    doc_content, error = get_document_content(doc_name)
-    if error:
-        return error
-
-
-    # Generate AI response using document context
-    full_prompt = f"Document Content:\n{doc_content}\n\nUser Query: {query}\n\nResponse:"
-    try:
-        llm = ChatOpenAI(model="gpt-4o", temperature=0.5, api_key=openai_api_key)
-        response = llm.invoke(full_prompt)
-        return response.content
-    except Exception as e:
-        logger.error(f"Error generating response using the document: {e}")
-        return f"Error generating response using the document: {e}"
-        
 
 if "missing_trustbucket_content" not in st.session_state:
     st.session_state["missing_trustbucket_content"] = None