Update provider typing. Ignore PyCharm static type checking bug.

Author: djcopley

shelloracle/provider.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncIterator
 from typing import Protocol
 
 
@@ -17,7 +17,7 @@ class Provider(Protocol):
     name: str
 
     @abstractmethod
-    def generate(self, prompt: str) -> AsyncGenerator[str, None, None]:
+    async def generate(self, prompt: str) -> AsyncIterator[str]:
         """
         This is an asynchronous generator method which defines the protocol that a provider implementation
         should adhere to. The method takes a prompt as an argument and produces an asynchronous stream
@@ -36,6 +36,6 @@ def get_provider(name: str) -> type[Provider]:
     """
     from .providers import Ollama
     providers = {
-        "Ollama": Ollama
+        Ollama.name: Ollama,
     }
     return providers[name]

shelloracle/providers/ollama.py

@@ -1,7 +1,6 @@
 import json
-from collections.abc import AsyncGenerator
 from dataclasses import dataclass, asdict
-from typing import Any
+from typing import Any, AsyncIterator
 
 import httpx
 
@@ -70,7 +69,7 @@ def endpoint(self) -> str:
         # computed property because python descriptors need to be bound to an instance before access
         return f"http://{self.host}:{self.port}/api/generate"
 
-    async def generate(self, prompt: str) -> AsyncGenerator[str, None, None]:
+    async def generate(self, prompt: str) -> AsyncIterator[str]:
         request = GenerateRequest(self.model, prompt, system=self.system_prompt, stream=True)
         data = dataclass_to_json(request)
         try: