polyfill anext

Kahlain · Sep 4, 2024 · de9f366 · de9f366
1 parent 8948d94
commit de9f366
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/guardrails/integrations/databricks/ml_flow_instrumentor.py b/guardrails/integrations/databricks/ml_flow_instrumentor.py
@@ -1,5 +1,6 @@
 from functools import wraps
 import inspect
+import sys
 from typing import (
     Any,
     AsyncIterable,
@@ -39,6 +40,10 @@
     raise ImportError("Please install mlflow to use this instrumentor")
 
 
+if sys.version_info.minor < 10:
+    from guardrails.utils.polyfills import anext
+
+
 # TODO: Abstract these methods and common logic into a base class
 #   that can be extended by other instrumentors
 class MlFlowInstrumentor:
@@ -300,7 +305,7 @@ async def trace_async_stream_step_wrapper(
                     next_exists = True
                     while next_exists:
                         try:
-                            res = await anext(gen)  # type: ignore
+                            res = await anext(gen)
                             yield res
                         except StopIteration:
                             next_exists = False

diff --git a/guardrails/utils/polyfills.py b/guardrails/utils/polyfills.py
@@ -0,0 +1,2 @@
+def anext(aiter):
+    return aiter.__anext__()