From 72d1caf92115d90ae789de1cffed29406f2a0a39 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sat, 13 Apr 2024 02:26:01 +0200 Subject: Wip chat ui --- rag/generator/ollama.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) (limited to 'rag/generator/ollama.py') 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"] -- cgit v1.2.3-70-g09d2