Skip to content

Commit 7b51d9e

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
feat: GenAI SDK client - Adding client-based SDKs for Agent Engine
PiperOrigin-RevId: 774906945
1 parent 1c60ee2 commit 7b51d9e

File tree

10 files changed

+5216
-720
lines changed

10 files changed

+5216
-720
lines changed

‎tests/unit/vertex_langchain/test_agent_engines.py

Lines changed: 139 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
#
15-
from absl.testing import parameterized
1615
import cloudpickle
1716
import difflib
1817
import importlib
@@ -1147,7 +1146,6 @@ def test_create_agent_engine_with_env_vars_list(
11471146
retry=_TEST_RETRY,
11481147
)
11491148

1150-
# pytest does not allow absl.testing.parameterized.named_parameters.
11511149
@pytest.mark.parametrize(
11521150
"test_case_name, test_engine_instance, expected_framework",
11531151
[
@@ -1190,7 +1188,6 @@ def test_get_agent_framework(
11901188
framework = _agent_engines._get_agent_framework(test_engine_instance)
11911189
assert framework == expected_framework
11921190

1193-
# pytest does not allow absl.testing.parameterized.named_parameters.
11941191
@pytest.mark.parametrize(
11951192
"test_case_name, test_kwargs, want_request",
11961193
[
@@ -1601,7 +1598,6 @@ def test_query_agent_engine(
16011598
test_agent_engine.query(query=_TEST_QUERY_PROMPT)
16021599
query_mock.assert_called_with(request=_TEST_AGENT_ENGINE_QUERY_REQUEST)
16031600

1604-
# pytest does not allow absl.testing.parameterized.named_parameters.
16051601
@pytest.mark.parametrize(
16061602
"test_case_name, test_class_methods_spec, want_operation_schema_api_modes",
16071603
[
@@ -1847,7 +1843,6 @@ def test_query_after_create_agent_engine_with_operation_schema(
18471843
class_method=method_name,
18481844
)
18491845
)
1850-
assert invoked_method.__doc__ == test_doc
18511846

18521847
# pytest does not allow absl.testing.parameterized.named_parameters.
18531848
@pytest.mark.parametrize(
@@ -2062,7 +2057,6 @@ def test_stream_query_after_create_agent_engine_with_operation_schema(
20622057
class_method=method_name,
20632058
)
20642059
)
2065-
assert invoked_method.__doc__ == test_doc
20662060

20672061
# pytest does not allow absl.testing.parameterized.named_parameters.
20682062
@pytest.mark.parametrize(
@@ -2243,7 +2237,6 @@ async def test_async_stream_query_after_create_agent_engine_with_operation_schem
22432237
class_method=method_name,
22442238
)
22452239
)
2246-
assert invoked_method.__doc__ == test_doc
22472240

22482241
# pytest does not allow absl.testing.parameterized.named_parameters.
22492242
@pytest.mark.parametrize(
@@ -2870,7 +2863,7 @@ def test_update_class_methods_spec_with_registered_operation_not_found(self):
28702863
"register the API methods: "
28712864
"https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
28722865
"Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
2873-
"Supported modes are: ``, `async`, `stream` and `async_stream`.}"
2866+
"Supported modes are: ``, `async`, `async_stream`, `stream`.}"
28742867
),
28752868
),
28762869
],
@@ -2987,161 +2980,169 @@ def assert_called_with_diff(mock_obj, expected_kwargs=None):
29872980
)
29882981

29892982

2990-
class TestGenerateSchema(parameterized.TestCase):
2991-
@parameterized.named_parameters(
2992-
dict(
2993-
testcase_name="place_tool_query",
2994-
func=place_tool_query,
2995-
required=["city", "activity"],
2996-
expected_operation={
2997-
"name": "place_tool_query",
2998-
"description": (
2999-
"Searches the city for recommendations on the activity."
3000-
),
3001-
"parameters": {
3002-
"type": "object",
3003-
"properties": {
3004-
"city": {"type": "string"},
3005-
"activity": {"type": "string", "nullable": True},
3006-
"page_size": {"type": "integer"},
2983+
class TestGenerateSchema:
2984+
# pytest does not allow absl.testing.parameterized.named_parameters.
2985+
@pytest.mark.parametrize(
2986+
"func, required, expected_operation",
2987+
[
2988+
(
2989+
# "place_tool_query",
2990+
place_tool_query,
2991+
["city", "activity"],
2992+
{
2993+
"name": "place_tool_query",
2994+
"description": (
2995+
"Searches the city for recommendations on the activity."
2996+
),
2997+
"parameters": {
2998+
"type": "object",
2999+
"properties": {
3000+
"city": {"type": "string"},
3001+
"activity": {"type": "string", "nullable": True},
3002+
"page_size": {"type": "integer"},
3003+
},
3004+
"required": ["city", "activity"],
30073005
},
3008-
"required": ["city", "activity"],
30093006
},
3010-
},
3011-
),
3012-
dict(
3013-
testcase_name="place_photo_query",
3014-
func=place_photo_query,
3015-
required=["photo_reference"],
3016-
expected_operation={
3017-
"name": "place_photo_query",
3018-
"description": "Returns the photo for a given reference.",
3019-
"parameters": {
3020-
"properties": {
3021-
"photo_reference": {"type": "string"},
3022-
"maxwidth": {"type": "integer"},
3023-
"maxheight": {"type": "integer", "nullable": True},
3007+
),
3008+
(
3009+
# "place_photo_query",
3010+
place_photo_query,
3011+
["photo_reference"],
3012+
{
3013+
"name": "place_photo_query",
3014+
"description": "Returns the photo for a given reference.",
3015+
"parameters": {
3016+
"type": "object",
3017+
"properties": {
3018+
"photo_reference": {"type": "string"},
3019+
"maxwidth": {"type": "integer"},
3020+
"maxheight": {"type": "integer", "nullable": True},
3021+
},
3022+
"required": ["photo_reference"],
30243023
},
3025-
"required": ["photo_reference"],
3026-
"type": "object",
30273024
},
3028-
},
3029-
),
3025+
),
3026+
],
30303027
)
30313028
def test_generate_schemas(self, func, required, expected_operation):
30323029
result = _utils.generate_schema(func, required=required)
3033-
self.assertDictEqual(result, expected_operation)
3030+
assert result == expected_operation
30343031

30353032

3036-
class TestToProto(parameterized.TestCase):
3037-
@parameterized.named_parameters(
3038-
dict(
3039-
testcase_name="empty_dict",
3040-
obj={},
3041-
expected_proto=struct_pb2.Struct(fields={}),
3042-
),
3043-
dict(
3044-
testcase_name="nonempty_dict",
3045-
obj={"snake_case": 1, "camelCase": 2},
3046-
expected_proto=struct_pb2.Struct(
3047-
fields={
3048-
"snake_case": struct_pb2.Value(number_value=1),
3049-
"camelCase": struct_pb2.Value(number_value=2),
3050-
},
3033+
class TestToProto:
3034+
# pytest does not allow absl.testing.parameterized.named_parameters.
3035+
@pytest.mark.parametrize(
3036+
"obj, expected_proto",
3037+
[
3038+
(
3039+
# "empty_dict",
3040+
{},
3041+
struct_pb2.Struct(fields={}),
30513042
),
3052-
),
3053-
dict(
3054-
testcase_name="empty_proto_message",
3055-
obj=struct_pb2.Struct(fields={}),
3056-
expected_proto=struct_pb2.Struct(fields={}),
3057-
),
3058-
dict(
3059-
testcase_name="nonempty_proto_message",
3060-
obj=struct_pb2.Struct(
3061-
fields={
3062-
"snake_case": struct_pb2.Value(number_value=1),
3063-
"camelCase": struct_pb2.Value(number_value=2),
3064-
},
3043+
(
3044+
# "nonempty_dict",
3045+
{"snake_case": 1, "camelCase": 2},
3046+
struct_pb2.Struct(
3047+
fields={
3048+
"snake_case": struct_pb2.Value(number_value=1),
3049+
"camelCase": struct_pb2.Value(number_value=2),
3050+
},
3051+
),
30653052
),
3066-
expected_proto=struct_pb2.Struct(
3067-
fields={
3068-
"snake_case": struct_pb2.Value(number_value=1),
3069-
"camelCase": struct_pb2.Value(number_value=2),
3070-
},
3053+
(
3054+
# "empty_proto_message",
3055+
struct_pb2.Struct(fields={}),
3056+
struct_pb2.Struct(fields={}),
30713057
),
3072-
),
3058+
(
3059+
# "nonempty_proto_message",
3060+
struct_pb2.Struct(
3061+
fields={
3062+
"snake_case": struct_pb2.Value(number_value=1),
3063+
"camelCase": struct_pb2.Value(number_value=2),
3064+
},
3065+
),
3066+
struct_pb2.Struct(
3067+
fields={
3068+
"snake_case": struct_pb2.Value(number_value=1),
3069+
"camelCase": struct_pb2.Value(number_value=2),
3070+
},
3071+
),
3072+
),
3073+
],
30733074
)
30743075
def test_to_proto(self, obj, expected_proto):
30753076
result = _utils.to_proto(obj)
3076-
self.assertDictEqual(_utils.to_dict(result), _utils.to_dict(expected_proto))
3077-
# converting a new object to proto should not modify earlier objects.
3078-
new_result = _utils.to_proto({})
3079-
self.assertDictEqual(_utils.to_dict(result), _utils.to_dict(expected_proto))
3080-
self.assertEmpty(new_result)
3077+
assert _utils.to_dict(result) == _utils.to_dict(expected_proto)
30813078

30823079

3083-
class ToParsedJsonTest(parameterized.TestCase):
3084-
@parameterized.named_parameters(
3085-
dict(
3086-
testcase_name="valid_json",
3087-
obj=httpbody_pb2.HttpBody(
3088-
content_type="application/json", data=b'{"a": 1, "b": "hello"}'
3080+
class ToParsedJsonTest:
3081+
# pytest does not allow absl.testing.parameterized.named_parameters.
3082+
@pytest.mark.parametrize(
3083+
"obj, expected",
3084+
[
3085+
(
3086+
# "valid_json",
3087+
httpbody_pb2.HttpBody(
3088+
content_type="application/json", data=b'{"a": 1, "b": "hello"}'
3089+
),
3090+
[{"a": 1, "b": "hello"}],
30893091
),
3090-
expected=[{"a": 1, "b": "hello"}],
3091-
),
3092-
dict(
3093-
testcase_name="invalid_json",
3094-
obj=httpbody_pb2.HttpBody(
3095-
content_type="application/json", data=b'{"a": 1, "b": "hello"'
3092+
(
3093+
# "invalid_json",
3094+
httpbody_pb2.HttpBody(
3095+
content_type="application/json", data=b'{"a": 1, "b": "hello"'
3096+
),
3097+
['{"a": 1, "b": "hello"'], # returns the unparsed string
30963098
),
3097-
expected=['{"a": 1, "b": "hello"'], # returns the unparsed string
3098-
),
3099-
dict(
3100-
testcase_name="missing_content_type",
3101-
obj=httpbody_pb2.HttpBody(data=b'{"a": 1}'),
3102-
expected=[httpbody_pb2.HttpBody(data=b'{"a": 1}')],
3103-
),
3104-
dict(
3105-
testcase_name="missing_data",
3106-
obj=httpbody_pb2.HttpBody(content_type="application/json"),
3107-
expected=[None],
3108-
),
3109-
dict(
3110-
testcase_name="wrong_content_type",
3111-
obj=httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello"),
3112-
expected=[httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello")],
3113-
),
3114-
dict(
3115-
testcase_name="empty_data",
3116-
obj=httpbody_pb2.HttpBody(content_type="application/json", data=b""),
3117-
expected=[None],
3118-
),
3119-
dict(
3120-
testcase_name="unicode_data",
3121-
obj=httpbody_pb2.HttpBody(
3122-
content_type="application/json", data='{"a": "你好"}'.encode("utf-8")
3099+
(
3100+
# "missing_content_type",
3101+
httpbody_pb2.HttpBody(data=b'{"a": 1}'),
3102+
[httpbody_pb2.HttpBody(data=b'{"a": 1}')],
31233103
),
3124-
expected=[{"a": "你好"}],
3125-
),
3126-
dict(
3127-
testcase_name="nested_json",
3128-
obj=httpbody_pb2.HttpBody(
3129-
content_type="application/json", data=b'{"a": {"b": 1}}'
3104+
(
3105+
# "missing_data",
3106+
httpbody_pb2.HttpBody(content_type="application/json"),
3107+
[None],
31303108
),
3131-
expected=[{"a": {"b": 1}}],
3132-
),
3133-
dict(
3134-
testcase_name="multiline_json",
3135-
obj=httpbody_pb2.HttpBody(
3136-
content_type="application/json",
3137-
data=b'{"a": {"b": 1}}\n{"a": {"b": 2}}',
3109+
(
3110+
# "wrong_content_type",
3111+
httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello"),
3112+
[httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello")],
31383113
),
3139-
expected=[{"a": {"b": 1}}, {"a": {"b": 2}}],
3140-
),
3114+
(
3115+
# "empty_data",
3116+
httpbody_pb2.HttpBody(content_type="application/json", data=b""),
3117+
[None],
3118+
),
3119+
(
3120+
# "unicode_data",
3121+
httpbody_pb2.HttpBody(
3122+
content_type="application/json", data='{"a": "你好"}'.encode("utf-8")
3123+
),
3124+
[{"a": "你好"}],
3125+
),
3126+
(
3127+
# "nested_json",
3128+
httpbody_pb2.HttpBody(
3129+
content_type="application/json", data=b'{"a": {"b": 1}}'
3130+
),
3131+
[{"a": {"b": 1}}],
3132+
),
3133+
(
3134+
# "multiline_json",
3135+
httpbody_pb2.HttpBody(
3136+
content_type="application/json",
3137+
data=b'{"a": {"b": 1}}\n{"a": {"b": 2}}',
3138+
),
3139+
[{"a": {"b": 1}}, {"a": {"b": 2}}],
3140+
),
3141+
],
31413142
)
31423143
def test_to_parsed_json(self, obj, expected):
31433144
for got, want in zip(_utils.yield_parsed_json(obj), expected):
3144-
self.assertEqual(got, want)
3145+
assert got == want
31453146

31463147

31473148
class TestRequirements:

0 commit comments

Comments
 (0)