Fix aiter & anext implementations.

Unity-google-new · Dec 15, 2023 · 3cd1d8b · 3cd1d8b
1 parent 7f1e165
commit 3cd1d8b
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/google/generativeai/types/generation_types.py b/google/generativeai/types/generation_types.py
@@ -30,18 +30,16 @@
 if sys.version_info < (3, 10):
 
     def aiter(obj):
-        yield obj.__next__()
+        return obj.__aiter__()
 
     async def anext(obj, default=None):
         try:
-            return obj.__anext__()
+            return await obj.__anext__()
         except StopAsyncIteration:
             if default is not None:
                 return default
             else:
-                raise StopAsyncIteration("stop iteration") from None
-        except Exception as e:
-            return obj.__next__()
+                raise
 
 
 class BlockedPromptException(Exception):