1 files changed, 4 insertions, 7 deletions
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"]