fix: make async custom functions work (#1255)

Author: georgeh0

python/cocoindex/runtime.py

@@ -69,6 +69,10 @@ def run(self, coro: Coroutine[Any, Any, T]) -> T:
 
 
 def to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
-    if inspect.iscoroutinefunction(call):
+    if isinstance(call, (staticmethod, classmethod)):
+        base_call = call.__func__
+    else:
+        base_call = call
+    if inspect.iscoroutinefunction(base_call):
         return call
     return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))

python/cocoindex/tests/test_transform_flow.py

@@ -268,3 +268,33 @@ def transform_flow_with_analyze_prepare(
 
     result = transform_flow_with_analyze_prepare.eval("Hello")
     expected = "Hello world!!"
+
+
+@cocoindex.op.function()
+async def async_custom_function(text: str) -> str:
+    """Append ' world' to the input text."""
+    return f"{text} world"
+
+
+class AsyncCustomFunctionSpec(cocoindex.op.FunctionSpec):
+    suffix: str
+
+
+@cocoindex.op.executor_class()
+class AsyncAppendSuffixExecutor:
+    spec: AsyncCustomFunctionSpec
+
+    async def __call__(self, text: str) -> str:
+        return f"{text}{self.spec.suffix}"
+
+
+def test_async_custom_function() -> None:
+    @cocoindex.transform_flow()
+    def transform_flow(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[str]:
+        return text.transform(async_custom_function).transform(
+            AsyncCustomFunctionSpec(suffix="!")
+        )
+
+    result = transform_flow.eval("Hello")
+    expected = "Hello world!"
+    assert result == expected, f"Expected {expected}, got {result}"