Add type hints for generators

jackmpcollins · May 16, 2024 · 5ca0246 · 5ca0246
1 parent 9ad92c2
commit 5ca0246
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/magentic/chat_model/openai_chat_model.py b/src/magentic/chat_model/openai_chat_model.py
@@ -306,7 +306,7 @@ def _create_usage_pointer(
     """Returns a pointer to a Usage object that is created at the end of the response."""
     usage_pointer = []
 
-    def generator():
+    def generator() -> Iterator[ChatCompletionChunk]:
         for chunk in response:
             if chunk.usage:
                 usage = Usage(
@@ -325,7 +325,7 @@ def _acreate_usage_pointer(
     """Async version of `_create_usage_pointer`."""
     usage_pointer = []
 
-    async def generator():
+    async def generator() -> AsyncIterator[ChatCompletionChunk]:
         async for chunk in response:
             if chunk.usage:
                 usage = Usage(