feat: Add support for system instruction and tools in tokenization.

PiperOrigin-RevId: 669058979

Author: happy-qiao

tests/system/vertexai/test_tokenization.py

@@ -22,9 +22,19 @@
 from vertexai.preview.tokenization import (
     get_tokenizer_for_model,
 )
-from vertexai.generative_models import GenerativeModel
+from vertexai.generative_models import (
+    GenerativeModel,
+    Part,
+    Tool,
+)
 from tests.system.aiplatform import e2e_base
 from google import auth
+from google.cloud.aiplatform_v1beta1.types import (
+    content as gapic_content_types,
+    tool as gapic_tool_types,
+    openapi,
+)
+from google.protobuf import struct_pb2
 
 
 _MODELS = ["gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash"]
@@ -39,6 +49,51 @@
     for model_name in _MODELS
     for (corpus_name, corpus_lib) in zip(_CORPUS, _CORPUS_LIB)
 ]
+_STRUCT = struct_pb2.Struct(
+    fields={
+        "string_key": struct_pb2.Value(string_value="value"),
+    }
+)
+_FUNCTION_CALL = gapic_tool_types.FunctionCall(name="test_function_call", args=_STRUCT)
+_FUNCTION_RESPONSE = gapic_tool_types.FunctionResponse(
+    name="function_response",
+    response=_STRUCT,
+)
+
+
+_SCHEMA_1 = openapi.Schema(format="schema1_format", description="schema1_description")
+_SCHEMA_2 = openapi.Schema(format="schema2_format", description="schema2_description")
+_EXAMPLE = struct_pb2.Value(string_value="value1")
+
+_FUNCTION_DECLARATION_1 = gapic_tool_types.FunctionDeclaration(
+    name="function_declaration_name",
+    description="function_declaration_description",
+    parameters=openapi.Schema(
+        format="schema_format",
+        description="schema_description",
+        enum=["schema_enum1", "schema_enum2"],
+        required=["schema_required1", "schema_required2"],
+        items=_SCHEMA_2,
+        properties={"property_key": _SCHEMA_1},
+        example=_EXAMPLE,
+    ),
+)
+_FUNCTION_DECLARATION_2 = gapic_tool_types.FunctionDeclaration(
+    parameters=openapi.Schema(
+        nullable=True,
+        default=struct_pb2.Value(string_value="value1"),
+        min_items=0,
+        max_items=0,
+        min_properties=0,
+        max_properties=0,
+        minimum=0,
+        maximum=0,
+        min_length=0,
+        max_length=0,
+        pattern="pattern",
+    ),
+    response=_SCHEMA_1,
+)
 
 STAGING_API_ENDPOINT = "STAGING_ENDPOINT"
 PROD_API_ENDPOINT = "PROD_ENDPOINT"
@@ -107,8 +162,107 @@ def test_compute_tokens(
             text = corpus_lib.raw(book)
             response = model.compute_tokens(text)
             local_result = tokenizer.compute_tokens(text)
-            for local, service in zip(
-                local_result.token_info_list, response.tokens_info
-            ):
+            for local, service in zip(local_result.tokens_info, response.tokens_info):
                 assert local.tokens == service.tokens
                 assert local.token_ids == service.token_ids
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_system_instruction(self, model_name):
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name, system_instruction=["You are a chatbot."])
+
+        assert (
+            tokenizer.count_tokens(
+                "hello", system_instruction=["You are a chatbot."]
+            ).total_tokens
+            == model.count_tokens("hello").total_tokens
+        )
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_system_instruction_is_function_call(self, model_name):
+        part = Part._from_gapic(gapic_content_types.Part(function_call=_FUNCTION_CALL))
+
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name, system_instruction=[part])
+
+        assert (
+            tokenizer.count_tokens("hello", system_instruction=[part]).total_tokens
+            == model.count_tokens("hello").total_tokens
+        )
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_system_instruction_is_function_response(self, model_name):
+        part = Part._from_gapic(
+            gapic_content_types.Part(function_response=_FUNCTION_RESPONSE)
+        )
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name, system_instruction=[part])
+
+        assert tokenizer.count_tokens(part, system_instruction=[part]).total_tokens
+        assert (
+            tokenizer.count_tokens("hello", system_instruction=[part]).total_tokens
+            == model.count_tokens("hello").total_tokens
+        )
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_tool_is_function_declaration(self, model_name):
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name)
+        tool1 = Tool._from_gapic(
+            gapic_tool_types.Tool(function_declarations=[_FUNCTION_DECLARATION_1])
+        )
+        tool2 = Tool._from_gapic(
+            gapic_tool_types.Tool(function_declarations=[_FUNCTION_DECLARATION_2])
+        )
+
+        assert tokenizer.count_tokens("hello", tools=[tool1]).total_tokens
+        with pytest.raises(ValueError):
+            tokenizer.count_tokens("hello", tools=[tool2]).total_tokens
+        assert (
+            tokenizer.count_tokens("hello", tools=[tool1]).total_tokens
+            == model.count_tokens("hello", tools=[tool1]).total_tokens
+        )
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_content_is_function_call(self, model_name):
+        part = Part._from_gapic(gapic_content_types.Part(function_call=_FUNCTION_CALL))
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name)
+
+        assert tokenizer.count_tokens(part).total_tokens
+        assert (
+            tokenizer.count_tokens(part).total_tokens
+            == model.count_tokens(part).total_tokens
+        )
+
+    @pytest.mark.parametrize(
+        "model_name",
+        _MODELS,
+    )
+    def test_count_tokens_content_is_function_response(self, model_name):
+        part = Part._from_gapic(
+            gapic_content_types.Part(function_response=_FUNCTION_RESPONSE)
+        )
+        tokenizer = get_tokenizer_for_model(model_name)
+        model = GenerativeModel(model_name)
+
+        assert tokenizer.count_tokens(part).total_tokens
+        assert (
+            tokenizer.count_tokens(part).total_tokens
+            == model.count_tokens(part).total_tokens
+        )