From 25ba1a928936495696e4ae773463b72a55679e28 Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Sat, 2 Nov 2024 17:14:53 -0400
Subject: [PATCH] update function annotations

---
 simplemind/providers/amazon.py    | 4 ++--
 simplemind/providers/anthropic.py | 6 ++++--
 simplemind/providers/gemini.py    | 4 ++--
 simplemind/providers/groq.py      | 4 ++--
 simplemind/providers/ollama.py    | 6 ++++--
 simplemind/providers/openai.py    | 5 ++---
 simplemind/providers/xai.py       | 6 ++++--
 7 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/simplemind/providers/amazon.py b/simplemind/providers/amazon.py
index afb6888..87d0b3c 100644
--- a/simplemind/providers/amazon.py
+++ b/simplemind/providers/amazon.py
@@ -1,4 +1,4 @@
-from typing import Type, TypeVar
+from typing import Type, TypeVar, Iterator
 from functools import cached_property
 
 import instructor
@@ -94,7 +94,7 @@ class Amazon(BaseProvider):
 
         return response.content[0].text
 
-    def generate_stream_text(self, prompt, *, llm_model, **kwargs):
+    def generate_stream_text(self, prompt, *, llm_model, **kwargs) -> Iterator[str]:
         """Generate streaming text using the Amazon API."""
 
         # Prepare the messages.
diff --git a/simplemind/providers/anthropic.py b/simplemind/providers/anthropic.py
index 0809913..ffa776d 100644
--- a/simplemind/providers/anthropic.py
+++ b/simplemind/providers/anthropic.py
@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from pydantic import BaseModel
@@ -110,7 +110,9 @@ class Anthropic(BaseProvider):
         return response.content[0].text
 
     @logger
-    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs):
+    def generate_stream_text(
+        self, prompt: str, *, llm_model: str, **kwargs
+    ) -> Iterator[str]:
         # Prepare the messages.
         messages = [
             {"role": "user", "content": prompt},
diff --git a/simplemind/providers/gemini.py b/simplemind/providers/gemini.py
index 78a014a..0123bf4 100644
--- a/simplemind/providers/gemini.py
+++ b/simplemind/providers/gemini.py
@@ -2,7 +2,7 @@
 # IT is not currently working as desired.
 
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from pydantic import BaseModel
@@ -110,7 +110,7 @@ class Gemini(BaseProvider):
         return response.text
 
     @logger
-    def generate_stream_text(self, prompt: str, **kwargs) -> str:
+    def generate_stream_text(self, prompt: str, **kwargs) -> Iterator[str]:
         """Generate streaming text using the Gemini API."""
         kwargs.pop("llm_model", None)
         try:
diff --git a/simplemind/providers/groq.py b/simplemind/providers/groq.py
index dfbf69d..45593aa 100644
--- a/simplemind/providers/groq.py
+++ b/simplemind/providers/groq.py
@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from pydantic import BaseModel
@@ -120,7 +120,7 @@ class Groq(BaseProvider):
         *,
         llm_model: str | None = None,
         **kwargs,
-    ) -> str:
+    ) -> Iterator[str]:
         """Generate streaming text using the Groq API."""
         messages = [
             {"role": "user", "content": prompt},
diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py
index 47b4268..84df560 100644
--- a/simplemind/providers/ollama.py
+++ b/simplemind/providers/ollama.py
@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from openai import OpenAI
@@ -119,7 +119,9 @@ class Ollama(BaseProvider):
         return response.get("message", {}).get("content", "")
 
     @logger
-    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+    def generate_stream_text(
+        self, prompt: str, *, llm_model: str, **kwargs
+    ) -> Iterator[str]:
         # Prepare the messages.
         messages = [
             {"role": "user", "content": prompt},
diff --git a/simplemind/providers/openai.py b/simplemind/providers/openai.py
index 25ed2af..97b4a9a 100644
--- a/simplemind/providers/openai.py
+++ b/simplemind/providers/openai.py
@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from pydantic import BaseModel
@@ -24,7 +24,6 @@ class OpenAI(BaseProvider):
     DEFAULT_MODEL = DEFAULT_MODEL
     DEFAULT_KWARGS = DEFAULT_KWARGS
     supports_streaming = True
-    
 
     def __init__(self, api_key: str | None = None):
         self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
@@ -112,7 +111,7 @@ class OpenAI(BaseProvider):
     @logger
     def generate_stream_text(
         self, prompt: str, *, llm_model: str | None = None, **kwargs
-    ):
+    ) -> Iterator[str]:
         """Generate streaming text using the OpenAI API.
 
         Yields chunks of text as they are generated by the model.
diff --git a/simplemind/providers/xai.py b/simplemind/providers/xai.py
index 4db0d64..279c150 100644
--- a/simplemind/providers/xai.py
+++ b/simplemind/providers/xai.py
@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import TYPE_CHECKING, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar, Iterator
 
 import instructor
 from pydantic import BaseModel
@@ -103,7 +103,9 @@ class XAI(BaseProvider):
         return str(response.choices[0].message.content)
 
     @logger
-    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+    def generate_stream_text(
+        self, prompt: str, *, llm_model: str, **kwargs
+    ) -> Iterator[str]:
         # Prepare the messages.
         messages = [
             {"role": "user", "content": prompt},