Async Function Support for Tools Parameter in GenerativeModel

google/generativeai/types/content_types.py

@@ -14,7 +14,8 @@
 
 
 from __future__ import annotations
-
+# we are importing asyncio here to use asynchronous functions
+import asyncio
 from collections.abc import Iterable, Mapping, Sequence
 import io
 import inspect
@@ -607,11 +608,7 @@ def from_function(function: Callable[..., Any], descriptions: dict[str, str] | N
 
 
 class CallableFunctionDeclaration(FunctionDeclaration):
-    """An extension of `FunctionDeclaration` that can be built from a python function, and is callable.
-
-    Note: The python function must have type annotations.
-    """
-
+    """An extension of `FunctionDeclaration` that can be built from a python function, and is callable."""
     def __init__(
         self,
         *,
@@ -622,12 +619,31 @@ def __init__(
     ):
         super().__init__(name=name, description=description, parameters=parameters)
         self.function = function
+        # class variable to check if the passed tool function is asynchronous function or not
+        self.is_async = inspect.iscoroutinefunction(function)
+
+    async def __call__(self, fc: protos.FunctionCall) -> protos.FunctionResponse:
+        try:
+            # handling async function seperately 
+            if self.is_async:
+                result = await self.function(**fc.args)
+            else:
+                result = self.function(**fc.args)
+
+            if not isinstance(result, dict):
+                result = {"result": result}
+
+            return protos.FunctionResponse(name=fc.name, response=result)
+        except Exception as e:
+            error_result = {
+                "error": str(e),
+                "type": type(e).__name__
+            }
+            return protos.FunctionResponse(name=fc.name, response=error_result)
+
+
+
 
-    def __call__(self, fc: protos.FunctionCall) -> protos.FunctionResponse:
-        result = self.function(**fc.args)
-        if not isinstance(result, dict):
-            result = {"result": result}
-        return protos.FunctionResponse(name=fc.name, response=result)
 
 
 FunctionDeclarationType = Union[
@@ -758,10 +774,14 @@ def __getitem__(
 
         return self._index[name]
 
-    def __call__(self, fc: protos.FunctionCall) -> protos.FunctionResponse | None:
+    async def __call__(self, fc: protos.FunctionCall) -> protos.Part | None:
         declaration = self[fc]
         if not callable(declaration):
             return None
+        response = await declaration(fc)
+        return protos.Part(function_response=response)
+
+
 
         return declaration(fc)
 
@@ -833,10 +853,8 @@ def _make_tool(tool: ToolType) -> Tool:
                 f"Object Value: {tool}"
             ) from e
 
-
 class FunctionLibrary:
     """A container for a set of `Tool` objects, manages lookup and execution of their functions."""
-
     def __init__(self, tools: Iterable[ToolType]):
         tools = _make_tools(tools)
         self._tools = list(tools)
@@ -856,21 +874,30 @@ def __getitem__(
     ) -> FunctionDeclaration | protos.FunctionDeclaration:
         if not isinstance(name, str):
             name = name.name
-
         return self._index[name]
 
-    def __call__(self, fc: protos.FunctionCall) -> protos.Part | None:
+    def __call__(self, fc: protos.FunctionCall) -> protos.Part:
         declaration = self[fc]
         if not callable(declaration):
             return None
-
-        response = declaration(fc)
+
+        # Handle both sync and async functions
+        if inspect.iscoroutinefunction(declaration.__call__):
+            # For async functions, run in an event loop
+            loop = asyncio.get_event_loop()
+            response = loop.run_until_complete(declaration(fc))
+        else:
+            # For sync functions, call directly
+            response = declaration(fc)
+
+        # Convert response to Part
+        if response is None:
+            return None
         return protos.Part(function_response=response)
 
     def to_proto(self):
         return [tool.to_proto() for tool in self._tools]
-
-
+
 ToolsType = Union[Iterable[ToolType], ToolType]