Skip to content

Commit 0ac315e

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add client.memory_banks.update method.
PiperOrigin-RevId: 999011108
1 parent 07c7fae commit 0ac315e

5 files changed

Lines changed: 786 additions & 0 deletions

File tree

‎agentplatform/_genai/memory_banks.py‎

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -590,6 +590,44 @@ def _StructuredMemorySchemaConfig_to_vertex(
590590
return to_object
591591

592592

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+
593631
class MemoryBanks(_api_module.BaseModule):
594632

595633
def _create(
@@ -669,6 +707,85 @@ def _create(
669707
self._api_client._verify_response(return_value)
670708
return return_value
671709

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+
672789
def _delete(
673790
self,
674791
*,
@@ -1082,6 +1199,41 @@ def create(
10821199
memory_bank = types.MemoryBank(name=operation.response.name)
10831200
return memory_bank
10841201

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+
10851237
def delete(
10861238
self,
10871239
*,
@@ -1332,6 +1484,87 @@ async def _create(
13321484
self._api_client._verify_response(return_value)
13331485
return return_value
13341486

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+
13351568
async def _delete(
13361569
self,
13371570
*,
@@ -1755,6 +1988,41 @@ async def create(
17551988
memory_bank = types.MemoryBank(name=operation.response.name)
17561989
return memory_bank
17571990

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+
17582026
async def delete(
17592027
self,
17602028
*,

‎agentplatform/_genai/types/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,7 @@
183183
from .common import _UndeployModelRequestParameters
184184
from .common import _UpdateDatasetParameters
185185
from .common import _UpdateEvaluationExperimentParameters
186+
from .common import _UpdateMemoryBankRequestParameters
186187
from .common import _UpdateMemoryRequestParameters
187188
from .common import _UpdateMultimodalDatasetParameters
188189
from .common import _UpdateRagConfigRequestParameters
@@ -2265,6 +2266,9 @@
22652266
from .common import UpdateEvaluationExperimentConfig
22662267
from .common import UpdateEvaluationExperimentConfigDict
22672268
from .common import UpdateEvaluationExperimentConfigOrDict
2269+
from .common import UpdateMemoryBankConfig
2270+
from .common import UpdateMemoryBankConfigDict
2271+
from .common import UpdateMemoryBankConfigOrDict
22682272
from .common import UpdateMemoryConfig
22692273
from .common import UpdateMemoryConfigDict
22702274
from .common import UpdateMemoryConfigOrDict
@@ -3024,6 +3028,9 @@
30243028
"MemoryBankOperation",
30253029
"MemoryBankOperationDict",
30263030
"MemoryBankOperationOrDict",
3031+
"UpdateMemoryBankConfig",
3032+
"UpdateMemoryBankConfigDict",
3033+
"UpdateMemoryBankConfigOrDict",
30273034
"DeleteMemoryBankConfig",
30283035
"DeleteMemoryBankConfigDict",
30293036
"DeleteMemoryBankConfigOrDict",
@@ -4493,6 +4500,7 @@
44934500
"_GetDeleteRuntimeRevisionOperationParameters",
44944501
"_QueryRuntimeRevisionRequestParameters",
44954502
"_CreateMemoryBankRequestParameters",
4503+
"_UpdateMemoryBankRequestParameters",
44964504
"_DeleteMemoryBankRequestParameters",
44974505
"_GetMemoryBankRequestParameters",
44984506
"_IngestEventsRequestParameters",

0 commit comments

Comments
 (0)