|
| 1 | +# pylint: disable=line-too-long,useless-suppression |
| 2 | +# ------------------------------------ |
| 3 | +# Copyright (c) Microsoft Corporation. |
| 4 | +# Licensed under the MIT License. |
| 5 | +# ------------------------------------ |
| 6 | +"""Customize generated code here. |
| 7 | +
|
| 8 | +Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize |
| 9 | +""" |
| 10 | +from typing import Union, Optional, Any, List, overload, IO, cast |
| 11 | +from azure.core.tracing.decorator_async import distributed_trace_async |
| 12 | +from azure.core.polling import AsyncNoPolling |
| 13 | +from azure.core.utils import case_insensitive_dict |
| 14 | +from ... import models as _models |
| 15 | +from ...models import ( |
| 16 | + MemoryStoreOperationUsage, |
| 17 | + MemoryStoreOperationUsageInputTokensDetails, |
| 18 | + MemoryStoreOperationUsageOutputTokensDetails, |
| 19 | + MemoryStoreUpdateCompletedResult, |
| 20 | + AsyncUpdateMemoriesLROPoller, |
| 21 | + AsyncUpdateMemoriesLROPollingMethod, |
| 22 | +) |
| 23 | +from ._operations import JSON, _Unset, ClsType, MemoryStoresOperations as GenerateMemoryStoresOperations |
| 24 | +from ..._validation import api_version_validation |
| 25 | +from ..._utils.model_base import _deserialize |
| 26 | + |
| 27 | + |
| 28 | +class MemoryStoresOperations(GenerateMemoryStoresOperations): |
| 29 | + |
| 30 | + @overload |
| 31 | + async def begin_update_memories( |
| 32 | + self, |
| 33 | + name: str, |
| 34 | + *, |
| 35 | + scope: str, |
| 36 | + content_type: str = "application/json", |
| 37 | + items: Optional[List[_models.ItemParam]] = None, |
| 38 | + previous_update_id: Optional[str] = None, |
| 39 | + update_delay: Optional[int] = None, |
| 40 | + **kwargs: Any, |
| 41 | + ) -> AsyncUpdateMemoriesLROPoller: |
| 42 | + """Update memory store with conversation memories. |
| 43 | +
|
| 44 | + :param name: The name of the memory store to update. Required. |
| 45 | + :type name: str |
| 46 | + :keyword scope: The namespace that logically groups and isolates memories, such as a user ID. |
| 47 | + Required. |
| 48 | + :paramtype scope: str |
| 49 | + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. |
| 50 | + Default value is "application/json". |
| 51 | + :paramtype content_type: str |
| 52 | + :keyword items: Conversation items from which to extract memories. Default value is None. |
| 53 | + :paramtype items: list[~azure.ai.projects.models.ItemParam] |
| 54 | + :keyword previous_update_id: The unique ID of the previous update request, enabling incremental |
| 55 | + memory updates from where the last operation left off. Default value is None. |
| 56 | + :paramtype previous_update_id: str |
| 57 | + :keyword update_delay: Timeout period before processing the memory update in seconds. |
| 58 | + If a new update request is received during this period, it will cancel the current request and |
| 59 | + reset the timeout. |
| 60 | + Set to 0 to immediately trigger the update without delay. |
| 61 | + Defaults to 300 (5 minutes). Default value is None. |
| 62 | + :paramtype update_delay: int |
| 63 | + :return: An instance of AsyncUpdateMemoriesLROPoller that returns MemoryStoreUpdateCompletedResult. The |
| 64 | + MemoryStoreUpdateCompletedResult is compatible with MutableMapping |
| 65 | + :rtype: |
| 66 | + ~azure.ai.projects.models.AsyncUpdateMemoriesLROPoller |
| 67 | + :raises ~azure.core.exceptions.HttpResponseError: |
| 68 | + """ |
| 69 | + |
| 70 | + @overload |
| 71 | + async def begin_update_memories( |
| 72 | + self, name: str, body: JSON, *, content_type: str = "application/json", **kwargs: Any |
| 73 | + ) -> AsyncUpdateMemoriesLROPoller: |
| 74 | + """Update memory store with conversation memories. |
| 75 | +
|
| 76 | + :param name: The name of the memory store to update. Required. |
| 77 | + :type name: str |
| 78 | + :param body: Required. |
| 79 | + :type body: JSON |
| 80 | + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. |
| 81 | + Default value is "application/json". |
| 82 | + :paramtype content_type: str |
| 83 | + :return: An instance of AsyncUpdateMemoriesLROPoller that returns MemoryStoreUpdateCompletedResult. The |
| 84 | + MemoryStoreUpdateCompletedResult is compatible with MutableMapping |
| 85 | + :rtype: |
| 86 | + ~azure.ai.projects.models.AsyncUpdateMemoriesLROPoller |
| 87 | + :raises ~azure.core.exceptions.HttpResponseError: |
| 88 | + """ |
| 89 | + |
| 90 | + @overload |
| 91 | + async def begin_update_memories( |
| 92 | + self, name: str, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any |
| 93 | + ) -> AsyncUpdateMemoriesLROPoller: |
| 94 | + """Update memory store with conversation memories. |
| 95 | +
|
| 96 | + :param name: The name of the memory store to update. Required. |
| 97 | + :type name: str |
| 98 | + :param body: Required. |
| 99 | + :type body: IO[bytes] |
| 100 | + :keyword content_type: Body Parameter content-type. Content type parameter for binary body. |
| 101 | + Default value is "application/json". |
| 102 | + :paramtype content_type: str |
| 103 | + :return: An instance of AsyncUpdateMemoriesLROPoller that returns MemoryStoreUpdateCompletedResult. The |
| 104 | + MemoryStoreUpdateCompletedResult is compatible with MutableMapping |
| 105 | + :rtype: |
| 106 | + ~azure.ai.projects.models.AsyncUpdateMemoriesLROPoller |
| 107 | + :raises ~azure.core.exceptions.HttpResponseError: |
| 108 | + """ |
| 109 | + |
| 110 | + @distributed_trace_async |
| 111 | + @api_version_validation( |
| 112 | + method_added_on="2025-11-15-preview", |
| 113 | + params_added_on={"2025-11-15-preview": ["api_version", "name", "content_type", "accept"]}, |
| 114 | + api_versions_list=["2025-11-15-preview"], |
| 115 | + ) |
| 116 | + async def begin_update_memories( |
| 117 | + self, |
| 118 | + name: str, |
| 119 | + body: Union[JSON, IO[bytes]] = _Unset, |
| 120 | + *, |
| 121 | + scope: str = _Unset, |
| 122 | + items: Optional[List[_models.ItemParam]] = None, |
| 123 | + previous_update_id: Optional[str] = None, |
| 124 | + update_delay: Optional[int] = None, |
| 125 | + **kwargs: Any, |
| 126 | + ) -> AsyncUpdateMemoriesLROPoller: |
| 127 | + """Update memory store with conversation memories. |
| 128 | +
|
| 129 | + :param name: The name of the memory store to update. Required. |
| 130 | + :type name: str |
| 131 | + :param body: Is either a JSON type or a IO[bytes] type. Required. |
| 132 | + :type body: JSON or IO[bytes] |
| 133 | + :keyword scope: The namespace that logically groups and isolates memories, such as a user ID. |
| 134 | + Required. |
| 135 | + :paramtype scope: str |
| 136 | + :keyword items: Conversation items from which to extract memories. Default value is None. |
| 137 | + :paramtype items: list[~azure.ai.projects.models.ItemParam] |
| 138 | + :keyword previous_update_id: The unique ID of the previous update request, enabling incremental |
| 139 | + memory updates from where the last operation left off. Default value is None. |
| 140 | + :paramtype previous_update_id: str |
| 141 | + :keyword update_delay: Timeout period before processing the memory update in seconds. |
| 142 | + If a new update request is received during this period, it will cancel the current request and |
| 143 | + reset the timeout. |
| 144 | + Set to 0 to immediately trigger the update without delay. |
| 145 | + Defaults to 300 (5 minutes). Default value is None. |
| 146 | + :paramtype update_delay: int |
| 147 | + :return: An instance of AsyncLROPoller that returns MemoryStoreUpdateCompletedResult. The |
| 148 | + MemoryStoreUpdateCompletedResult is compatible with MutableMapping |
| 149 | + :rtype: |
| 150 | + ~azure.ai.projects.models.AsyncUpdateMemoriesLROPoller |
| 151 | + :raises ~azure.core.exceptions.HttpResponseError: |
| 152 | + """ |
| 153 | + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) |
| 154 | + _params = kwargs.pop("params", {}) or {} |
| 155 | + |
| 156 | + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) |
| 157 | + cls: ClsType[_models.MemoryStoreUpdateCompletedResult] = kwargs.pop("cls", None) |
| 158 | + polling: Union[bool, AsyncUpdateMemoriesLROPollingMethod] = kwargs.pop("polling", True) |
| 159 | + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) |
| 160 | + cont_token: Optional[str] = kwargs.pop("continuation_token", None) |
| 161 | + if cont_token is None: |
| 162 | + raw_result = await self._update_memories_initial( |
| 163 | + name=name, |
| 164 | + body=body, |
| 165 | + scope=scope, |
| 166 | + items=items, |
| 167 | + previous_update_id=previous_update_id, |
| 168 | + update_delay=update_delay, |
| 169 | + content_type=content_type, |
| 170 | + cls=lambda x, y, z: x, |
| 171 | + headers=_headers, |
| 172 | + params=_params, |
| 173 | + **kwargs, |
| 174 | + ) |
| 175 | + await raw_result.http_response.read() # type: ignore |
| 176 | + |
| 177 | + raw_result.http_response.status_code = 202 # type: ignore |
| 178 | + raw_result.http_response.headers["Operation-Location"] = ( # type: ignore |
| 179 | + f"{self._config.endpoint}/memory_stores/{name}/updates/{raw_result.http_response.json().get('update_id')}?api-version=2025-11-15-preview" # type: ignore |
| 180 | + ) |
| 181 | + |
| 182 | + kwargs.pop("error_map", None) |
| 183 | + |
| 184 | + def get_long_running_output(pipeline_response): |
| 185 | + response_headers = {} |
| 186 | + response = pipeline_response.http_response |
| 187 | + response_headers["Operation-Location"] = self._deserialize( |
| 188 | + "str", response.headers.get("Operation-Location") |
| 189 | + ) |
| 190 | + |
| 191 | + deserialized = _deserialize(MemoryStoreUpdateCompletedResult, response.json().get("result", None)) |
| 192 | + if deserialized is None: |
| 193 | + usage = MemoryStoreOperationUsage( |
| 194 | + embedding_tokens=0, |
| 195 | + input_tokens=0, |
| 196 | + input_tokens_details=MemoryStoreOperationUsageInputTokensDetails(cached_tokens=0), |
| 197 | + output_tokens=0, |
| 198 | + output_tokens_details=MemoryStoreOperationUsageOutputTokensDetails(reasoning_tokens=0), |
| 199 | + total_tokens=0, |
| 200 | + ) |
| 201 | + deserialized = MemoryStoreUpdateCompletedResult(memory_operations=[], usage=usage) |
| 202 | + if cls: |
| 203 | + return cls(pipeline_response, deserialized, response_headers) # type: ignore |
| 204 | + return deserialized |
| 205 | + |
| 206 | + path_format_arguments = { |
| 207 | + "endpoint": self._serialize.url("self._config.endpoint", self._config.endpoint, "str", skip_quote=True), |
| 208 | + } |
| 209 | + |
| 210 | + if polling is True: |
| 211 | + polling_method: AsyncUpdateMemoriesLROPollingMethod = AsyncUpdateMemoriesLROPollingMethod( |
| 212 | + lro_delay, path_format_arguments=path_format_arguments, **kwargs |
| 213 | + ) |
| 214 | + elif polling is False: |
| 215 | + polling_method = cast(AsyncUpdateMemoriesLROPollingMethod, AsyncNoPolling()) |
| 216 | + else: |
| 217 | + polling_method = polling |
| 218 | + if cont_token: |
| 219 | + return AsyncUpdateMemoriesLROPoller.from_continuation_token( |
| 220 | + polling_method=polling_method, |
| 221 | + continuation_token=cont_token, |
| 222 | + client=self._client, |
| 223 | + deserialization_callback=get_long_running_output, |
| 224 | + ) |
| 225 | + return AsyncUpdateMemoriesLROPoller( |
| 226 | + self._client, |
| 227 | + raw_result, # type: ignore[possibly-undefined] |
| 228 | + get_long_running_output, |
| 229 | + polling_method, # pylint: disable=possibly-used-before-assignment |
| 230 | + ) |
0 commit comments