Update app/main.py

app/main.py

@@ -9,15 +9,11 @@ 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, Depends
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
-import nltk  # Importing NLTK
+import 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,18 +43,10 @@ except Exception as e:
     raise
 
 def clean_response(response):
-    # Remove any leading/trailing whitespace, including newlines
     cleaned = response.strip()
-    
-    # Remove any enclosing quotation marks
-    cleaned = re.sub(r'^\"+|\"+$', '', cleaned)
-    
-    # Replace multiple newlines with a single newline
+    cleaned = re.sub(r'^\"|\"$', '', cleaned)
     cleaned = re.sub(r'\n+', '\n', cleaned)
-    
-    # Remove any remaining '\n' characters
     cleaned = cleaned.replace('\\n', '')
-    
     return cleaned
 
 app = FastAPI()
@@ -71,9 +59,6 @@ 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,
@@ -81,39 +66,28 @@ llm = ChatOpenAI(
     temperature=0.7
 )
 
+conversation_history = {}  # Dictionary to maintain contextual memory
+
 @app.get("/")
 def read_root():
     return {"Hello": "World"}
 
 class Query(BaseModel):
+    session_id: str  # Unique identifier for user session
     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_template = 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.'
 
-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"
 
-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>
-Question: {input}  
-"""
+    <context>
+    {context}
+    </context>
+    Question: {input}
+    """
 )
 
 def vector_embedding():
@@ -130,16 +104,16 @@ def vector_embedding():
 
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
         chunks = text_splitter.split_documents(documents)
-        
+
         print(f"Created {len(chunks)} chunks.")
 
         model_name = "BAAI/bge-base-en"
         encode_kwargs = {'normalize_embeddings': True}
         model_norm = HuggingFaceBgeEmbeddings(model_name=model_name, encode_kwargs=encode_kwargs)
-        
+
         db = FAISS.from_documents(chunks, model_norm)
         db.save_local("./vectors_db")
-        
+
         print("Vector store created and saved successfully.")
         return {"response": "Vector Store DB Is Ready"}
 
@@ -154,9 +128,12 @@ def get_embeddings():
     return model_norm
 
 @app.post("/chat")
-def read_item(query: Query, request: Request):
+def chat_endpoint(query: Query):
     try:
-        user_id, user_context = get_user_context(request)
+        session_id = query.session_id
+        if session_id not in conversation_history:
+            conversation_history[session_id] = []
+
         embeddings = get_embeddings()
         vectors = FAISS.load_local("./vectors_db", embeddings, allow_dangerous_deserialization=True)
     except Exception as e:
@@ -166,22 +143,21 @@ def read_item(query: Query, request: Request):
     prompt1 = query.query_text
     if prompt1:
         start = time.process_time()
-        document_chain = create_stuff_documents_chain(llm, prompt)
+        document_chain = create_stuff_documents_chain(llm, prompt_template)
         retriever = vectors.as_retriever()
         retrieval_chain = create_retrieval_chain(retriever, document_chain)
 
-        # 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)
+        # Combine context from conversation history
+        context = "\n".join(conversation_history[session_id])
+        response = retrieval_chain.invoke({'input': prompt1, 'context': context})
 
-        # Apply the cleaning function to the response
         cleaned_response = clean_response(response['answer'])
 
-        # Update context
-        user_context["context"] = cleaned_response
+        # Update conversation history
+        conversation_history[session_id].append(f"User: {prompt1}")
+        conversation_history[session_id].append(f"Assistant: {cleaned_response}")
 
-        print("Cleaned response:", repr(cleaned_response))
+        print("Response time:", time.process_time() - start)
         return {"response": cleaned_response}
     else:
         return {"response": "No Query Found"}