Improve prompt

3 files changed, 12 insertions, 7 deletions

diff --git a/rag/generator/cohere.py b/rag/generator/cohere.py
index 7028b21..2ed2cf5 100644
--- a/rag/generator/cohere.py
+++ b/rag/generator/cohere.py
@@ -6,7 +6,7 @@ import cohere
 from loguru import logger as log
 
 from .abstract import AbstractGenerator
-from .prompt import Prompt
+from .prompt import ANSWER_INSTRUCTION, Prompt
 
 
 class Cohere(metaclass=AbstractGenerator):
@@ -15,8 +15,9 @@ class Cohere(metaclass=AbstractGenerator):
 
     def generate(self, prompt: Prompt) -> Generator[Any, Any, Any]:
         log.debug("Generating answer from cohere")
+        query = f"{prompt.query}\n\n{ANSWER_INSTRUCTION}"
         for event in self.client.chat_stream(
-            message=prompt.query,
+            message=query,
             documents=[asdict(d) for d in prompt.documents],
             prompt_truncation="AUTO",
         ):
diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py
index abd5599..6340235 100644
--- a/rag/generator/ollama.py
+++ b/rag/generator/ollama.py
@@ -7,7 +7,7 @@ from loguru import logger as log
 from rag.retriever.vector import Document
 
 from .abstract import AbstractGenerator
-from .prompt import Prompt
+from .prompt import ANSWER_INSTRUCTION, Prompt
 
 
 class Ollama(metaclass=AbstractGenerator):
@@ -23,13 +23,11 @@ class Ollama(metaclass=AbstractGenerator):
 
     def __metaprompt(self, prompt: Prompt) -> str:
         metaprompt = (
-            "Answer the question based only on the following context:"
+            "Answer the question based only on the following context:\n"
             "<context>\n"
             f"{self.__context(prompt.documents)}\n\n"
             "</context>\n"
-            "Given the context information and not prior knowledge, answer the question."
-            "If the context is irrelevant to the question, answer that you do not know "
-            "the answer to the question given the context.\n"
+            f"{ANSWER_INSTRUCTION}"
             f"Question: {prompt.query.strip()}\n\n"
             "Answer:"
         )
diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py
index 41f264e..fa007db 100644
--- a/rag/generator/prompt.py
+++ b/rag/generator/prompt.py
@@ -3,6 +3,12 @@ from typing import List
 
 from rag.retriever.vector import Document
 
+ANSWER_INSTRUCTION = (
+    "Given the context information and not prior knowledge, answer the question."
+    "If the context is irrelevant to the question, answer that you do not know "
+    "the answer to the question given the context and stop.\n"
+)
+
 
 @dataclass
 class Prompt: