Skip to content

feat(google): add thought_signature support for Gemini 2.0+ #112

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Gabbosaur wants to merge 1 commit into
base: main
Choose a base branch
from
+61 −42
Open

feat(google): add thought_signature support for Gemini 2.0+ #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 49 additions & 41 deletions datapizza-ai-clients/datapizza-ai-clients-google/datapizza/clients/google/google_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -617,48 +617,56 @@ def _response_to_client_response(
self, response, tool_map: dict[str, Tool] | None = None
) -> ClientResponse:
blocks = []
# Handle function calls if present
if hasattr(response, "function_calls") and response.function_calls:
for fc in response.function_calls:
if not tool_map:
raise ValueError("Tool map is required")

tool = tool_map.get(fc.name, None)
if not tool:
raise ValueError(f"Tool {fc.name} not found in tool map")

blocks.append(
FunctionCallBlock(
name=fc.name,
arguments=fc.args,
id=f"fc_{id(fc)}",
tool=tool,
)
)
else:
if (
hasattr(response, "candidates")
and response.candidates
and response.candidates[0].content.parts
):
for part in response.candidates[0].content.parts:
# Handle inline_data (images from generation or code execution)
if hasattr(part, "inline_data") and part.inline_data is not None:
media = Media(
media_type="image",
source_type="base64",
source=base64.b64encode(part.inline_data.data).decode(
"utf-8"
),
extension=(part.inline_data.mime_type.split("/")[-1])
if part.inline_data.mime_type
else "png",

# Check if response has candidates with parts
if (
hasattr(response, "candidates")
and response.candidates
and response.candidates[0].content
and response.candidates[0].content.parts
):
for part in response.candidates[0].content.parts:
# Handle function calls - extract thought_signature from part
if hasattr(part, "function_call") and part.function_call:
fc = part.function_call
if not tool_map:
raise ValueError("Tool map is required")

tool = tool_map.get(fc.name, None)
if not tool:
raise ValueError(f"Tool {fc.name} not found in tool map")

# Extract thought_signature if present (required for Gemini 2.0+)
thought_sig = None
if hasattr(part, "thought_signature"):
thought_sig = part.thought_signature

blocks.append(
FunctionCallBlock(
name=fc.name,
arguments=dict(fc.args) if fc.args else {},
id=f"fc_{id(fc)}",
tool=tool,
thought_signature=thought_sig,
)
blocks.append(MediaBlock(media=media))
elif hasattr(part, "thought") and part.thought and part.text:
blocks.append(ThoughtBlock(content=part.text))
elif hasattr(part, "text") and part.text:
blocks.append(TextBlock(content=part.text))
)
# Handle inline_data (images from generation or code execution)
elif hasattr(part, "inline_data") and part.inline_data is not None:
media = Media(
media_type="image",
source_type="base64",
source=base64.b64encode(part.inline_data.data).decode(
"utf-8"
),
extension=(part.inline_data.mime_type.split("/")[-1])
if part.inline_data.mime_type
else "png",
)
blocks.append(MediaBlock(media=media))
elif hasattr(part, "thought") and part.thought and part.text:
blocks.append(ThoughtBlock(content=part.text))
elif hasattr(part, "text") and part.text:
blocks.append(TextBlock(content=part.text))

usage_metadata = getattr(response, "usage_metadata", None)
token_usage = self._token_usage_from_metadata(usage_metadata)
Expand Down
3 changes: 3 additions & 0 deletions datapizza-ai-clients/datapizza-ai-clients-google/datapizza/clients/google/memory_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ def _turn_to_message(self, turn: Turn) -> dict:
block_dict = {
"function_call": {"name": block.name, "args": block.arguments}
}
# Include thought_signature if present (required for Gemini 2.0+)
if block.thought_signature:
block_dict["thought_signature"] = block.thought_signature
case FunctionCallResultBlock():
block_dict = types.Part.from_function_response(
name=block.tool.name,
Expand Down
10 changes: 9 additions & 1 deletion datapizza-ai-core/datapizza/type/type.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ def from_dict(cls, data: dict):
arguments=data.get("arguments", {}),
name=data.get("name", ""),
tool=tool,
thought_signature=data.get("thought_signature"),
)
case "function_call_result":
tool = Tool.tool_from_dict(data.get("tool"))
Expand Down Expand Up @@ -170,6 +171,7 @@ def __init__(
name: str,
tool: Tool,
type: str = "function",
thought_signature: str | None = None,
):
"""
Initialize a FunctionCallBlock object.
Expand All @@ -179,11 +181,14 @@ def __init__(
arguments (dict[str, Any]): The arguments of the function call block.
name (str): The name of the function call block.
tool (Tool): The tool of the function call block.
thought_signature (str, optional): The thought signature for Gemini 2.0+ models.
Required for multi-turn function calling with thinking models.
"""
self.id = id
self.arguments = arguments
self.name = name
self.tool = tool
self.thought_signature = thought_signature
super().__init__(type)

def __eq__(self, other):
Expand All @@ -204,13 +209,16 @@ def __hash__(self) -> int:
return int(hashlib.sha256(self.id.encode("utf-8")).hexdigest(), 16)

def to_dict(self) -> dict:
return {
result = {
"type": self.type,
"id": self.id,
"arguments": self.arguments,
"name": self.name,
"tool": self.tool.to_dict(),
}
if self.thought_signature:
result["thought_signature"] = self.thought_signature
return result


class FunctionCallResultBlock(Block):
Expand Down
Loading