Update prompt

utils/llm_generation.py

@@ -3,6 +3,7 @@ from langchain.prompts import (
     SystemMessagePromptTemplate,
     HumanMessagePromptTemplate,
 )
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_community.chat_models import ChatOpenAI
 from langchain.chains import LLMChain
 from langchain.schema import Document
@@ -38,31 +39,43 @@ class LLMGeneration:
             },
         )
 
-        self.create_system_message()
+        self.create_initial_prompt()
 
-    def create_system_message(self):
+    def create_initial_prompt(self):
         # System message for the chain
-        system_message = (
-            "You are a helpful assistant. Use ONLY the provided context to answer the user's question. "
-            'If you cannot find the answer in the context, return answer key as  "Data Not Available".'
-        )
-        self.system_message_prompt = SystemMessagePromptTemplate.from_template(
-            system_message
+        system_message = SystemMessage(
+            content=(
+                "You are a helpful assistant. Use ONLY the provided context to answer the user's question mentioned in <question> tag"
+                """Answer the following question in a structured JSON format"""
+                "Adhere to the json schema stricly. Make sure the answer is exact match."
+                """If you do not find the answer in the provided contexts mentioned in <context> tag, return '{"answer": "Data Not Available"}'."""
+            )
         )
 
+        few_shots = [
+            HumanMessage(
+                content="<context>John traveled to Paris last summer. He stayed at a small boutique hotel and visited the Louvre museum.</context>"
+                "<question>Where did John travel?</question>"
+            ),
+            AIMessage(
+                content="""{"answer": "Paris"}"""
+            ),
+        ]
+
+        self.initial_prompt_messages = [system_message] + few_shots
+
     def create_human_message_prompt(self, query: str, docs: List[Document]):
 
         # Prepare the context from the retrieved chunks
-        context = "\n\n".join([doc.page_content for doc in docs])
+        context = "\n\n".join(
+            [f"<context>{doc.page_content}</context>" for doc in docs]
+        )
 
         # Human message, instructing the assistant to use the context and produce a structured JSON answer
-        human_message = f"""Given the following context:
-
+        human_message = f"""
         {context}
 
-        Answer the following question in a structured JSON format with a key called "answer". If you do not find the answer in the provided context, return "answer" as "No answer".
-
-        Question: {query}"""
+        <question>{query}</question>"""
 
         return HumanMessagePromptTemplate.from_template(human_message)
 
@@ -70,7 +83,8 @@ class LLMGeneration:
 
         # Create the prompt template
         prompt = ChatPromptTemplate.from_messages(
-            [self.system_message_prompt, self.create_human_message_prompt(query, docs)]
+            self.initial_prompt_messages
+            + [self.create_human_message_prompt(query, docs)]
         )
 
         # Create and run the chain with the hypothetical gpt-40-mini model
@@ -81,4 +95,4 @@ class LLMGeneration:
         result = chain.invoke({})
 
         print(f"Query: {query} Result: {result}")
-        return result.get("answer", "")
+        return json.loads(result.get("text", "{}")).get("answer", "")