diff options
Diffstat (limited to 'rag/generator/cohere.py')
| -rw-r--r-- | rag/generator/cohere.py | 22 |
1 files changed, 20 insertions, 2 deletions
diff --git a/rag/generator/cohere.py b/rag/generator/cohere.py index 2ed2cf5..16dfe88 100644 --- a/rag/generator/cohere.py +++ b/rag/generator/cohere.py @@ -1,6 +1,6 @@ import os from dataclasses import asdict -from typing import Any, Generator +from typing import Any, Dict, Generator, List import cohere from loguru import logger as log @@ -14,7 +14,7 @@ class Cohere(metaclass=AbstractGenerator): self.client = cohere.Client(os.environ["COHERE_API_KEY"]) def generate(self, prompt: Prompt) -> Generator[Any, Any, Any]: - log.debug("Generating answer from cohere") + log.debug("Generating answer from cohere...") query = f"{prompt.query}\n\n{ANSWER_INSTRUCTION}" for event in self.client.chat_stream( message=query, @@ -27,3 +27,21 @@ class Cohere(metaclass=AbstractGenerator): yield event.citations elif event.event_type == "stream-end": yield event.finish_reason + + def chat( + self, prompt: Prompt, messages: List[Dict[str, str]] + ) -> 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=query, + documents=[asdict(d) for d in prompt.documents], + chat_history=messages, + prompt_truncation="AUTO", + ): + if event.event_type == "text-generation": + yield event.text + # elif event.event_type == "citation-generation": + # yield event.citations + elif event.event_type == "stream-end": + yield event.finish_reason |