Update src/simple_rag.py

src/simple_rag.py

@@ -40,21 +40,21 @@ CHROMA_PATH = os.path.join(WRITABLE_DIR, "src", "chroma")
 
 embedding_model = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-base")
 
-PROMPT_TEMPLATE = """
-You are a helpful assistant.
-Answer the question based ONLY on the context below.
-If the user asks in Khmer, respond in Khmer.
-If the user asks in English, respond in English.
-Use clear, concise sentences, no more than 50 word. Do not mention the existence of context.
+# PROMPT_TEMPLATE = """
+# You are a helpful assistant.
+# Answer the question based ONLY on the context below.
+# If the user asks in Khmer, respond in Khmer.
+# If the user asks in English, respond in English.
+# Use clear, concise sentences, no more than 50 word. Do not mention the existence of context.
 
-Context:
-{context}
+# Context:
+# {context}
 
-Question:
-{question}
+# Question:
+# {question}
 
-Answer:
-"""
+# Answer:
+# """
 
 def load_documents():
     loader = PyPDFDirectoryLoader(DATA_PATH)
@@ -102,24 +102,39 @@ def ask_question(query_text: str, k: int = 3):
         })
 
     context_text = "\n\n".join(chunk["text"] for chunk in context_chunks)
-    prompt = PROMPT_TEMPLATE.format(context=context_text, question=query_text)
-    logging.info(f"Prompt: {prompt}")
+    #prompt = PROMPT_TEMPLATE.format(context=context_text, question=query_text)
+    #logging.info(f"Prompt: {prompt}")
     
-    messages = [{"role": "user", "content": prompt}]
-    
-    logging.info("Sending prompt to model...")
-    logging.info(f"Messages: {messages}")
+    # Construct structured messages instead of using PROMPT_TEMPLATE
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are a helpful assistant. "
+                "Answer the question based ONLY on the context provided. "
+                "If the user asks in Khmer, respond in Khmer. "
+                "If the user asks in English, respond in English. "
+                "Use clear, concise sentences, max 50 words. "
+                "Do not mention context or metadata."
+            )
+        },
+        {
+            "role": "user",
+            "content": f"Context:\n{context_text}\n\nQuestion:\n{query_text}"
+        }
+    ]
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=False,
+        thinking_mode="off"
+    )
     
-    prompts = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=False,
-            thinking_mode="off"
-        )
-    logging.info(f"Prompts: {prompts}")
+    logging.info(f"Prompts: {prompt}")
 
     output = pipeline(
-        prompts,
+        prompt,
         max_new_tokens=128,
         do_sample=False,
         return_full_text=False,