Update app.py

app.py

@@ -2,26 +2,24 @@ from fastapi import FastAPI, UploadFile,File,HTTPException
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from dotenv import load_dotenv
-from langchain_community.document_loaders import PyMuPDFLoader
+from langchain_community.document_loaders import PyMuPDFLoader, UnstructuredPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_groq import ChatGroq
-from langchain_pinecone import PineconeVectorStore
 from langchain_core.runnables import RunnablePassthrough
 from pathlib import Path
 import uvicorn
 import shutil
 import os
 import hashlib
-from pinecone import Pinecone
 import fitz
 import pytesseract
 from PIL import Image
 from langchain.schema import Document
+from langchain_community.vectorstores import Chroma
 import io
-import time
 
 app = FastAPI()
 
@@ -36,21 +34,17 @@ app.add_middleware(
 UPLOAD_DIR = "uploads"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
-
-index_name = "pinecone-chatbot"
+persist_directory = "./chroma_db"
 
 load_dotenv()
 os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
-os.environ["PINECONE_API_KEY"] = os.getenv("PINECONE_API_KEY")
 os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
 
-llm = ChatGroq(model_name = "qwen-2.5-32b")
+llm = ChatGroq(model_name = "Llama3-8b-8192")
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 prompt = '''You are given a context below. Use it to answer the question that follows.  
-Provide a concise and factual response based on the context as below". If user mentions keywords such as "file","pdf", "document", please refer them as context. 
-If you cannot find the answer, please reply *the answer cannot be found in the given context*
+Provide a concise and factual response. If the answer is not in the context, simply state "I don't know based on context provided."  
 
 <context>  
 {context}  
@@ -61,9 +55,6 @@ Answer:'''
 
 parser = StrOutputParser()
 
-pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
-index = pc.Index(name=index_name)
-
 def generate_file_id(file_path):   
     hasher = hashlib.md5()
     with open(file_path, "rb") as f:
@@ -71,16 +62,15 @@ def generate_file_id(file_path):
     return hasher.hexdigest()
 
 def delete_existing_embedding(file_id):
-    index_stats = index.describe_index_stats()
-    if index_stats["total_vector_count"] > 0:            
-        index.delete(delete_all=True)
+    if os.path.exists(persist_directory):
+        shutil.rmtree(persist_directory)
 
 def tempUploadFile(filePath,file):
     with open(filePath,'wb') as buffer:
         shutil.copyfileobj(file.file, buffer)
 
 def loadAndSplitDocuments(filePath):
-    loader = PyMuPDFLoader(filePath)
+    loader = UnstructuredPDFLoader(filePath)
     docs = loader.load()
     
     splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=500)
@@ -118,30 +108,16 @@ def loadAndSplitPdfFile(filePath):
     final_chunks = splitter.split_documents(documents)
     return final_chunks
 
-def prepare_retriever(filePath = "", load_from_pinecone = False):
-    if load_from_pinecone:
-        vector_store = PineconeVectorStore.from_existing_index(index_name, embeddings)
+def prepare_retriever(filePath = "", load_from_chromadb = False):
+    if load_from_chromadb:
+        vector_store = Chroma(persist_directory=persist_directory, embedding_function= embeddings)
         return vector_store.as_retriever(search_kwargs={"k": 5})    
     elif filePath:   
-        doc_chunks = loadAndSplitPdfFile(filePath)        
-        vector_data = []      
+        doc_chunks = loadAndSplitPdfFile(filePath)
+        print(f"Loaded {len(doc_chunks)} documents from {filePath}") 
     
-        for i, doc in enumerate(doc_chunks):
-            embedding = embeddings.embed_query(doc.page_content)              
-            if embedding:  
-                metadata = {
-                "text": doc.page_content,
-                "source": str(doc.metadata.get("source", "unknown")),
-                "page": int(doc.metadata.get("page", i)),  # Add page info if available
-        }
-            vector_data.append((str(i), embedding, metadata))          
-        print(f"Upserting {len(vector_data)} records into Pinecone...")
-
-        index.describe_index_stats()
-        time.sleep(2)
-        
-        index.upsert(vectors=vector_data)  
-        print("Upsert complete")
+        vector_store = Chroma.from_documents(documents= doc_chunks, persist_directory=persist_directory, embedding_function= embeddings)
+        vector_store.persist()
 
 def get_retriever_chain(retriever):    
     chat_prompt = ChatPromptTemplate.from_template(prompt)
@@ -156,6 +132,7 @@ def UploadFileInStore(file: UploadFile = File(...)):
     filePath = Path(UPLOAD_DIR) / file.filename
     tempUploadFile(filePath,file)
     file_id = generate_file_id(filePath)
+    
     delete_existing_embedding(file_id)
     prepare_retriever(filePath)
     
@@ -166,7 +143,7 @@ def UploadFileInStore(file: UploadFile = File(...)):
 
 @app.get("/QnAFromPdf")
 async def QnAFromPdf(query: str):   
-    retriever = prepare_retriever(load_from_pinecone=True) 
+    retriever = prepare_retriever(load_from_chromadb=True) 
     chain = get_retriever_chain(retriever)
     response = chain.invoke(query)    
     return response