Async Function Support for Tools Parameter in GenerativeModel

[somwrks]

Description

The current implementation of the GenerativeModel class doesn't properly handle async functions when passed as tools, resulting in coroutine objects never being awaited and causing runtime errors.

Problem

When passing async functions as tools to the GenerativeModel, the following errors occur:

1. RuntimeWarning: coroutine was never awaited
2. Error: Unable to coerce value: <coroutine object>
3. Parameter to MergeFrom() must be instance of same class: expected <class 'Part'> got <class 'coroutine'>

Root Cause

The CallableFunctionDeclaration class and FunctionLibrary class in content_types.py don't properly handle coroutines returned by async functions. The current implementation attempts to use async function returns directly without awaiting them.

Solution

Modified the following classes in content_types.py to properly handle async functions:

class CallableFunctionDeclaration(FunctionDeclaration):
    def __init__(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
        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)

And the FunctionLibrary class:

class FunctionLibrary:
    def __call__(self, fc: protos.FunctionCall) -> protos.Part | None:
        declaration = self[fc]
        if not callable(declaration):
            return None
            
        if inspect.iscoroutinefunction(declaration.__call__):
            loop = asyncio.get_event_loop()
            response = loop.run_until_complete(declaration(fc))
        else:
            response = declaration(fc)
            
        return protos.Part(function_response=response)

Key Changes

1. Added async function detection using inspect.iscoroutinefunction()
2. Made CallableFunctionDeclaration.__call__ an async method
3. Added proper event loop handling for async functions
4. Ensured correct protobuf message type conversion
5. Added proper error handling for both sync and async functions

Testing

The solution was tested with async functions passed as tools to the GenerativeModel:

gemini_model = genai.GenerativeModel(
    model_name='gemini-1.5-flash',
    tools=[
        async_function1,
        async_function2,
        async_function3
    ]
)

Impact

This fix allows developers to use async functions as tools in the GenerativeModel, enabling integration with asynchronous APIs and services while maintaining proper coroutine handling.

References

1. content_types.py implementation in google.generativeai.types
2. Google GenerativeAI Python SDK documentation
3. Python asyncio documentation

Actual vs expected behavior:

Expected Behavior:

• Async functions passed as tools to GenerativeModel should be properly awaited and executed
• The model should handle both synchronous and asynchronous functions seamlessly
• Function responses should be properly converted to protobuf messages

Actual Behavior:

• Async functions passed as tools result in coroutine objects never being awaited
• Runtime errors occur with messages like:
  • RuntimeWarning: coroutine was never awaited
  • Error: Unable to coerce value: <coroutine object>
  • Parameter to MergeFrom() must be instance of same class: expected <class 'Part'> got <class 'coroutine'>

Any other information you'd like to share?

1. Environment Details:

• Python SDK Version: google-generativeai latest
• Python Version: 3.8+
• Platform: Cross-platform issue

2. Technical Details:

• The issue occurs in the content_types.py module, specifically in the CallableFunctionDeclaration and FunctionLibrary classes
• The root cause is the lack of proper coroutine handling in the function execution pipeline
• The fix maintains backward compatibility while adding async support

3. Impact:

• This fix enables developers to use async functions with the GenerativeModel's tools parameter
• Important for applications that need to integrate with asynchronous APIs or services
• Improves the overall flexibility of the SDK's function calling capabilities

4. Testing:

• The solution has been tested with both sync and async functions
• Verified proper error handling and protobuf message conversion
• Tested with multiple async functions in the tools parameter

5. Related Issues:

• This may be related to other async/await support requests in the repository
• Could improve integration with async frameworks and libraries

[MarkDaoust]

Thanks for pointing this out. Can you create a pull request with these changes?

[somwrks]

sure, which branch should i fork before editing and creating a pull request?

[MarkDaoust]

which branch

main please

[somwrks]

which branch

main please

submitted here