Update app.py

app.py

@@ -1,19 +1,59 @@
 import gradio as gr
 import utils
+from langchain_mistralai import ChatMistralAI
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_community.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.runnables import RunnablePassthrough
+import torch
 
-# Placeholder functions and variables for RAG components
+import os
+os.environ['MISTRAL_API_KEY'] = 'XuyOObDE7trMbpAeI7OXYr3dnmoWy3L0'
 
-# Function to ingest a new file into the system
 class VectorData():
     def __init__(self):
-        self.retriever = None
+        embedding_model_name = 'nomic-ai/nomic-embed-text-v1.5'
+
+        model_kwargs = {'device':'cuda' if torch.cuda.is_available() else 'cpu',"trust_remote_code": True}
+
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=embedding_model_name,
+            model_kwargs=model_kwargs
+        )
+
+        self.vectorstore = Chroma(persist_directory="chroma_db", embedding_function=self.embeddings)
+        self.retriever = self.vectorstore.as_retriever()
         self.ingested_files = []
-        self.vectorstore = None
+        self.prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    """Answer the question based on the given context. Dont give any ans if context is not valid to question. Always give the source of context: 
+                    {context}
+                    """,
+                ),
+                ("human", "{question}"),
+            ]
+        )
+        self.llm = ChatMistralAI(model="mistral-large-latest")
+        self.rag_chain = (
+                {"context": self.retriever, "question": RunnablePassthrough()}
+                | self.prompt
+                | self.llm
+                | StrOutputParser()
+            )
 
     def add_file(self,file):
         if file is not None:
             self.ingested_files.append(file.name.split('/')[-1])
             self.retriever, self.vectorstore = utils.add_doc(file,self.vectorstore)
+            self.rag_chain = (
+                {"context": self.retriever, "question": RunnablePassthrough()}
+                | self.prompt
+                | self.llm
+                | StrOutputParser()
+            )
         return [[name] for name in self.ingested_files]
 
     def delete_file_by_name(self,file_name):
@@ -26,14 +66,15 @@ class VectorData():
         self.ingested_files.clear()
         self.retriever, self.vectorstore = utils.delete_all_doc(self.vectorstore)
         return []
+    
+data_obj = VectorData()
 
 # Function to handle question answering
 def answer_question(question):
     if question.strip():
-        return f"Generated answer for the question: '{question}'"
+        return f'{data_obj.rag_chain.invoke(question)}'
     return "Please enter a question."
 
-data_obj = VectorData()
 
 # Define the Gradio interface
 with gr.Blocks() as rag_interface: