chore: Add support for abstract types in AFC

PiperOrigin-RevId: 831873580
This commit is contained in:
Sara Robinson
2025-11-13 08:38:25 -08:00
committed by Copybara-Service
parent 6bb0b7417e
commit 2efc184a46
3 changed files with 218 additions and 20 deletions
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from collections.abc import Sequence
from typing import Any
from typing import Dict
@@ -192,3 +193,52 @@ def test_from_function_with_options_no_params():
# VERTEX_AI should have response schema for None return
assert declaration.response is not None
assert declaration.response.type == types.Type.NULL
def test_from_function_with_collections_type_parameter():
"""Test from_function_with_options with collections type parameter."""
def test_function(
artifact_key: str,
input_edit_ids: Sequence[str],
) -> str:
"""Saves a sequence of edit IDs."""
return f'Saved {len(input_edit_ids)} edit IDs for artifact {artifact_key}'
declaration = _automatic_function_calling_util.from_function_with_options(
test_function, GoogleLLMVariant.VERTEX_AI
)
assert declaration.name == 'test_function'
assert declaration.parameters.type == types.Type.OBJECT
assert (
declaration.parameters.properties['artifact_key'].type
== types.Type.STRING
)
assert (
declaration.parameters.properties['input_edit_ids'].type
== types.Type.ARRAY
)
assert (
declaration.parameters.properties['input_edit_ids'].items.type
== types.Type.STRING
)
assert declaration.response.type == types.Type.STRING
def test_from_function_with_collections_return_type():
"""Test from_function_with_options with collections return type."""
def test_function(
names: list[str],
) -> Sequence[str]:
"""Returns a sequence of names."""
return names
declaration = _automatic_function_calling_util.from_function_with_options(
test_function, GoogleLLMVariant.VERTEX_AI
)
assert declaration.name == 'test_function'
assert declaration.response.type == types.Type.ARRAY
assert declaration.response.items.type == types.Type.STRING