Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
20 changes: 20 additions & 0 deletions astrbot/core/agent/tool.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import copy
import json
from collections.abc import AsyncGenerator, Awaitable, Callable
from typing import Any, Generic

Expand Down Expand Up @@ -303,6 +304,25 @@ def convert_schema(schema: dict) -> dict:
}
result.update({k: schema[k] for k in support_fields if k in schema})

# Gemini only accepts string enums ("enum: only allowed for STRING
# type"). Other fields keep their type, so the model still sends
# e.g. an integer, and list the allowed values in the description.
if "enum" in result:
enum = result.pop("enum")
values = [v for v in enum if v is not None]
if None in enum:
result["nullable"] = True
if result["type"] == "string":
result["enum"] = [
v if isinstance(v, str) else json.dumps(v) for v in values
]
else:
allowed = ", ".join(json.dumps(v) for v in values)
description = result.get("description", "")
result["description"] = (
f"{description}\nAllowed values: {allowed}".strip()
)
Comment thread
sourcery-ai[bot] marked this conversation as resolved.

if "properties" in schema:
properties = {}
for key, value in schema["properties"].items():
Expand Down
45 changes: 45 additions & 0 deletions tests/unit/test_tool_google_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
from pathlib import Path
from typing import Generic, TypeVar

from google.genai import types as genai_types

REPO_ROOT = Path(__file__).resolve().parents[2]
TOOL_MODULE_PATH = REPO_ROOT / "astrbot/core/agent/tool.py"

Expand Down Expand Up @@ -77,6 +79,49 @@ def test_google_schema_fills_missing_array_items_with_string_schema():
assert source_uuids["items"] == {"type": "string"}


def test_google_schema_keeps_enums_only_on_string_fields():
tool_module = load_tool_module()
FunctionTool = tool_module.FunctionTool
ToolSet = tool_module.ToolSet

tool = FunctionTool(
name="list_items",
description="List items.",
parameters={
"type": "object",
"properties": {
"page_size": {
"type": "integer",
"description": "Items per page.",
"enum": [10, 20, 50],
},
"strict": {"type": "boolean", "enum": [True]},
"order": {"type": ["string", "null"], "enum": ["asc", "desc", None]},
},
},
)

schema = ToolSet([tool]).google_schema()
properties = schema["function_declarations"][0]["parameters"]["properties"]

assert properties["page_size"] == {
"type": "integer",
"description": "Items per page.\nAllowed values: 10, 20, 50",
}
assert properties["strict"] == {
"type": "boolean",
"description": "Allowed values: true",
}
assert properties["order"] == {
"type": "string",
"enum": ["asc", "desc"],
"nullable": True,
}
assert tool.parameters["properties"]["page_size"]["enum"] == [10, 20, 50]
# The Gemini SDK rejects non-string enum values when the request is built.
genai_types.Tool(function_declarations=schema["function_declarations"])


def test_openai_schema_sorts_tools_by_name_without_mutating_toolset_order():
tool_module = load_tool_module()
FunctionTool = tool_module.FunctionTool
Expand Down
Loading