Update app.py

app.py

@@ -2,10 +2,11 @@ import streamlit as st
 import asyncio
 import os
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain.chat_models import ChatOpenAI
+from langchain.schema import SystemMessage, HumanMessage
 from PyPDF2 import PdfReader
 import aiohttp
 from io import BytesIO
@@ -15,10 +16,12 @@ os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
 
 # Set up prompts
 system_template = "Use the following context to answer a user's question. If you cannot find the answer in the context, say you don't know the answer."
-system_role_prompt = SystemMessagePromptTemplate.from_template(system_template)
+system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)
 
-user_prompt_template = "Context:\n{context}\n\nQuestion:\n{question}"
-user_role_prompt = HumanMessagePromptTemplate.from_template(user_prompt_template)
+human_template = "Context:\n{context}\n\nQuestion:\n{question}"
+human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
+
+chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
 
 # Define RetrievalAugmentedQAPipeline class
 class RetrievalAugmentedQAPipeline:
@@ -35,10 +38,9 @@ class RetrievalAugmentedQAPipeline:
         if len(context_prompt) > max_context_length:
             context_prompt = context_prompt[:max_context_length]
         
-        formatted_system_prompt = system_role_prompt.format()
-        formatted_user_prompt = user_role_prompt.format(question=user_query, context=context_prompt)
+        messages = chat_prompt.format_prompt(context=context_prompt, question=user_query).to_messages()
 
-        response = await self.llm.agenerate([formatted_system_prompt, formatted_user_prompt])
+        response = await self.llm.agenerate([messages])
         return {"response": response.generations[0][0].text, "context": context_list}
 
 # PDF processing functions