From 1da5811415250befc8c726d61cc17376a6fbb0f5 Mon Sep 17 00:00:00 2001 From: "Som S." Date: Tue, 12 Nov 2024 11:49:30 -0700 Subject: [PATCH 1/3] Async Function Support for Tools Parameter in GenerativeModel --- google/generativeai/types/content_types.py | 64 +++++++++++++++------- 1 file changed, 44 insertions(+), 20 deletions(-) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index f3db610e1..7ad23b664 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -14,7 +14,7 @@ from __future__ import annotations - +import asyncio from collections.abc import Iterable, Mapping, Sequence import io import inspect @@ -607,11 +607,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 +618,29 @@ def __init__( ): super().__init__(name=name, description=description, parameters=parameters) self.function = function + self.is_async = inspect.iscoroutinefunction(function) + + async def __call__(self, fc: protos.FunctionCall) -> protos.FunctionResponse: + try: + 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 +771,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 +850,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 +871,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] From 3cd8702c2bc5fd003a452f765768478c868370df Mon Sep 17 00:00:00 2001 From: "Som S." Date: Tue, 12 Nov 2024 11:53:25 -0700 Subject: [PATCH 2/3] Async Function Support for Tools Parameter in GenerativeModel --- google/generativeai/types/content_types.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index 7ad23b664..4dee56546 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -14,6 +14,7 @@ from __future__ import annotations +# we are importing asyncio here to use asynchronous functions import asyncio from collections.abc import Iterable, Mapping, Sequence import io @@ -618,10 +619,12 @@ 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: From b0ec6bc30724fbace66e97177218da727813c1f7 Mon Sep 17 00:00:00 2001 From: "Som S." Date: Sat, 16 Nov 2024 17:51:10 -0700 Subject: [PATCH 3/3] asyncionesting loop approach --- google/generativeai/types/content_types.py | 54 +++++++++++----------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index 4dee56546..def601a9a 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -608,40 +608,47 @@ 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.""" def __init__( - self, - *, - name: str, - description: str, + self, + *, + name: str, + description: str, parameters: dict[str, Any] | None = None, function: Callable[..., Any], ): 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: + def __call__(self, fc: protos.FunctionCall) -> protos.FunctionResponse: + """Handles both sync and async function calls transparently""" try: - # handling async function seperately + # Get or create event loop + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + # Execute function based on type if self.is_async: - result = await self.function(**fc.args) + result = loop.run_until_complete(self._run_async(fc)) else: result = self.function(**fc.args) + # Format response 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) - + return protos.FunctionResponse( + name=fc.name, + response={"error": str(e), "type": type(e).__name__} + ) + async def _run_async(self, fc: protos.FunctionCall): + """Helper method to run async functions""" + return await self.function(**fc.args) @@ -854,11 +861,11 @@ def _make_tool(tool: ToolType) -> 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) self._index = {} + for tool in self._tools: for declaration in tool.function_declarations: name = declaration.name @@ -869,6 +876,7 @@ def __init__(self, tools: Iterable[ToolType]): ) self._index[declaration.name] = declaration + def __getitem__( self, name: str | protos.FunctionCall ) -> FunctionDeclaration | protos.FunctionDeclaration: @@ -881,18 +889,10 @@ def __call__(self, fc: protos.FunctionCall) -> protos.Part: if not callable(declaration): return None - # 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 + response = declaration(fc) if response is None: return None + return protos.Part(function_response=response) def to_proto(self):