diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-04-13 02:26:01 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-04-13 02:26:01 +0200 |
commit | 72d1caf92115d90ae789de1cffed29406f2a0a39 (patch) | |
tree | ca4f3b755dccdd94894f18e7ce599cfd7eb28e58 /rag/generator/ollama.py | |
parent | 36722903391ec42d5458112bc0549eb843548d90 (diff) |
Wip chat ui
Diffstat (limited to 'rag/generator/ollama.py')
-rw-r--r-- | rag/generator/ollama.py | 9 |
1 files changed, 8 insertions, 1 deletions
diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py index 6340235..b475dcf 100644 --- a/rag/generator/ollama.py +++ b/rag/generator/ollama.py @@ -1,5 +1,5 @@ import os -from typing import Any, Generator, List +from typing import Any, Dict, Generator, List import ollama from loguru import logger as log @@ -38,3 +38,10 @@ class Ollama(metaclass=AbstractGenerator): metaprompt = self.__metaprompt(prompt) for chunk in ollama.generate(model=self.model, prompt=metaprompt, stream=True): yield chunk["response"] + + def chat(self, prompt: Prompt, messages: List[Dict[str, str]]) -> Generator[Any, Any, Any]: + log.debug("Generating answer with ollama...") + metaprompt = self.__metaprompt(prompt) + messages.append({"role": "user", "content": metaprompt}) + for chunk in ollama.chat(model=self.model, messages=messages, stream=True): + yield chunk["message"]["content"] |