Update app.py

app.py

@@ -7,6 +7,9 @@ from google.genai.types import GenerateContentConfig
 from datasets import load_dataset
 from huggingface_hub import login
 from typing import List, Dict, Any
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
 
 app = FastAPI()
 
@@ -16,10 +19,23 @@ google_api_key = os.environ.get("GOOGLE_API_KEY")
 
 login(token=hf_token)
 
-# grab dataset
-dataset = load_dataset("Lhumpal/youtube-hunting-beast-transcripts", data_files={"concise": "concise/*", "raw": "raw/*"})
-concise_text = dataset["concise"]["text"]
-concise_text_string = "".join(concise_text)
+def chunk_text(text, chunk_size=50, chunk_overlap=10):
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=[" ", "\n"]
+    )
+    chunks = splitter.split_text(text)
+    return chunks
+
+# Function to build FAISS index
+def build_faiss_vectorstore(chunks):
+    vectorstore = FAISS.from_texts(chunks, embedding_model)
+    return vectorstore
+
+# Function to retrieve similar text
+def retrieve(query, vectorstore, top_k=5):
+    docs = vectorstore.similarity_search(query, k=top_k)
+    return [doc.page_content for doc in docs]
+    
 
 class ChatRequest(BaseModel):
     message: str
@@ -33,28 +49,57 @@ class ChatRequest(BaseModel):
     chat_history: List[Dict[str, Any]] = []
     model_choice: str = "google"
 
+
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# grab dataset
+dataset = load_dataset("Lhumpal/youtube-hunting-beast-transcripts", data_files={"concise": "concise/*", "raw": "raw/*"})
+concise_text = dataset["concise"]["text"]
+concise_text_string = "".join(concise_text)
+
+# Chunk and index the documents
+chunks = chunk_text(concise_text_string, chunk_size=30))
+# Build the vectorsore
+vectorstore = build_faiss_vectorstore(chunks)
+
 @app.post("/chat")
 async def chat(request: ChatRequest):
     try:
         if request.model_choice == "google":
             client = genai.Client(api_key=google_api_key)
 
-            # summarize chat history
-            summary_thresh = 4
-            if len(request.chat_history) > summary_thresh:
-                summarize_prompt = f"""Please summarize the following chat history concisely, focusing on the key points and main topics discussed. Avoid 
-                unnecessary details and provide a clear, straightforward summary. {request.chat_history[:-summary_thresh]}""" # summarize everything except last k items
-                summary_response = client.models.generate_content(
-                    model="gemini-2.0-flash",
-                    contents=summarize_prompt,
-                    config=GenerateContentConfig(
-                        system_instruction=["You are a helpful assistant who is an expert at summarization."]
-                    ),
-                )
-                request.chat_history = request.chat_history[-(summary_thresh+1):] # keep last k items
-                request.chat_history.insert(0, {"role": "user", "parts": [{"text": f"Here is a summary of this conversation so far: {summary_response.text}"}]})
+            # Retrieve relevant text
+            results = retrieve(request.message, vectorstore)
 
+            rag_prompt = f"""You have access to the following relevant information retrieved based on the user's query:
             
+            {"\n\n".join(results)}
+
+            Using the information above, answer the user's query as accurately as possible:
+
+            User's Query: {request.message}
+            """
+
+            # # summarize chat history
+            # summary_thresh = 2
+            # if len(request.chat_history) > summary_thresh:
+            #     summarize_prompt = f"""Please summarize the following chat history concisely, focusing on the key points and main topics discussed. Avoid 
+            #     unnecessary details and provide a clear, straightforward summary. {request.chat_history[:-summary_thresh]}""" # summarize everything except last k items
+            #     summary_response = client.models.generate_content(
+            #         model="gemini-2.0-flash",
+            #         contents=summarize_prompt,
+            #         config=GenerateContentConfig(
+            #             system_instruction=["You are a helpful assistant who is an expert at summarization."]
+            #         ),
+            #     )
+            #     request.chat_history = request.chat_history[-(summary_thresh+1):] # keep last k items
+            #     request.chat_history.insert(0, {"role": "user", "parts": [{"text": f"Here is a summary of this conversation so far: {summary_response.text}"}]})
+            
+            # remove the unfformatted user message     
+            del request.chat_history[-1]
+            # add the user message with RAG data
+            request.chat_history.append({"role": "user", "parts": [{"text": rag_prompt}]})
+
             response = client.models.generate_content(
                 model="gemini-2.0-flash",
                 contents=request.chat_history,
@@ -62,7 +107,11 @@ async def chat(request: ChatRequest):
                     system_instruction=[request.system_message]
                 ),
             )
-            return {"response": response.text, "dataset_str": concise_text_string[:75], "history": request.chat_history}
+
+            # put back the unformatted user message
+            request.chat_history.append({"role": "user", "parts": [{"text": request.message}]})
+
+            return {"response": response.text, "dataset_str": concise_text_string[:75], "history": request.chat_history, "RAG_prompt": rag_prompt_template}
 
         if request.model_choice == "HF":
             if hf_token: