From 0fc49c7e132d0aa3f55cfb67a91837aa4e87c561 Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Sat, 2 Nov 2024 16:44:32 -0400
Subject: [PATCH] Add streaming support to Ollama provider and implement
 generate_stream_text method

---
 simplemind/providers/ollama.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py
index 3e00c25..47b4268 100644
--- a/simplemind/providers/ollama.py
+++ b/simplemind/providers/ollama.py
@@ -26,6 +26,7 @@ class Ollama(BaseProvider):
     DEFAULT_MODEL = DEFAULT_MODEL
     DEFAULT_KWARGS = DEFAULT_KWARGS
     TIMEOUT = DEFAULT_TIMEOUT
+    supports_streaming = True
 
     def __init__(self, host_url: str | None = None):
         self.host_url = host_url or settings.OLLAMA_HOST_URL
@@ -116,3 +117,21 @@ class Ollama(BaseProvider):
         )
 
         return response.get("message", {}).get("content", "")
+
+    @logger
+    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+        # Prepare the messages.
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+
+        # Iterate over the response and yield the content.
+        for chunk in response:
+            yield chunk["message"]["content"]