Update app.py

app.py

@@ -11,31 +11,32 @@ from fastapi.middleware.cors import CORSMiddleware
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_community.chat_message_histories import ChatMessageHistory
-from langchain_core.chat_history import BaseChatMessageHistory
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain.schema import BaseChatMessageHistory
+from langchain.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnableWithMessageHistory
 from pinecone import Pinecone
 from pinecone_text.sparse import BM25Encoder
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.retrievers import PineconeHybridSearchRetriever
 from langchain.retrievers import ContextualCompressionRetriever
-from langchain_community.chat_models import ChatPerplexity
+from langchain_community.chat_models import ChatOpenAI
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
-from langchain_core.prompts import PromptTemplate
+from langchain.prompts import PromptTemplate
 import re
+from langchain_huggingface import HuggingFaceEmbeddings
 
 # Load environment variables
 load_dotenv(".env")
 USER_AGENT = os.getenv("USER_AGENT")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 SECRET_KEY = os.getenv("SECRET_KEY")
 PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
 SESSION_ID_DEFAULT = "abc123"
 
 # Set environment variables
 os.environ['USER_AGENT'] = USER_AGENT
-os.environ["GROQ_API_KEY"] = GROQ_API_KEY
+os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 os.environ["TOKENIZERS_PARALLELISM"] = 'true'
 
 # Initialize FastAPI app and CORS
@@ -74,6 +75,7 @@ bm25 = BM25Encoder().load("./mbzuai-policies.json")
 
 # Initialize models and retriever
 embed_model = HuggingFaceEmbeddings(model_name="jinaai/jina-embeddings-v3", model_kwargs={"trust_remote_code":True})
+
 retriever = PineconeHybridSearchRetriever(
     embeddings=embed_model, 
     sparse_encoder=bm25, 
@@ -83,11 +85,11 @@ retriever = PineconeHybridSearchRetriever(
 )
 
 # Initialize LLM
-llm = ChatPerplexity(temperature=0, pplx_api_key=GROQ_API_KEY, model="llama-3.1-sonar-large-128k-chat", max_tokens=512, max_retries=2)
+llm = ChatOpenAI(temperature=0, model_name="gpt-4o-mini", max_tokens=512)
 
 # Initialize Reranker
 model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
-compressor = CrossEncoderReranker(model=model, top_n=20)
+compressor = CrossEncoderReranker(model=model, top_n=10)
 
 compression_retriever = ContextualCompressionRetriever(
     base_compressor=compressor, base_retriever=retriever
@@ -122,15 +124,15 @@ When responding to queries, follow these guidelines:
 
 2. Formatting for Readability: 
    - Provide the entire response in proper markdown format.
-   - Use structured Maekdown elements such as headings, subheading, lists, tables, and links.
-   - Use emaphsis on headings, important texts and phrases.
+   - Use structured Markdown elements such as headings, subheadings, lists, tables, and links.
+   - Use emphasis on headings, important texts, and phrases.
 
 3. Proper Citations:
-   - ALWAYS USE INLINE CITATIONS with embed source URLs where users can verify information or explore further. 
-   - The inline citations should be in the format [1], [2], etc., in the response with links to reference sources.
+   - ALWAYS USE INLINE CITATIONS with embedded source URLs where users can verify information or explore further. 
+   - The inline citations should be in the format [[1]], [[2]], etc., in the response with links to reference sources.
+   - Then at the end of the response, list out the citations with their sources.
 
 FOLLOW ALL THE GIVEN INSTRUCTIONS, FAILURE TO DO SO WILL RESULT IN TERMINATION OF THE CHAT.
-
 {context}
 """
 qa_prompt = ChatPromptTemplate.from_messages(
@@ -165,7 +167,6 @@ conversational_rag_chain = RunnableWithMessageHistory(
     output_messages_key="answer",
 )
 
-
 # WebSocket endpoint with streaming
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):