UPDATE: follow ups

functions.py

@@ -6,9 +6,9 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_qdrant import QdrantVectorStore
 from langchain_qdrant import RetrievalMode
 from langchain_core.prompts.chat import ChatPromptTemplate
-from uuid import uuid4
-from langchain_core.output_parsers import StrOutputParser
-from langchain.retrievers import ParentDocumentRetriever
+from langchain_core.prompts import PromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain.memory import ChatMessageHistory
 from pandasai import SmartDataframe
@@ -16,9 +16,7 @@ from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_community.document_loaders import YoutubeLoader
 from langchain.docstore.document import Document
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.retrievers import ContextualCompressionRetriever
 from langchain_qdrant import FastEmbedSparse
-from langchain.retrievers.document_compressors import FlashrankRerank
 from supabase.client import create_client
 from qdrant_client import QdrantClient
 from langchain_groq import ChatGroq
@@ -73,6 +71,24 @@ NOTE: Generate responses directly without using phrases like "Response:" or "Ans
 prompt = ChatPromptTemplate.from_template(prompt)
 chatHistoryStore = dict()
 
+class FollowUps(BaseModel):
+    q1: str = Field(description="First Follow-up Question")
+    q2: str = Field(description="Second Follow-up Question")
+    q3: str = Field(description="Third Follow-up Question")
+
+followUpPrompt = """
+You are an expert chatbot at framing follow up questions using some given text such that their answers can be found in the text itself and have been given the task of doing the same. Frame appropriate and meaningful questions out of it. Also, if no the given text says NO CONTEXT FOUND, please return an empty string for each question asked.
+\n{format_instructions}
+\n{context}
+"""
+jsonParser = JsonOutputParser(pydantic_object=FollowUps)
+followUpPrompt = PromptTemplate(
+    template=followUpPrompt,
+    input_variables=["context"],
+    partial_variables={"format_instructions": jsonParser.get_format_instructions()},
+)
+
+
 
 def createUser(user_id: str, username: str, email: str) -> dict:
     userData = client.table("ConversAI_UserInfo").select("*").execute().data
@@ -157,6 +173,7 @@ def addDocuments(texts: list[tuple[str]], vectorstore: str):
 
 def format_docs(docs: str):
     global sources
+    global tempContext
     sources = []
     context = ""
     for doc in docs:
@@ -169,6 +186,7 @@ def format_docs(docs: str):
     else:
         pass
     sources = list(set(sources))
+    tempContext = context
     return context
 
 
@@ -194,6 +212,9 @@ def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192")
     global prompt
     global client
     global sources
+    global jsonParser
+    global tempContext
+    global followUpPrompt
     global vectorEmbeddings
     global sparseEmbeddings
     vectorStoreName = vectorstore
@@ -211,7 +232,7 @@ def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192")
              "question": RunnableLambda(lambda x: x["question"]),
              "chatHistory": RunnableLambda(lambda x: x["chatHistory"])}
             | prompt
-            | ChatGroq(model=llmModel, temperature=0.75, max_tokens=512)
+            | ChatGroq(model_name=llmModel, temperature=0.75, max_tokens=512)
             | StrOutputParser()
     )
     messageChain = RunnableWithMessageHistory(
@@ -221,11 +242,15 @@ def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192")
         history_messages_key="chatHistory"
     )
     chain = RunnablePassthrough.assign(messages_trimmed=trimMessages) | messageChain
-    return {
-        "output": chain.invoke(
+    followUpChain = prompt | ChatGroq(model_name="gemma2-9b-it", temperature=0) | jsonParser
+    output = chain.invoke(
             {"question": query},
             {"configurable": {"session_id": vectorStoreName}}
-        ),
+        )
+    followUpQuestions = followUpChain.invoke({"context": tempContext})
+    return {
+        "output": output,
+        "followUpQuestions": followUpQuestions,
         "sources": sources
     }