Update app/main.py

app/main.py

@@ -9,11 +9,15 @@ from langchain.chains import create_retrieval_chain
 from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import UnstructuredWordDocumentLoader as DocxLoader
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends
 from pydantic import BaseModel
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 import nltk  # Importing NLTK
 import time
+from typing import Dict, Optional
+from fastapi.sessions import SessionMiddleware
+from fastapi.requests import Request
+from fastapi.responses import JSONResponse
 
 # Set writable paths for cache and data
 cache_dir = '/tmp'
@@ -47,7 +51,7 @@ def clean_response(response):
     cleaned = response.strip()
     
     # Remove any enclosing quotation marks
-    cleaned = re.sub(r'^["\']+|["\']+$', '', cleaned)
+    cleaned = re.sub(r'^\"+|\"+$', '', cleaned)
     
     # Replace multiple newlines with a single newline
     cleaned = re.sub(r'\n+', '\n', cleaned)
@@ -67,10 +71,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Adding session middleware for contextual memory
+app.add_middleware(SessionMiddleware, secret_key="your-secret-key")
+
 openai_api_key = os.environ.get('OPENAI_API_KEY')
 llm = ChatOpenAI(
     api_key=openai_api_key,
-    model_name="gpt-4-turbo-preview",  # or "gpt-3.5-turbo" for a more economical option
+    model_name="gpt-4-turbo-preview",
     temperature=0.7
 )
 
@@ -81,13 +88,27 @@ def read_root():
 class Query(BaseModel):
     query_text: str
 
+# In-memory storage for contextual memory
+user_sessions: Dict[str, Dict[str, str]] = {}
+
+def get_user_context(request: Request):
+    user_id = request.client.host
+    if user_id not in user_sessions:
+        user_sessions[user_id] = {}
+    return user_id, user_sessions[user_id]
+
 prompt = ChatPromptTemplate.from_template(
 """
 You are a helpful assistant designed specifically for the Thapar Institute of Engineering and Technology (TIET), a renowned technical college. Your task is to answer all queries related to TIET. Every response you provide should be relevant to the context of TIET. If a question falls outside of this context, please decline by stating, 'Sorry, I cannot help with that.' If you do not know the answer to a question, do not attempt to fabricate a response; instead, politely decline.
 You may elaborate on your answers slightly to provide more information, but avoid sounding boastful or exaggerating. Stay focused on the context provided.
+
 If the query is not related to TIET or falls outside the context of education, respond with:
         "Sorry, I cannot help with that. I'm specifically designed to answer questions about the Thapar Institute of Engineering and Technology. 
         For more information, please contact at our toll-free number: 18002024100 or E-mail us at admissions@thapar.edu
+
+Previous Context:
+{previous_context}
+
 <context>
 {context}
 </context>
@@ -132,33 +153,38 @@ def get_embeddings():
     model_norm = HuggingFaceBgeEmbeddings(model_name=model_name, encode_kwargs=encode_kwargs)
     return model_norm
 
-@app.post("/chat")  # Changed from /anthropic to /chat
-def read_item(query: Query):
+@app.post("/chat")
+def read_item(query: Query, request: Request):
     try:
+        user_id, user_context = get_user_context(request)
         embeddings = get_embeddings()
         vectors = FAISS.load_local("./vectors_db", embeddings, allow_dangerous_deserialization=True)
     except Exception as e:
         print(f"Error loading vector store: {str(e)}")
         return {"response": "Vector Store Not Found or Error Loading. Please run /setup first."}
-    
+
     prompt1 = query.query_text
     if prompt1:
         start = time.process_time()
         document_chain = create_stuff_documents_chain(llm, prompt)
         retriever = vectors.as_retriever()
         retrieval_chain = create_retrieval_chain(retriever, document_chain)
-        response = retrieval_chain.invoke({'input': prompt1})
+
+        # Add previous context
+        previous_context = user_context.get("context", "None")
+        response = retrieval_chain.invoke({'input': prompt1, 'previous_context': previous_context})
         print("Response time:", time.process_time() - start)
-        
+
         # Apply the cleaning function to the response
         cleaned_response = clean_response(response['answer'])
-        
-        # For debugging, print the cleaned response
+
+        # Update context
+        user_context["context"] = cleaned_response
+
         print("Cleaned response:", repr(cleaned_response))
-        
-        return cleaned_response
+        return {"response": cleaned_response}
     else:
-        return "No Query Found"
+        return {"response": "No Query Found"}
 
 @app.get("/setup")
 def setup():