google-cloud-agentplatform 2.1.0__tar.gz → 2.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (74) hide show
  1. {google_cloud_agentplatform-2.1.0/google_cloud_agentplatform.egg-info → google_cloud_agentplatform-2.3.0}/PKG-INFO +1 -1
  2. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/__init__.py +11 -2
  3. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_common.py +23 -30
  4. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_gcs_utils.py +12 -2
  5. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/client.py +26 -0
  6. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/evals.py +14 -0
  7. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/memory_banks.py +7 -0
  8. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/runtimes.py +7 -0
  9. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/sandboxes.py +63 -13
  10. google_cloud_agentplatform-2.3.0/agentplatform/_genai/serving_profiles.py +1251 -0
  11. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/__init__.py +86 -0
  12. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/common.py +627 -27
  13. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/version.py +1 -1
  14. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0/google_cloud_agentplatform.egg-info}/PKG-INFO +1 -1
  15. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/google_cloud_agentplatform.egg-info/SOURCES.txt +1 -0
  16. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/LICENSE +0 -0
  17. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/README.md +0 -0
  18. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/__init__.py +0 -0
  19. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_bigquery_utils.py +0 -0
  20. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_datasets_utils.py +0 -0
  21. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_builtin_tools.py +0 -0
  22. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_constant.py +0 -0
  23. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_data_converters.py +0 -0
  24. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_metric_handlers.py +0 -0
  25. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_metric_loaders.py +0 -0
  26. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_utils.py +0 -0
  27. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_evals_visualization.py +0 -0
  28. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_logging_utils.py +0 -0
  29. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_memory_bank_utils.py +0 -0
  30. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_observability_data_converter.py +0 -0
  31. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_operations_utils.py +0 -0
  32. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_prompt_management_utils.py +0 -0
  33. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_prompt_optimizer_utils.py +0 -0
  34. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_runtimes_utils.py +0 -0
  35. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_skills_utils.py +0 -0
  36. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/_transformers.py +0 -0
  37. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/a2a_task_events.py +0 -0
  38. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/a2a_tasks.py +0 -0
  39. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/datasets.py +0 -0
  40. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/endpoints.py +0 -0
  41. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/example_stores.py +0 -0
  42. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/feedback_contexts.py +0 -0
  43. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/feedback_entries.py +0 -0
  44. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/live.py +0 -0
  45. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/live_runtimes.py +0 -0
  46. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/memories.py +0 -0
  47. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/memory_revisions.py +0 -0
  48. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/model_garden.py +0 -0
  49. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/prompt_optimizer.py +0 -0
  50. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/prompts.py +0 -0
  51. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/rag.py +0 -0
  52. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/runtime_revisions.py +0 -0
  53. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/sandbox_snapshots.py +0 -0
  54. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/sandbox_templates.py +0 -0
  55. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/session_events.py +0 -0
  56. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/sessions.py +0 -0
  57. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/skill_revisions.py +0 -0
  58. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/skills.py +0 -0
  59. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/agent_engines.py +0 -0
  60. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/evals.py +0 -0
  61. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/prompt_optimizer.py +0 -0
  62. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/_genai/types/prompts.py +0 -0
  63. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/__init__.py +0 -0
  64. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/a2a.py +0 -0
  65. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/adk.py +0 -0
  66. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/ag2.py +0 -0
  67. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/langchain.py +0 -0
  68. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/langgraph.py +0 -0
  69. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/agentplatform/frameworks/llama_index.py +0 -0
  70. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/google_cloud_agentplatform.egg-info/dependency_links.txt +0 -0
  71. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/google_cloud_agentplatform.egg-info/requires.txt +0 -0
  72. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/google_cloud_agentplatform.egg-info/top_level.txt +0 -0
  73. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/pyproject.toml +0 -0
  74. {google_cloud_agentplatform-2.1.0 → google_cloud_agentplatform-2.3.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-cloud-agentplatform
3
- Version: 2.1.0
3
+ Version: 2.3.0
4
4
  Summary: Google Cloud Agent Platform Client Library
5
5
  Author-email: Google LLC <googleapis-packages@google.com>
6
6
  License: Apache-2.0
@@ -54,8 +54,17 @@ def __getattr__(name): # type: ignore[no-untyped-def]
54
54
  global _genai_types
55
55
  if _genai_types is None:
56
56
  _genai_types = importlib.import_module("._genai.types", __name__)
57
- if "vertexai.types" not in sys.modules:
58
- sys.modules["vertexai.types"] = _genai_types
57
+ # `types` is an alias for `._genai.types` rather than a real submodule,
58
+ # so register it to keep
59
+ # `from agentplatform.types import TypeName`
60
+ # working without a prior attribute access. Spell it the google3 way:
61
+ # Copybara rewrites that prefix in both directions, and the external
62
+ # spelling fails its reversibility check. This key was `vertexai.types`,
63
+ # which both misnamed this package's own alias and shadowed the real
64
+ # `vertexai.types` module.
65
+ types_module_name = f"{__name__}.types"
66
+ if types_module_name not in sys.modules:
67
+ sys.modules[types_module_name] = _genai_types
59
68
  return _genai_types
60
69
 
61
70
  raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
@@ -15,7 +15,6 @@
15
15
  """Common utilities for evals."""
16
16
 
17
17
  import asyncio
18
- import base64
19
18
  import collections
20
19
  import concurrent.futures
21
20
  import contextlib
@@ -38,6 +37,11 @@ from google.genai._gaos.types.interactions import functioncallstep
38
37
  from google.genai._gaos.types.interactions import functionresultstep
39
38
  from google.genai._gaos.types.interactions import modeloutputstep
40
39
  from google.genai._gaos.types.interactions import userinputstep
40
+
41
+ try:
42
+ from google.genai._gaos.utils.serializers import ALLOW_UNKNOWN_UNION_VARIANTS
43
+ except (ImportError, AttributeError):
44
+ ALLOW_UNKNOWN_UNION_VARIANTS = "speakeasy_allow_unknown_union_variants"
41
45
  from google.genai.models import Models
42
46
  import pandas as pd
43
47
  from tqdm import tqdm
@@ -776,7 +780,16 @@ def _interaction_dict_to_agent_data(
776
780
  Returns:
777
781
  An AgentData object with one or more ConversationTurns.
778
782
  """
779
- typed_interaction = interaction_types.Interaction.model_validate(interaction)
783
+ # A server response, so unknown step types must degrade rather than raise.
784
+ # `Step` is an open discriminated union: parse_open_union only falls back to
785
+ # the Unknown variant when the validation context carries
786
+ # ALLOW_UNKNOWN_UNION_VARIANTS, and raises without it so that a
787
+ # user-constructed request payload surfaces its mistakes locally. This
788
+ # payload comes off the Interactions API, so it is on the tolerant side of
789
+ # that line. _interaction_steps_to_events already drops steps it cannot map.
790
+ typed_interaction = interaction_types.Interaction.model_validate(
791
+ interaction, context={ALLOW_UNKNOWN_UNION_VARIANTS: True}
792
+ )
780
793
  all_events = _interaction_steps_to_events(typed_interaction.steps or [])
781
794
 
782
795
  # Group events into turns. Each UserInputStep starts a new turn.
@@ -1690,8 +1703,12 @@ def _resolve_interactions_to_eval_cases(
1690
1703
  break
1691
1704
  interaction_dict = json.loads(response.body)
1692
1705
  try:
1706
+ # Also a server response -- see _interaction_dict_to_agent_data.
1707
+ # Without the context a single step type this client does not
1708
+ # know about raised, and the except below turned that into a
1709
+ # `break`, silently truncating the interaction history.
1693
1710
  typed_interaction = interaction_types.Interaction.model_validate(
1694
- interaction_dict
1711
+ interaction_dict, context={ALLOW_UNKNOWN_UNION_VARIANTS: True}
1695
1712
  )
1696
1713
  except Exception as e:
1697
1714
  logger.warning("Failed to validate interaction model: %s", e)
@@ -4078,32 +4095,6 @@ async def _convert_evaluation_run_results_async(
4078
4095
  return eval_result, eval_item_map
4079
4096
 
4080
4097
 
4081
- def _object_to_dict(obj: Any) -> Union[dict[str, Any], Any]:
4082
- """Converts an object to a dictionary."""
4083
- if obj is None:
4084
- return obj
4085
- if isinstance(obj, (int, float, str, bool)):
4086
- return obj
4087
- if isinstance(obj, datetime.datetime):
4088
- return obj.isoformat()
4089
- if isinstance(obj, bytes):
4090
- return base64.b64encode(obj).decode("utf-8")
4091
- if isinstance(obj, (list, tuple)):
4092
- return [_object_to_dict(item) for item in obj]
4093
- if isinstance(obj, dict):
4094
- return {k: _object_to_dict(v) for k, v in obj.items()}
4095
-
4096
- if not hasattr(obj, "__dict__"):
4097
- return obj # Not an object with attributes, return as is (e.g., set)
4098
-
4099
- result: dict[str, Any] = {}
4100
- for key, value in obj.__dict__.items():
4101
- if value is None:
4102
- continue
4103
- result[key] = _object_to_dict(value)
4104
- return result
4105
-
4106
-
4107
4098
  def _get_content(row: dict[str, Any], column: str) -> Optional[genai_types.Content]:
4108
4099
  if isinstance(row[column], str):
4109
4100
  return genai_types.Content(
@@ -4248,7 +4239,9 @@ def _create_evaluation_set_from_dataframe(
4248
4239
  eval_items = []
4249
4240
  for eval_item_request in eval_item_requests:
4250
4241
  gcs_uri = gcs_utils.upload_json_to_prefix(
4251
- data=_object_to_dict(eval_item_request),
4242
+ data=eval_item_request.model_dump(
4243
+ mode="json", by_alias=True, exclude_none=True
4244
+ ),
4252
4245
  gcs_dest_prefix=gcs_dest_prefix,
4253
4246
  filename_prefix="request",
4254
4247
  )
@@ -17,13 +17,15 @@ from importlib.metadata import version as get_version
17
17
  import io
18
18
  import json
19
19
  import logging
20
- from typing import Any, Union
20
+ from typing import Any, TYPE_CHECKING, Union
21
21
  import uuid
22
22
 
23
23
  from google.cloud import storage # type: ignore[attr-defined]
24
24
  from google.genai._api_client import BaseApiClient
25
25
  from packaging.version import Version
26
- import pandas as pd
26
+
27
+ if TYPE_CHECKING:
28
+ import pandas as pd
27
29
 
28
30
 
29
31
  logger = logging.getLogger(__name__)
@@ -211,6 +213,14 @@ class GcsUtils:
211
213
  self, gcs_filepath: str, file_type: str
212
214
  ) -> "pd.DataFrame":
213
215
  """Reads a file from Google Cloud Storage into a Pandas DataFrame."""
216
+ try:
217
+ import pandas as pd # pylint: disable=g-import-not-at-top
218
+ except ImportError as e:
219
+ raise ImportError(
220
+ "Reading a Cloud Storage file into a DataFrame requires 'pandas'. "
221
+ "Please install it using pip install "
222
+ "google-cloud-aiplatform[evaluation]"
223
+ ) from e
214
224
  file_contents = self.read_file_contents(gcs_filepath)
215
225
  if file_type == "csv":
216
226
  return pd.read_csv(io.StringIO(file_contents), encoding="utf-8")
@@ -62,6 +62,10 @@ if TYPE_CHECKING:
62
62
  from agentplatform._genai import (
63
63
  memory_banks as memory_banks_module,
64
64
  )
65
+ from agentplatform._genai import (
66
+ serving_profiles as serving_profiles_module,
67
+ )
68
+
65
69
 
66
70
  _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules"
67
71
 
@@ -112,6 +116,7 @@ class AsyncClient:
112
116
  self._sessions: Optional[ModuleType] = None
113
117
  self._sandboxes: Optional[ModuleType] = None
114
118
  self._memory_banks: Optional[ModuleType] = None
119
+ self._serving_profiles: Optional[ModuleType] = None
115
120
 
116
121
  @property
117
122
  @_common.experimental_warning(
@@ -287,6 +292,17 @@ class AsyncClient:
287
292
  self._memory_banks = importlib.import_module(".memory_banks", __package__)
288
293
  return self._memory_banks.AsyncMemoryBanks(self._api_client) # type: ignore[no-any-return]
289
294
 
295
+ @property
296
+ def serving_profiles(
297
+ self,
298
+ ) -> "serving_profiles_module.AsyncServingProfiles":
299
+ if self._serving_profiles is None:
300
+ self._serving_profiles = importlib.import_module(
301
+ ".serving_profiles",
302
+ __package__,
303
+ )
304
+ return self._serving_profiles.AsyncServingProfiles(self._api_client) # type: ignore[no-any-return]
305
+
290
306
  async def aclose(self) -> None:
291
307
  """Closes the async client explicitly.
292
308
 
@@ -399,6 +415,7 @@ class Client:
399
415
  self._sessions: Optional[ModuleType] = None
400
416
  self._sandboxes: Optional[ModuleType] = None
401
417
  self._memory_banks: Optional[ModuleType] = None
418
+ self._serving_profiles: Optional[ModuleType] = None
402
419
 
403
420
  @property
404
421
  def evals(self) -> "evals_module.Evals":
@@ -598,3 +615,12 @@ class Client:
598
615
  if self._memory_banks is None:
599
616
  self._memory_banks = importlib.import_module(".memory_banks", __package__)
600
617
  return self._memory_banks.MemoryBanks(self._api_client) # type: ignore[no-any-return]
618
+
619
+ @property
620
+ def serving_profiles(self) -> "serving_profiles_module.ServingProfiles":
621
+ if self._serving_profiles is None:
622
+ self._serving_profiles = importlib.import_module(
623
+ ".serving_profiles",
624
+ __package__,
625
+ )
626
+ return self._serving_profiles.ServingProfiles(self._api_client) # type: ignore[no-any-return]
@@ -215,6 +215,13 @@ def _CustomCodeExecutionSpec_from_vertex(
215
215
  getv(from_object, ["evaluation_function"]),
216
216
  )
217
217
 
218
+ if getv(from_object, ["codeExecutionRegion"]) is not None:
219
+ setv(
220
+ to_object,
221
+ ["code_execution_region"],
222
+ getv(from_object, ["codeExecutionRegion"]),
223
+ )
224
+
218
225
  return to_object
219
226
 
220
227
 
@@ -237,6 +244,13 @@ def _CustomCodeExecutionSpec_to_vertex(
237
244
  getv(from_object, ["remote_custom_function"]),
238
245
  )
239
246
 
247
+ if getv(from_object, ["code_execution_region"]) is not None:
248
+ setv(
249
+ to_object,
250
+ ["codeExecutionRegion"],
251
+ getv(from_object, ["code_execution_region"]),
252
+ )
253
+
240
254
  return to_object
241
255
 
242
256
 
@@ -504,6 +504,13 @@ def _ReasoningEngine_from_vertex(
504
504
  if getv(from_object, ["trafficConfig"]) is not None:
505
505
  setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"]))
506
506
 
507
+ if getv(from_object, ["revisionGarbageCollectionStrategy"]) is not None:
508
+ setv(
509
+ to_object,
510
+ ["revision_garbage_collection_strategy"],
511
+ getv(from_object, ["revisionGarbageCollectionStrategy"]),
512
+ )
513
+
507
514
  return to_object
508
515
 
509
516
 
@@ -484,6 +484,13 @@ def _ReasoningEngine_from_vertex(
484
484
  if getv(from_object, ["trafficConfig"]) is not None:
485
485
  setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"]))
486
486
 
487
+ if getv(from_object, ["revisionGarbageCollectionStrategy"]) is not None:
488
+ setv(
489
+ to_object,
490
+ ["revision_garbage_collection_strategy"],
491
+ getv(from_object, ["revisionGarbageCollectionStrategy"]),
492
+ )
493
+
487
494
  return to_object
488
495
 
489
496
 
@@ -26,7 +26,6 @@ from typing import Any, Iterator, Optional, Union
26
26
  from urllib.parse import urlencode
27
27
 
28
28
  from google import auth as google_auth
29
- from google import genai
30
29
  from google.auth.transport import requests as google_auth_requests
31
30
  from google.genai import _api_module
32
31
  from google.genai import _common
@@ -34,6 +33,7 @@ from google.genai import types as genai_types
34
33
  from google.genai._common import get_value_by_path as getv
35
34
  from google.genai._common import set_value_by_path as setv
36
35
  from google.genai.pagers import Pager
36
+ import requests
37
37
 
38
38
  from . import _runtimes_utils
39
39
  from . import types
@@ -1362,6 +1362,7 @@ class Sandboxes(_api_module.BaseModule):
1362
1362
  query_params: Optional[dict[str, object]] = None,
1363
1363
  headers: Optional[dict[str, str]] = None,
1364
1364
  request_dict: Optional[dict[str, object]] = None,
1365
+ psc_endpoint: Optional[str] = None,
1365
1366
  ) -> genai_types.HttpResponse:
1366
1367
  """Sends a command to the sandbox.
1367
1368
 
@@ -1382,6 +1383,16 @@ class Sandboxes(_api_module.BaseModule):
1382
1383
  Optional. The headers to include in the command.
1383
1384
  request_dict (dict[str, object]):
1384
1385
  Optional. The request body to include in the command.
1386
+ psc_endpoint (str):
1387
+ Optional. Host (RFC1918 IP or resolvable internal hostname) of the
1388
+ customer-side PSC endpoint forwarding rule that targets
1389
+ ``connection_info.service_attachment``. Needed for VPC-SC sandboxes
1390
+ (``ingressControlConfig.enablePrivateServiceConnect=true``) where
1391
+ ``connection_info.load_balancer_hostname`` and ``load_balancer_ip``
1392
+ are both unset because ingress goes through PSC in the customer's
1393
+ VPC rather than a Google-managed load balancer. When provided, it
1394
+ is used as the data-plane host in place of the load balancer
1395
+ address.
1385
1396
 
1386
1397
  Returns:
1387
1398
  genai_types.HttpResponse: The response from the sandbox.
@@ -1391,12 +1402,21 @@ class Sandboxes(_api_module.BaseModule):
1391
1402
  connection_info = sandbox_environment.connection_info
1392
1403
  if not connection_info:
1393
1404
  raise ValueError("Connection info is not available.")
1394
- if connection_info.load_balancer_hostname:
1405
+ if psc_endpoint:
1406
+ # VPC-SC path: caller has provisioned a PSC endpoint against
1407
+ # connection_info.service_attachment and passes its host here.
1408
+ endpoint = "https://" + psc_endpoint
1409
+ elif connection_info.load_balancer_hostname:
1395
1410
  endpoint = "https://" + connection_info.load_balancer_hostname
1396
1411
  elif connection_info.load_balancer_ip:
1397
1412
  endpoint = "http://" + connection_info.load_balancer_ip
1398
1413
  else:
1399
- raise ValueError("Load balancer hostname or ip is not available.")
1414
+ raise ValueError(
1415
+ "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
1416
+ " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
1417
+ " the caller to pass psc_endpoint (the host of the forwarding"
1418
+ " rule targeting connection_info.service_attachment)."
1419
+ )
1400
1420
 
1401
1421
  routing_token = connection_info.routing_token
1402
1422
  if not routing_token:
@@ -1408,15 +1428,27 @@ class Sandboxes(_api_module.BaseModule):
1408
1428
  headers["Authorization"] = f"Bearer {access_token}"
1409
1429
  headers["X-Sandbox-Routing-Token"] = routing_token
1410
1430
  headers["X-Sandbox-Port"] = port
1411
- endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path
1412
- http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
1413
- http_client = genai.Client(vertexai=True, http_options=http_options)
1414
- # Full path is constructed in this function. The passed in path into request
1415
- # function will not be used.
1416
- response = http_client._api_client.request(http_method, path, request_dict)
1431
+ url = endpoint + path if path.startswith("/") else endpoint + "/" + path
1432
+ # The sandbox data plane authenticates via the Authorization,
1433
+ # X-Sandbox-Routing-Token and X-Sandbox-Port headers set above, not via
1434
+ # ADC. Routing this through genai.Client(vertexai=True) builds an
1435
+ # authorized session with no credentials, whose request() then fails with
1436
+ # "'NoneType' object has no attribute 'before_request'". Issue the request
1437
+ # directly instead, mirroring generate_access_token above which also talks
1438
+ # to a non-aiplatform endpoint with plain requests.
1439
+ body = None
1440
+ if request_dict:
1441
+ body = json.dumps(request_dict)
1442
+ headers.setdefault("Content-Type", "application/json")
1443
+ response = requests.request(
1444
+ http_method,
1445
+ url,
1446
+ headers=headers,
1447
+ data=body,
1448
+ )
1417
1449
  return genai_types.HttpResponse(
1418
- headers=response.headers,
1419
- body=response.body,
1450
+ headers=dict(response.headers),
1451
+ body=response.text,
1420
1452
  )
1421
1453
 
1422
1454
  def generate_browser_ws_headers(
@@ -1425,6 +1457,7 @@ class Sandboxes(_api_module.BaseModule):
1425
1457
  service_account_email: str,
1426
1458
  port: str = "8080",
1427
1459
  timeout: int = 3600,
1460
+ psc_endpoint: Optional[str] = None,
1428
1461
  ) -> tuple[str, dict[str, str]]:
1429
1462
  """Generates the websocket upgrade headers for the browser.
1430
1463
 
@@ -1438,6 +1471,15 @@ class Sandboxes(_api_module.BaseModule):
1438
1471
  Defaults to "8080". This should be one of the ports specified during template creation.
1439
1472
  timeout (int):
1440
1473
  Optional. The timeout in seconds for the token. Defaults to 3600.
1474
+ psc_endpoint (str):
1475
+ Optional. Host (RFC1918 IP or resolvable internal hostname) of the
1476
+ customer-side PSC endpoint forwarding rule that targets
1477
+ ``connection_info.service_attachment``. Needed for VPC-SC
1478
+ sandboxes; see ``send_command`` for details. When provided, it is
1479
+ used as the websocket host in place of the load balancer address,
1480
+ and is propagated to the internal ``send_command`` call that
1481
+ fetches the CDP endpoint.
1482
+
1441
1483
  Returns:
1442
1484
  tuple[str, dict[str, str]]: A tuple containing the websocket URL and
1443
1485
  the headers for websocket upgrade.
@@ -1446,12 +1488,19 @@ class Sandboxes(_api_module.BaseModule):
1446
1488
  raise ValueError("Connection info is not available.")
1447
1489
 
1448
1490
  connection_info = sandbox_environment.connection_info
1449
- if connection_info.load_balancer_hostname:
1491
+ if psc_endpoint:
1492
+ ws_base_url = "wss://" + psc_endpoint
1493
+ elif connection_info.load_balancer_hostname:
1450
1494
  ws_base_url = "wss://" + connection_info.load_balancer_hostname
1451
1495
  elif connection_info.load_balancer_ip:
1452
1496
  ws_base_url = "ws://" + connection_info.load_balancer_ip
1453
1497
  else:
1454
- raise ValueError("Load balancer hostname or ip is not available.")
1498
+ raise ValueError(
1499
+ "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
1500
+ " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
1501
+ " the caller to pass psc_endpoint (the host of the forwarding"
1502
+ " rule targeting connection_info.service_attachment)."
1503
+ )
1455
1504
 
1456
1505
  http_access_token = self.generate_access_token(service_account_email, timeout)
1457
1506
  response = self.send_command(
@@ -1460,6 +1509,7 @@ class Sandboxes(_api_module.BaseModule):
1460
1509
  sandbox_environment=sandbox_environment,
1461
1510
  port=port,
1462
1511
  path="/cdp_ws_endpoint",
1512
+ psc_endpoint=psc_endpoint,
1463
1513
  )
1464
1514
  if not response:
1465
1515
  raise ValueError("Failed to get the websocket endpoint.")