google-cloud-agentplatform 2.5.0__tar.gz → 2.6.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.5.0/google_cloud_agentplatform.egg-info → google_cloud_agentplatform-2.6.0}/PKG-INFO +1 -1
  2. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_common.py +43 -7
  3. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_data_converters.py +17 -0
  4. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_visualization.py +2 -2
  5. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/memory_banks.py +268 -0
  6. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/model_garden.py +1027 -61
  7. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/__init__.py +72 -0
  8. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/common.py +785 -2
  9. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/evals.py +1 -1
  10. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/adk.py +218 -85
  11. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/version.py +1 -1
  12. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0/google_cloud_agentplatform.egg-info}/PKG-INFO +1 -1
  13. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/LICENSE +0 -0
  14. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/README.md +0 -0
  15. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/__init__.py +0 -0
  16. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/__init__.py +0 -0
  17. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_bigquery_utils.py +0 -0
  18. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_datasets_utils.py +0 -0
  19. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_builtin_tools.py +0 -0
  20. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_constant.py +0 -0
  21. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_metric_handlers.py +0 -0
  22. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_metric_loaders.py +0 -0
  23. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_evals_utils.py +0 -0
  24. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_gcs_utils.py +0 -0
  25. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_logging_utils.py +0 -0
  26. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_memory_bank_utils.py +0 -0
  27. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_observability_data_converter.py +0 -0
  28. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_operations_utils.py +0 -0
  29. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_prompt_management_utils.py +0 -0
  30. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_prompt_optimizer_utils.py +0 -0
  31. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_runtimes_utils.py +0 -0
  32. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_skills_utils.py +0 -0
  33. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/_transformers.py +0 -0
  34. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/a2a_task_events.py +0 -0
  35. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/a2a_tasks.py +0 -0
  36. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/client.py +0 -0
  37. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/datasets.py +0 -0
  38. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/endpoints.py +0 -0
  39. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/evals.py +0 -0
  40. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/example_stores.py +0 -0
  41. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/feedback_contexts.py +0 -0
  42. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/feedback_entries.py +0 -0
  43. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/live.py +0 -0
  44. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/live_runtimes.py +0 -0
  45. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/memories.py +0 -0
  46. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/memory_revisions.py +0 -0
  47. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/prompt_optimizer.py +0 -0
  48. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/prompts.py +0 -0
  49. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/rag.py +0 -0
  50. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/runtime_revisions.py +0 -0
  51. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/runtimes.py +0 -0
  52. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/sandbox_snapshots.py +0 -0
  53. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/sandbox_templates.py +0 -0
  54. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/sandboxes.py +0 -0
  55. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/serving_profiles.py +0 -0
  56. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/session_events.py +0 -0
  57. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/sessions.py +0 -0
  58. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/skill_revisions.py +0 -0
  59. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/skills.py +0 -0
  60. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/agent_engines.py +0 -0
  61. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/prompt_optimizer.py +0 -0
  62. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/_genai/types/prompts.py +0 -0
  63. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/__init__.py +0 -0
  64. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/a2a.py +0 -0
  65. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/ag2.py +0 -0
  66. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/langchain.py +0 -0
  67. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/langgraph.py +0 -0
  68. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/agentplatform/frameworks/llama_index.py +0 -0
  69. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/google_cloud_agentplatform.egg-info/SOURCES.txt +0 -0
  70. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/google_cloud_agentplatform.egg-info/dependency_links.txt +0 -0
  71. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/google_cloud_agentplatform.egg-info/requires.txt +0 -0
  72. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/google_cloud_agentplatform.egg-info/top_level.txt +0 -0
  73. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/pyproject.toml +0 -0
  74. {google_cloud_agentplatform-2.5.0 → google_cloud_agentplatform-2.6.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-cloud-agentplatform
3
- Version: 2.5.0
3
+ Version: 2.6.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
@@ -296,9 +296,8 @@ def _build_generate_content_config(
296
296
  for key in [
297
297
  "system_instruction",
298
298
  "tools",
299
- "tools_config",
299
+ "tool_config",
300
300
  "safety_settings",
301
- "labels",
302
301
  ]:
303
302
  if key in request_dict:
304
303
  merged_config_dict[key] = request_dict[key]
@@ -3735,7 +3734,10 @@ def _convert_gcs_to_evaluation_item_result(
3735
3734
  return types.EvaluationItemResult(**eval_item_data)
3736
3735
  except Exception as e:
3737
3736
  logger.error(
3738
- "Failed to load evaluation result from GCS: %s. Error: %s", gcs_uri, e
3737
+ "Failed to load evaluation result from GCS: %s. Error: %s",
3738
+ gcs_uri,
3739
+ e,
3740
+ exc_info=True,
3739
3741
  )
3740
3742
  return types.EvaluationItemResult()
3741
3743
 
@@ -3752,7 +3754,10 @@ def _convert_gcs_to_evaluation_item_request(
3752
3754
  return types.EvaluationItemRequest(**eval_item_data)
3753
3755
  except Exception as e:
3754
3756
  logger.error(
3755
- "Failed to load evaluation request from GCS: %s. Error: %s", gcs_uri, e
3757
+ "Failed to load evaluation request from GCS: %s. Error: %s",
3758
+ gcs_uri,
3759
+ e,
3760
+ exc_info=True,
3756
3761
  )
3757
3762
  return types.EvaluationItemRequest()
3758
3763
 
@@ -3926,16 +3931,16 @@ def _get_eval_cases_eval_dfs_from_eval_items(
3926
3931
  item.
3927
3932
  - eval_dfs: A list of EvaluationDatasets, one for each candidate.
3928
3933
  """
3929
- dataset_rows = []
3934
+ dataset_rows: list[dict[str, Any]] = []
3930
3935
  eval_case_results = []
3931
- for index, eval_item in enumerate(eval_items):
3936
+ for eval_item in eval_items:
3932
3937
  if (
3933
3938
  eval_item
3934
3939
  and eval_item.evaluation_response
3935
3940
  and eval_item.evaluation_response.request
3936
3941
  ):
3937
3942
  eval_case_results.append(
3938
- _get_eval_case_result_from_eval_item(index, eval_item)
3943
+ _get_eval_case_result_from_eval_item(len(dataset_rows), eval_item)
3939
3944
  )
3940
3945
  dataset_rows.append(
3941
3946
  _convert_request_to_dataset_row(eval_item.evaluation_response.request)
@@ -4120,6 +4125,36 @@ def _get_content(row: dict[str, Any], column: str) -> Optional[genai_types.Conte
4120
4125
  )
4121
4126
 
4122
4127
 
4128
+ # Columns mapped to EvaluationItemRequest fields; all others go to extra_context.
4129
+ _EVALUATION_SET_STANDARD_COLUMNS = _evals_constant.COMMON_DATASET_COLUMNS | {
4130
+ _evals_constant.RESPONSE,
4131
+ "interaction",
4132
+ "gemini_agent",
4133
+ }
4134
+
4135
+
4136
+ def _to_extra_context_content(value: Any) -> Optional[genai_types.Content]:
4137
+ if isinstance(value, str):
4138
+ return genai_types.Content(parts=[genai_types.Part(text=value)])
4139
+ if isinstance(value, genai_types.Content):
4140
+ return value
4141
+ parts = []
4142
+ for content in _evals_metric_handlers._value_to_content_list(value):
4143
+ parts.extend(content.parts or [])
4144
+ return genai_types.Content(parts=parts) if parts else None
4145
+
4146
+
4147
+ def _get_extra_context(row: "pd.Series") -> Optional[dict[str, genai_types.Content]]:
4148
+ extra_context = {}
4149
+ for column, value in row.items():
4150
+ if column in _EVALUATION_SET_STANDARD_COLUMNS or _is_missing(value):
4151
+ continue
4152
+ content = _to_extra_context_content(value)
4153
+ if content:
4154
+ extra_context[str(column)] = content
4155
+ return extra_context or None
4156
+
4157
+
4123
4158
  def _create_evaluation_set_from_dataframe(
4124
4159
  api_client: BaseApiClient,
4125
4160
  gcs_dest_prefix: str,
@@ -4244,6 +4279,7 @@ def _create_evaluation_set_from_dataframe(
4244
4279
  candidate_responses=(
4245
4280
  candidate_responses if candidate_responses else None
4246
4281
  ),
4282
+ extra_context=_get_extra_context(row),
4247
4283
  )
4248
4284
  )
4249
4285
  logger.info("Writing evaluation item requests to GCS.")
@@ -230,6 +230,18 @@ class _FlattenEvalDataConverter(_evals_utils.EvalDataConverter):
230
230
  "Invalid prompt type for case %s: %s" % (i, type(prompt_data))
231
231
  )
232
232
 
233
+ if (
234
+ isinstance(conversation_history_data, str)
235
+ and conversation_history_data.strip()
236
+ ):
237
+ try:
238
+ conversation_history_data = json.loads(conversation_history_data)
239
+ except json.JSONDecodeError:
240
+ logger.warning(
241
+ "Could not decode conversation_history JSON string for case %s.",
242
+ eval_case_id,
243
+ )
244
+
233
245
  conversation_history: Optional[list[types.evals.Message]] = None
234
246
  if isinstance(conversation_history_data, list):
235
247
  conversation_history = []
@@ -242,6 +254,11 @@ class _FlattenEvalDataConverter(_evals_utils.EvalDataConverter):
242
254
  )
243
255
  )
244
256
  elif isinstance(content, dict):
257
+ if isinstance(content.get("content"), str):
258
+ content = {
259
+ "role": content.get("role"),
260
+ "parts": [{"text": content["content"]}],
261
+ }
245
262
  try:
246
263
  validated_content = genai_types.Content.model_validate(
247
264
  content
@@ -23,7 +23,7 @@ import textwrap
23
23
  from typing import Any, Optional
24
24
 
25
25
  import pandas as pd
26
- from pydantic import errors
26
+ import pydantic_core
27
27
 
28
28
  from . import _evals_common
29
29
  from . import types
@@ -1483,7 +1483,7 @@ def display_evaluation_result(
1483
1483
  result_dump = eval_result_obj.model_dump(
1484
1484
  mode="json", exclude_none=True, exclude={"evaluation_dataset"}
1485
1485
  )
1486
- except errors.PydanticSerializationError as e:
1486
+ except pydantic_core.PydanticSerializationError as e:
1487
1487
  logger.error(
1488
1488
  "Serialization Error: %s\nCould not display the evaluation "
1489
1489
  "result due to a data serialization issue. Please check the "
@@ -590,6 +590,44 @@ def _StructuredMemorySchemaConfig_to_vertex(
590
590
  return to_object
591
591
 
592
592
 
593
+ def _UpdateMemoryBankConfig_to_vertex(
594
+ from_object: Union[dict[str, Any], object],
595
+ parent_object: Optional[dict[str, Any]] = None,
596
+ ) -> dict[str, Any]:
597
+ to_object: dict[str, Any] = {}
598
+
599
+ if getv(from_object, ["display_name"]) is not None:
600
+ setv(parent_object, ["displayName"], getv(from_object, ["display_name"]))
601
+
602
+ if getv(from_object, ["description"]) is not None:
603
+ setv(parent_object, ["description"], getv(from_object, ["description"]))
604
+
605
+ return to_object
606
+
607
+
608
+ def _UpdateMemoryBankRequestParameters_to_vertex(
609
+ from_object: Union[dict[str, Any], object],
610
+ parent_object: Optional[dict[str, Any]] = None,
611
+ ) -> dict[str, Any]:
612
+ to_object: dict[str, Any] = {}
613
+ if getv(from_object, ["name"]) is not None:
614
+ setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
615
+
616
+ if getv(from_object, ["memory_bank_config"]) is not None:
617
+ setv(
618
+ to_object,
619
+ ["context_spec", "memory_bank_config"],
620
+ _ReasoningEngineContextSpecMemoryBankConfig_to_vertex(
621
+ getv(from_object, ["memory_bank_config"]), to_object
622
+ ),
623
+ )
624
+
625
+ if getv(from_object, ["config"]) is not None:
626
+ _UpdateMemoryBankConfig_to_vertex(getv(from_object, ["config"]), to_object)
627
+
628
+ return to_object
629
+
630
+
593
631
  class MemoryBanks(_api_module.BaseModule):
594
632
 
595
633
  def _create(
@@ -669,6 +707,85 @@ class MemoryBanks(_api_module.BaseModule):
669
707
  self._api_client._verify_response(return_value)
670
708
  return return_value
671
709
 
710
+ def _update(
711
+ self,
712
+ *,
713
+ name: str,
714
+ memory_bank_config: Optional[
715
+ types.ReasoningEngineContextSpecMemoryBankConfigOrDict
716
+ ] = None,
717
+ config: Optional[types.UpdateMemoryBankConfigOrDict] = None,
718
+ ) -> types.MemoryBankOperation:
719
+ """
720
+ Updates a Memory Bank.
721
+ """
722
+
723
+ parameter_model = types._UpdateMemoryBankRequestParameters(
724
+ name=name,
725
+ memory_bank_config=memory_bank_config,
726
+ config=config,
727
+ )
728
+
729
+ request_url_dict: Optional[dict[str, str]]
730
+ if not self._api_client.vertexai:
731
+ raise ValueError(
732
+ "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
733
+ )
734
+ else:
735
+ request_dict = _UpdateMemoryBankRequestParameters_to_vertex(parameter_model)
736
+ request_url_dict = request_dict.get("_url")
737
+ if request_url_dict:
738
+ path = "{name}".format_map(request_url_dict)
739
+ else:
740
+ path = "{name}"
741
+
742
+ query_params = request_dict.get("_query")
743
+ if query_params:
744
+ path = f"{path}?{urlencode(query_params)}"
745
+ # TODO: remove the hack that pops config.
746
+ request_dict.pop("config", None)
747
+
748
+ http_options: Optional[types.HttpOptions] = None
749
+ if (
750
+ parameter_model.config is not None
751
+ and parameter_model.config.http_options is not None
752
+ ):
753
+ http_options = parameter_model.config.http_options
754
+
755
+ request_dict = _common.convert_to_dict(request_dict)
756
+ request_dict = _common.encode_unserializable_types(request_dict)
757
+
758
+ response = self._api_client.request("patch", path, request_dict, http_options)
759
+
760
+ response_dict = {} if not response.body else json.loads(response.body)
761
+
762
+ if self._api_client.vertexai:
763
+ response_dict = _MemoryBankOperation_from_vertex(response_dict)
764
+
765
+ return_value = types.MemoryBankOperation._from_response(
766
+ response=response_dict,
767
+ kwargs=(
768
+ {
769
+ "config": {
770
+ "response_schema": getattr(
771
+ parameter_model.config, "response_schema", None
772
+ ),
773
+ "response_json_schema": getattr(
774
+ parameter_model.config, "response_json_schema", None
775
+ ),
776
+ "include_all_fields": getattr(
777
+ parameter_model.config, "include_all_fields", None
778
+ ),
779
+ }
780
+ }
781
+ if getattr(parameter_model, "config", None)
782
+ else {}
783
+ ),
784
+ )
785
+
786
+ self._api_client._verify_response(return_value)
787
+ return return_value
788
+
672
789
  def _delete(
673
790
  self,
674
791
  *,
@@ -1082,6 +1199,41 @@ class MemoryBanks(_api_module.BaseModule):
1082
1199
  memory_bank = types.MemoryBank(name=operation.response.name)
1083
1200
  return memory_bank
1084
1201
 
1202
+ def update(
1203
+ self,
1204
+ *,
1205
+ name: str,
1206
+ managed_semantic_memory_config: Optional[
1207
+ types.ManagedSemanticMemoryConfigOrDict
1208
+ ] = None,
1209
+ config: Optional[types.UpdateMemoryBankConfigOrDict] = None,
1210
+ ) -> types.MemoryBank:
1211
+ """Updates a Memory Bank resource."""
1212
+ memory_bank_config = (
1213
+ _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config(
1214
+ managed_semantic_memory_config
1215
+ )
1216
+ )
1217
+
1218
+ operation = self._update(
1219
+ name=name,
1220
+ memory_bank_config=memory_bank_config,
1221
+ config=config,
1222
+ )
1223
+
1224
+ operation = _memory_bank_utils._await_operation(
1225
+ operation_name=operation.name,
1226
+ get_operation_fn=self._get_memory_bank_operation,
1227
+ )
1228
+ if operation.error:
1229
+ raise RuntimeError(operation.error)
1230
+ if not operation.response:
1231
+ raise RuntimeError(
1232
+ f"Memory Bank update completed without a response: {operation}"
1233
+ )
1234
+ memory_bank = types.MemoryBank(name=operation.response.name)
1235
+ return memory_bank
1236
+
1085
1237
  def delete(
1086
1238
  self,
1087
1239
  *,
@@ -1332,6 +1484,87 @@ class AsyncMemoryBanks(_api_module.BaseModule):
1332
1484
  self._api_client._verify_response(return_value)
1333
1485
  return return_value
1334
1486
 
1487
+ async def _update(
1488
+ self,
1489
+ *,
1490
+ name: str,
1491
+ memory_bank_config: Optional[
1492
+ types.ReasoningEngineContextSpecMemoryBankConfigOrDict
1493
+ ] = None,
1494
+ config: Optional[types.UpdateMemoryBankConfigOrDict] = None,
1495
+ ) -> types.MemoryBankOperation:
1496
+ """
1497
+ Updates a Memory Bank.
1498
+ """
1499
+
1500
+ parameter_model = types._UpdateMemoryBankRequestParameters(
1501
+ name=name,
1502
+ memory_bank_config=memory_bank_config,
1503
+ config=config,
1504
+ )
1505
+
1506
+ request_url_dict: Optional[dict[str, str]]
1507
+ if not self._api_client.vertexai:
1508
+ raise ValueError(
1509
+ "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
1510
+ )
1511
+ else:
1512
+ request_dict = _UpdateMemoryBankRequestParameters_to_vertex(parameter_model)
1513
+ request_url_dict = request_dict.get("_url")
1514
+ if request_url_dict:
1515
+ path = "{name}".format_map(request_url_dict)
1516
+ else:
1517
+ path = "{name}"
1518
+
1519
+ query_params = request_dict.get("_query")
1520
+ if query_params:
1521
+ path = f"{path}?{urlencode(query_params)}"
1522
+ # TODO: remove the hack that pops config.
1523
+ request_dict.pop("config", None)
1524
+
1525
+ http_options: Optional[types.HttpOptions] = None
1526
+ if (
1527
+ parameter_model.config is not None
1528
+ and parameter_model.config.http_options is not None
1529
+ ):
1530
+ http_options = parameter_model.config.http_options
1531
+
1532
+ request_dict = _common.convert_to_dict(request_dict)
1533
+ request_dict = _common.encode_unserializable_types(request_dict)
1534
+
1535
+ response = await self._api_client.async_request(
1536
+ "patch", path, request_dict, http_options
1537
+ )
1538
+
1539
+ response_dict = {} if not response.body else json.loads(response.body)
1540
+
1541
+ if self._api_client.vertexai:
1542
+ response_dict = _MemoryBankOperation_from_vertex(response_dict)
1543
+
1544
+ return_value = types.MemoryBankOperation._from_response(
1545
+ response=response_dict,
1546
+ kwargs=(
1547
+ {
1548
+ "config": {
1549
+ "response_schema": getattr(
1550
+ parameter_model.config, "response_schema", None
1551
+ ),
1552
+ "response_json_schema": getattr(
1553
+ parameter_model.config, "response_json_schema", None
1554
+ ),
1555
+ "include_all_fields": getattr(
1556
+ parameter_model.config, "include_all_fields", None
1557
+ ),
1558
+ }
1559
+ }
1560
+ if getattr(parameter_model, "config", None)
1561
+ else {}
1562
+ ),
1563
+ )
1564
+
1565
+ self._api_client._verify_response(return_value)
1566
+ return return_value
1567
+
1335
1568
  async def _delete(
1336
1569
  self,
1337
1570
  *,
@@ -1755,6 +1988,41 @@ class AsyncMemoryBanks(_api_module.BaseModule):
1755
1988
  memory_bank = types.MemoryBank(name=operation.response.name)
1756
1989
  return memory_bank
1757
1990
 
1991
+ async def update(
1992
+ self,
1993
+ *,
1994
+ name: str,
1995
+ managed_semantic_memory_config: Optional[
1996
+ types.ManagedSemanticMemoryConfigOrDict
1997
+ ] = None,
1998
+ config: Optional[types.UpdateMemoryBankConfigOrDict] = None,
1999
+ ) -> types.MemoryBank:
2000
+ """Updates a Memory Bank."""
2001
+ memory_bank_config = (
2002
+ _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config(
2003
+ managed_semantic_memory_config
2004
+ )
2005
+ )
2006
+
2007
+ operation = await self._update(
2008
+ name=name,
2009
+ memory_bank_config=memory_bank_config,
2010
+ config=config,
2011
+ )
2012
+
2013
+ operation = await _memory_bank_utils._await_async_operation(
2014
+ operation_name=operation.name,
2015
+ get_operation_fn=self._get_memory_bank_operation,
2016
+ )
2017
+ if operation.error:
2018
+ raise RuntimeError(operation.error)
2019
+ if not operation.response:
2020
+ raise RuntimeError(
2021
+ f"Memory Bank update completed without a response: {operation}"
2022
+ )
2023
+ memory_bank = types.MemoryBank(name=operation.response.name)
2024
+ return memory_bank
2025
+
1758
2026
  async def delete(
1759
2027
  self,
1760
2028
  *,