From aac821b148c6c0d35b940609dc9b0ddcb053b28e Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Wed, 19 Jun 2024 02:07:06 +0200 Subject: Still wip on rewrite --- rag/generator/ollama.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) (limited to 'rag/generator/ollama.py') diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py index ff5402b..84563bb 100644 --- a/rag/generator/ollama.py +++ b/rag/generator/ollama.py @@ -4,10 +4,10 @@ from typing import Any, Generator, List import ollama from loguru import logger as log -from rag.rag import Message +from rag.message import Message +from rag.retriever.vector import Document from .abstract import AbstractGenerator -from .prompt import Prompt class Ollama(metaclass=AbstractGenerator): @@ -16,12 +16,9 @@ class Ollama(metaclass=AbstractGenerator): log.debug(f"Using {self.model} for generator...") def generate( - self, prompt: Prompt, messages: List[Message] + self, messages: List[Message], documents: List[Document] ) -> Generator[Any, Any, Any]: log.debug("Generating answer with ollama...") - messages = messages.append( - Message(role="user", content=prompt.to_str(), client="ollama") - ) messages = [m.as_dict() for m in messages] for chunk in ollama.chat(model=self.model, messages=messages, stream=True): - yield chunk["response"] + yield chunk["message"]["content"] -- cgit v1.2.3-70-g09d2