Update src/rag_methods.py

src/rag_methods.py

@@ -2,6 +2,21 @@ import os
 import dotenv
 from time import time
 import streamlit as st
+import logging
+
+# Configure environment for Hugging Face Spaces
+os.environ["HF_HOME"] = "/tmp/.cache/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache/huggingface"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/.cache/huggingface"
+
+# Create necessary directories
+os.makedirs("/tmp/.cache/huggingface", exist_ok=True)
+os.makedirs("/tmp/chroma_persistent_db", exist_ok=True)
+os.makedirs("/tmp/source_files", exist_ok=True)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 from langchain_community.document_loaders.text import TextLoader
 from langchain_community.document_loaders import (
@@ -21,7 +36,17 @@ dotenv.load_dotenv()
 os.environ["USER_AGENT"] = "myagent"
 DB_DOCS_LIMIT = 10
 
-# Stream non-RAG LLM response
+def clean_temp_files():
+    """Clean up temporary files to prevent storage issues"""
+    try:
+        for folder in ["/tmp/source_files"]:
+            for filename in os.listdir(folder):
+                file_path = os.path.join(folder, filename)
+                if os.path.isfile(file_path):
+                    os.unlink(file_path)
+    except Exception as e:
+        logger.warning(f"Error cleaning temp files: {e}")
+
 def stream_llm_response(llm_stream, messages):
     response_message = ""
     for chunk in llm_stream.stream(messages):
@@ -29,18 +54,17 @@ def stream_llm_response(llm_stream, messages):
         yield chunk
     st.session_state.messages.append({"role": "assistant", "content": response_message})
 
-# --- Document Loading and Indexing ---
 def load_doc_to_db():
     if "rag_docs" in st.session_state and st.session_state.rag_docs:
         docs = []
         for doc_file in st.session_state.rag_docs:
             if doc_file.name not in st.session_state.rag_sources:
                 if len(st.session_state.rag_sources) < DB_DOCS_LIMIT:
-                    os.makedirs("source_files", exist_ok=True)
-                    file_path = f"./source_files/{doc_file.name}"
-                    with open(file_path, "wb") as file:
-                        file.write(doc_file.read())
                     try:
+                        file_path = f"/tmp/source_files/{doc_file.name}"
+                        with open(file_path, "wb") as file:
+                            file.write(doc_file.getbuffer())
+                        
                         if doc_file.type == "application/pdf":
                             loader = PyPDFLoader(file_path)
                         elif doc_file.name.endswith(".docx"):
@@ -50,17 +74,22 @@ def load_doc_to_db():
                         else:
                             st.warning(f"Unsupported document type: {doc_file.type}")
                             continue
+                            
                         docs.extend(loader.load())
                         st.session_state.rag_sources.append(doc_file.name)
+                        logger.info(f"Successfully loaded document: {doc_file.name}")
                     except Exception as e:
-                        st.toast(f"Error loading document {doc_file.name}: {e}", icon="⚠️")
+                        st.toast(f"Error loading document {doc_file.name}: {str(e)}", icon="⚠️")
+                        logger.error(f"Error loading document: {e}")
                     finally:
-                        os.remove(file_path)
+                        if os.path.exists(file_path):
+                            os.remove(file_path)
                 else:
                     st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
         if docs:
             _split_and_load_docs(docs)
-            st.toast(f"Documents loaded successfully.", icon="✅")
+            st.toast("Documents loaded successfully.", icon="✅")
+            clean_temp_files()
 
 def load_url_to_db():
     if "rag_url" in st.session_state and st.session_state.rag_url:
@@ -72,8 +101,10 @@ def load_url_to_db():
                     loader = WebBaseLoader(url)
                     docs.extend(loader.load())
                     st.session_state.rag_sources.append(url)
+                    logger.info(f"Successfully loaded URL: {url}")
                 except Exception as e:
-                    st.error(f"Error loading from URL {url}: {e}")
+                    st.error(f"Error loading from URL {url}: {str(e)}")
+                    logger.error(f"Error loading URL: {e}")
                 if docs:
                     _split_and_load_docs(docs)
                     st.toast(f"Loaded content from URL: {url}", icon="✅")
@@ -81,18 +112,16 @@ def load_url_to_db():
                 st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
 
 def initialize_vector_db(docs):
-    # Initialize HuggingFace embeddings
     embedding = HuggingFaceEmbeddings(
         model_name="BAAI/bge-large-en-v1.5",
         model_kwargs={'device': 'cpu'},
-        encode_kwargs={'normalize_embeddings': False}
+        encode_kwargs={'normalize_embeddings': False},
+        cache_folder="/tmp/.cache"
     )
 
-    # Shared persistent directory for long-term storage
-    persist_dir = "./chroma_persistent_db"
+    persist_dir = "/tmp/chroma_persistent_db"
     collection_name = "persistent_collection"
 
-    # Create the persistent Chroma vector store
     vector_db = Chroma.from_documents(
         documents=docs,
         embedding=embedding,
@@ -100,12 +129,10 @@ def initialize_vector_db(docs):
         collection_name=collection_name
     )
 
-    # Persist to disk
     vector_db.persist()
-
+    logger.info("Vector database initialized and persisted")
     return vector_db
 
-
 def _split_and_load_docs(docs):
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=1000,
@@ -118,9 +145,8 @@ def _split_and_load_docs(docs):
         st.session_state.vector_db = initialize_vector_db(chunks)
     else:
         st.session_state.vector_db.add_documents(chunks)
-        st.session_state.vector_db.persist()  # Save changes
-
-# --- RAG Chain ---
+        st.session_state.vector_db.persist()
+        logger.info("Added new documents to existing vector database")
 
 def _get_context_retriever_chain(vector_db, llm):
     retriever = vector_db.as_retriever()
@@ -132,43 +158,4 @@ def _get_context_retriever_chain(vector_db, llm):
     return create_history_aware_retriever(llm, retriever, prompt)
 
 def get_conversational_rag_chain(llm):
-    retriever_chain = _get_context_retriever_chain(st.session_state.vector_db, llm)
-    prompt = ChatPromptTemplate.from_messages([
-        ("system",
-         """You are a helpful assistant answering the user's queries using the provided context if available.\n
-         {context}"""),
-        MessagesPlaceholder(variable_name="messages"),
-        ("user", "{input}")
-    ])
-    stuff_documents_chain = create_stuff_documents_chain(llm, prompt)
-    return create_retrieval_chain(retriever_chain, stuff_documents_chain)
-
-# Stream RAG LLM response
-def stream_llm_rag_response(llm_stream, messages):
-    rag_chain = get_conversational_rag_chain(llm_stream)
-
-    # Extract latest user input and prior messages
-    input_text = messages[-1].content
-    history = messages[:-1]
-
-    # --- DEBUG: Show context retrieved ---
-    if st.session_state.get("debug_mode"):
-        retriever = st.session_state.vector_db.as_retriever()
-        retrieved_docs = retriever.get_relevant_documents(input_text)
-        st.markdown("### 🔍 Retrieved Context (Debug Mode)")
-        for i, doc in enumerate(retrieved_docs):
-            st.markdown(f"**Chunk {i+1}:**\n```\n{doc.page_content.strip()}\n```")
-
-    response_message = "*(RAG Response)*\n"
-    response = rag_chain.stream({
-        "messages": history,
-        "input": input_text
-    })
-
-    for chunk in response:
-        if 'answer' in chunk:
-            response_message += chunk['answer']
-            yield chunk['answer']
-
-    st.session_state.messages.append({"role": "assistant", "content": response_message})
-
+    retriever_chain = _get_context_retriever_chain