Update app.py

changed pinecone to chromadb to avoid limits

app.py

@@ -12,10 +12,7 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_core.runnables import RunnablePassthrough
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
-from langchain_community.retrievers import PineconeHybridSearchRetriever
-from langchain_pinecone import PineconeVectorStore
-from pinecone import Pinecone, ServerlessSpec
-from pinecone import PineconeApiException, NotFoundException
+from langchain_chroma import Chroma
 import shutil
 import uuid
 
@@ -23,7 +20,7 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # Set page configuration
-st.set_page_config(page_title="Document Analyzer", layout="wide", )
+st.set_page_config(page_title="Document Analyzer", layout="wide")
 
 st.title("📚 Document Analyzer")
 
@@ -41,46 +38,65 @@ if 'initialized' not in st.session_state:
     st.session_state.initialized = False
 if 'processing' not in st.session_state:
     st.session_state.processing = False
-if 'last_processed_files' not in st.session_state:
-    st.session_state.last_processed_files = set()
-if 'chat_history' not in st.session_state:
-    st.session_state.chat_history = []
 if 'chat_enabled' not in st.session_state:
     st.session_state.chat_enabled = False
 if 'session_id' not in st.session_state:
     # Generate a unique session ID using UUID
     st.session_state.session_id = str(uuid.uuid4())[:8]
 
-def get_session_index_name():
-    """Get unique index name for current session"""
-    base_name = "docdb"  # Using a short base name to leave room for the unique identifier
-    unique_id = st.session_state.session_id
-    # Combine base name with unique ID, ensuring total length is under 45 chars
-    return f"{base_name}-{unique_id}"  # This will be like "docdb-12345678"
+def get_chroma_directory():
+    """Get unique directory name for current session's ChromaDB"""
+    base_dir = "vectorstores"
+    if not os.path.exists(base_dir):
+        os.makedirs(base_dir)
+    return os.path.join(base_dir, f"chroma_db_{st.session_state.session_id}")
 
-def cleanup_pinecone_index():
-    """Clean up existing Pinecone index for the current session"""
+def cleanup_chroma_db():
+    """Clean up existing ChromaDB for the current session"""
     try:
-        pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
-        index_name = get_session_index_name()
-        if index_name in pc.list_indexes().names():
-            pc.delete_index(index_name)
+        chroma_dir = get_chroma_directory()
+        if os.path.exists(chroma_dir):
+            shutil.rmtree(chroma_dir)
     except Exception as e:
-        print(f"Error cleaning up index: {str(e)}")  # Log error internally
+        print(f"Error cleaning up ChromaDB: {str(e)}")  # Log error internally
+
+def cleanup_old_vectorstores():
+    """Clean up vector stores that are older than 24 hours"""
+    try:
+        base_dir = "vectorstores"
+        if not os.path.exists(base_dir):
+            return
+            
+        current_time = time.time()
+        one_day_in_seconds = 24 * 60 * 60
+        
+        # Get all directories in vectorstores
+        for dir_name in os.listdir(base_dir):
+            dir_path = os.path.join(base_dir, dir_name)
+            if os.path.isdir(dir_path):
+                # Get directory's last modification time
+                last_modified = os.path.getmtime(dir_path)
+                if current_time - last_modified > one_day_in_seconds:
+                    shutil.rmtree(dir_path)
+    except Exception as e:
+        print(f"Error cleaning up old vector stores: {str(e)}")  # Log error internally
 
 if not st.session_state.initialized:
+    # Clean up old vector stores first
+    cleanup_old_vectorstores()
+    
     # Clear everything only on first run or page refresh
     if os.path.exists("data"):
         shutil.rmtree("data")
     os.makedirs("data")
+    
+    # Clear vectorstores directory for current session
+    if os.path.exists("vectorstores"):
+        os.makedirs("vectorstores", exist_ok=True)
+    
     st.session_state.uploaded_files = {}
     st.session_state.previous_files = set()
-    st.session_state.vectorstore = None
-    st.session_state.retriever = None
     st.session_state.initialized = True
-    
-    # Clean up any existing index
-    cleanup_pinecone_index()
 
 def save_uploaded_file(uploaded_file):
     """Save uploaded file to the data directory"""
@@ -105,15 +121,15 @@ def save_uploaded_file(uploaded_file):
         return None
 
 def process_documents(uploaded_files_dict):
-    """Process documents and store in Pinecone"""
+    """Process documents and store in ChromaDB"""
     warning_placeholder = st.empty()
     warning_placeholder.warning("⚠️ Document processing in progress. Please wait before adding or removing files.")
     success_placeholder = st.empty()
     
     try:
         with st.spinner('Processing documents...'):
-            # Clean up existing index before processing
-            cleanup_pinecone_index()
+            # Clean up existing ChromaDB before processing
+            cleanup_chroma_db()
             
             docs = []
             # Process each file
@@ -152,27 +168,14 @@ def process_documents(uploaded_files_dict):
             # Initialize embeddings
             embed_func = OpenAIEmbeddings(model='text-embedding-3-small', dimensions=512)
             
-            # Initialize Pinecone
-            pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
-            index_name = get_session_index_name()
-
             try:
-                pc.create_index(
-                    name=index_name,
-                    dimension=512,
-                    metric='cosine',
-                    spec=ServerlessSpec(cloud='aws', region='us-east-1')
-                )
-
-                # Wait for index to be ready
-                while not pc.describe_index(index_name).status['ready']:
-                    time.sleep(1)
-
-                pc_index = pc.Index(index_name)
-                
                 # Create vectorstore and add documents
-                vectorstore = PineconeVectorStore(index=pc_index, embedding=embed_func)
-                vectorstore.add_documents(documents=chunks)
+                vectorstore = Chroma.from_documents(
+                    collection_name="collection",
+                    documents=chunks,
+                    embedding=embed_func,
+                    persist_directory=get_chroma_directory()
+                )
 
                 st.session_state.chat_enabled = True
                 success_placeholder.success('Documents processed successfully!')
@@ -180,8 +183,8 @@ def process_documents(uploaded_files_dict):
                 success_placeholder.empty()  # Clear the success message
                 return True
                 
-            except PineconeApiException as e:
-                print(f"Pinecone API error: {str(e)}")  # Log error internally
+            except Exception as e:
+                print(f"ChromaDB error: {str(e)}")  # Log error internally
                 st.warning("Unable to process documents at the moment. Please try again.")
                 st.session_state.chat_enabled = False
                 return False
@@ -195,38 +198,9 @@ def process_documents(uploaded_files_dict):
         warning_placeholder.empty()
 
 def doc2str(docs):
-    return "\n\n".join(doc for doc in docs)
-
-def format_reranked_docs(pc, retriever, question):
-    """Rerank documents using Pinecone's reranking model"""
-    # Get relevant docs and ensure they're not empty
-    relevant_docs = [doc.page_content for doc in retriever.invoke(question) if doc.page_content.strip()]
-
-    if not relevant_docs:
-        return "I don't have enough context to answer this question."
-
-    try:
-        # Format documents for reranking
-        formatted_docs = [{"text": doc} for doc in relevant_docs]
-        
-        reranked_docs = pc.inference.rerank(
-            model="pinecone-rerank-v0",
-            query=question,
-            documents=formatted_docs,
-            top_n=3,
-            return_documents=True
-        )
-
-        # Extract text from reranked documents
-        final_docs = [d.document["text"] for d in reranked_docs.data]
-        context = "\n\n".join(final_docs)
-        return context
-    except Exception as e:
-        print(f"Error during reranking: {str(e)}")  # Log error internally
-        # Fallback to using retrieved docs without reranking
-        return "\n\n".join(relevant_docs[:3])
+    return "\n\n".join(doc.page_content for doc in docs)
 
-def run_chatbot(retriever, pc, llm):
+def run_chatbot(retriever, llm):
     """Run the chatbot with the given components"""
     # Initialize chat prompt
     prompt = ChatPromptTemplate.from_template("""
@@ -245,9 +219,9 @@ def run_chatbot(retriever, pc, llm):
 
     {question}""")
 
-    # Create the QA chain with reranking
+    # Create the QA chain
     qa_chain = (
-        RunnablePassthrough.assign(context=lambda input: format_reranked_docs(pc, retriever, input["question"]))
+        RunnablePassthrough.assign(context=lambda input: doc2str(retriever.invoke(input["question"])))
         | prompt
         | llm
         | StrOutputParser()
@@ -305,8 +279,14 @@ def process_and_chat():
     # Check for removed files
     files_to_remove = set(st.session_state.uploaded_files.keys()) - current_uploaded_filenames
     if files_to_remove:
-        # Clean up index when files are removed
-        cleanup_pinecone_index()
+        # Set processing state immediately
+        st.session_state.processing = True
+        st.session_state.chat_enabled = False
+        if "messages" in st.session_state:
+            del st.session_state.messages
+            
+        # Clean up ChromaDB when files are removed
+        cleanup_chroma_db()
         for file_name in files_to_remove:
             # Remove file from session state
             if file_name in st.session_state.uploaded_files:
@@ -323,6 +303,12 @@ def process_and_chat():
         for file in uploaded_files:
             # Only process files that haven't been uploaded before
             if file.name not in st.session_state.uploaded_files:
+                # Set processing state immediately when new file is detected
+                st.session_state.processing = True
+                st.session_state.chat_enabled = False
+                if "messages" in st.session_state:
+                    del st.session_state.messages
+                    
                 file_path = save_uploaded_file(file)
                 if file_path:  # Only add to session state if file was saved successfully
                     st.session_state.uploaded_files[file.name] = {
@@ -336,45 +322,66 @@ def process_and_chat():
 
     # If files have changed (added or removed), reset chat and process documents
     if current_files != st.session_state.previous_files or files_to_remove:
-        # Reset chat state
-        st.session_state.chat_enabled = False
-        if "messages" in st.session_state:
-            del st.session_state.messages
-        
         st.session_state.previous_files = current_files
         
         if current_files:
-            st.session_state.processing = True
             # Process documents and enable chat if successful
             if process_documents(st.session_state.uploaded_files):
                 st.session_state.chat_enabled = True
             st.session_state.processing = False
         else:
             st.warning('Please upload a file to continue')
+            st.session_state.processing = False
     
     # If files exist and chat is enabled, show chat interface
     if current_files and st.session_state.chat_enabled:
         try:
             # Initialize components for chat
             llm = ChatGroq(temperature=0, model_name="llama-3.3-70b-versatile", groq_api_key=os.getenv("GROQ_API_KEY"), max_tokens=8000)
-            pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
-            index_name = get_session_index_name()
-            pc_index = pc.Index(index_name)
             
             # Create vectorstore
             embed_func = OpenAIEmbeddings(model='text-embedding-3-small', dimensions=512)
-            vectorstore = PineconeVectorStore(index=pc_index, embedding=embed_func)
+            vectorstore = Chroma(
+                collection_name="collection",
+                embedding_function=embed_func,
+                persist_directory=get_chroma_directory()
+            )
             
             # Create retrievers
             vectorstore_retriever = vectorstore.as_retriever(
-                search_type="similarity_score_threshold",
-                search_kwargs={"k": 5, "score_threshold": 0.6},
+                search_kwargs={"k": 3}
+            )
+            
+            # Create keyword retriever
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1500,
+                chunk_overlap=400,
+                length_function=len
+            )
+            docs = []
+            for file_info in st.session_state.uploaded_files.values():
+                if file_info["path"].endswith(".pdf"):
+                    docs.extend(PyMuPDFLoader(file_info["path"]).load())
+                elif file_info["path"].endswith(".txt"):
+                    docs.extend(TextLoader(file_info["path"]).load())
+                elif file_info["path"].endswith(".docx"):
+                    docs.extend(Docx2txtLoader(file_info["path"]).load())
+            
+            chunks = text_splitter.split_documents(docs)
+            keyword_retriever = BM25Retriever.from_documents(chunks)
+            keyword_retriever.k = 3
+            
+            # Combine retrievers
+            ensemble_retriever = EnsembleRetriever(
+                retrievers=[vectorstore_retriever, keyword_retriever],
+                weights=[0.5, 0.5]
             )
             
             # Run chatbot with fresh components
-            run_chatbot(vectorstore_retriever, pc, llm)
-        except NotFoundException:
-            st.error("Vector database not found. Please try uploading your documents again.")
+            run_chatbot(ensemble_retriever, llm)
+        except Exception as e:
+            print(f"Chat interface error: {str(e)}")  # Log error internally
+            st.warning("Please try uploading your documents again.")
             st.session_state.chat_enabled = False
             # Clear the previous files to force reprocessing
             st.session_state.previous_files = set()
@@ -382,5 +389,4 @@ def process_and_chat():
                 del st.session_state.messages
 
 # Call the main function
-process_and_chat()
-
+process_and_chat()