From 5b7d1cb49bd473c7dbcf6e89f7d1b6fc8be1f5b1 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Thu, 11 Apr 2024 08:59:41 +0200
Subject: Improve prompt

---
 rag/generator/cohere.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

(limited to 'rag/generator/cohere.py')

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",
         ):
-- 
cgit v1.2.3-70-g09d2