Skip to content

Commit

Permalink
chore: GenAI - Part.function_call.args is now a proper dict
Browse files Browse the repository at this point in the history
* Added `generative_models.FunctionCall` class (returned by `Part.FunctionCall`).

Fixes #4079

PiperOrigin-RevId: 680897095
  • Loading branch information
Ark-kun authored and copybara-github committed Oct 8, 2024
1 parent 30cf221 commit 5b3624c
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 8 deletions.
2 changes: 2 additions & 0 deletions vertexai/generative_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
ChatSession,
Content,
FinishReason,
FunctionCall,
FunctionDeclaration,
GenerationConfig,
GenerativeModel,
Expand All @@ -44,6 +45,7 @@
"ChatSession",
"Content",
"FinishReason",
"FunctionCall",
"FunctionDeclaration",
"HarmCategory",
"HarmBlockThreshold",
Expand Down
53 changes: 45 additions & 8 deletions vertexai/generative_models/_generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2302,13 +2302,13 @@ def text(self) -> str:
) from e

@property
def function_calls(self) -> Sequence[gapic_tool_types.FunctionCall]:
def function_calls(self) -> Sequence["FunctionCall"]:
if not self.content or not self.content.parts:
return []
return [
part.function_call
for part in self.content.parts
if part and part.function_call
if part._raw_part._pb.WhichOneof("data") == "function_call"
]


Expand Down Expand Up @@ -2483,8 +2483,12 @@ def file_data(self) -> gapic_content_types.FileData:
return self._raw_part.file_data

@property
def function_call(self) -> gapic_tool_types.FunctionCall:
return self._raw_part.function_call
def function_call(self) -> "FunctionCall":
return (
FunctionCall._from_gapic(self._raw_part.function_call)
if self._raw_part.function_call
else None
)

@property
def function_response(self) -> gapic_tool_types.FunctionResponse:
Expand All @@ -2495,6 +2499,40 @@ def _image(self) -> "Image":
return Image.from_bytes(data=self._raw_part.inline_data.data)


class FunctionCall:
"""Function call."""

def __init__(
self,
*,
name: str,
args: Dict[str, Any],
):
self._raw_message = aiplatform_types.FunctionCall(name=name, args=args)

@classmethod
def _from_gapic(cls, raw_message: aiplatform_types.FunctionCall) -> "FunctionCall":
response = cls.__new__(cls)
response._raw_message = raw_message
return response

def to_dict(self) -> Dict[str, Any]:
return _proto_to_dict(self._raw_message)

def __repr__(self) -> str:
return self._raw_message.__repr__()

@property
def name(self) -> str:
return self._raw_message.name

@property
def args(self) -> Dict[str, Any]:
# We cannot use `type(self.args).to_dict(self.args)`
# due to: AttributeError: type object 'MapComposite' has no attribute 'to_dict'
return self.to_dict().get("args")


class SafetySetting:
"""Parameters for the generation."""

Expand Down Expand Up @@ -2953,10 +2991,9 @@ def respond_to_model_response(
)

try:
# We cannot use `function_args = type(function_call.args).to_dict(function_call.args)`
# due to: AttributeError: type object 'MapComposite' has no attribute 'to_dict'
function_args = type(function_call).to_dict(function_call)["args"]
function_call_result = callable_function._function(**function_args)
function_call_result = callable_function._function(
**function_call.args
)
if not isinstance(function_call_result, Mapping):
# If the function returns a single value, wrap it in the
# format that Part.from_function_response can accept.
Expand Down
2 changes: 2 additions & 0 deletions vertexai/preview/generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
Candidate,
Content,
FinishReason,
FunctionCall,
FunctionDeclaration,
HarmCategory,
HarmBlockThreshold,
Expand Down Expand Up @@ -59,6 +60,7 @@ class ChatSession(_PreviewChatSession):
"ChatSession",
"Content",
"FinishReason",
"FunctionCall",
"FunctionDeclaration",
"HarmCategory",
"HarmBlockThreshold",
Expand Down

0 comments on commit 5b3624c

Please sign in to comment.