diff --git a/sdk/ai/azure-ai-agents/azure/ai/agents/aio/operations/_patch.py b/sdk/ai/azure-ai-agents/azure/ai/agents/aio/operations/_patch.py index cddf3a7c3225..b9db5b57a607 100644 --- a/sdk/ai/azure-ai-agents/azure/ai/agents/aio/operations/_patch.py +++ b/sdk/ai/azure-ai-agents/azure/ai/agents/aio/operations/_patch.py @@ -529,7 +529,7 @@ async def create_and_process( RunStatus.IN_PROGRESS, RunStatus.REQUIRES_ACTION, ]: - time.sleep(sleep_interval) + await asyncio.sleep(sleep_interval) run = await self.get(thread_id=thread_id, run_id=run.id) if run.status == "requires_action" and isinstance(run.required_action, _models.SubmitToolOutputsAction): @@ -1389,7 +1389,9 @@ async def upload( raise ValueError("Invalid parameters for upload_file. Please provide the necessary arguments.") @overload - async def upload_and_poll(self, body: JSON, *, sleep_interval: float = 1, **kwargs: Any) -> _models.FileInfo: + async def upload_and_poll( + self, body: JSON, *, sleep_interval: float = 1, timeout: Optional[float] = None, **kwargs: Any + ) -> _models.FileInfo: """Uploads a file for use by other operations. :param body: Required. @@ -1397,9 +1399,12 @@ async def upload_and_poll(self, body: JSON, *, sleep_interval: float = 1, **kwar :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1410,6 +1415,7 @@ async def upload_and_poll( purpose: Union[str, _models.FilePurpose], filename: Optional[str] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.FileInfo: """Uploads a file for use by other operations. @@ -1424,14 +1430,23 @@ async def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload async def upload_and_poll( - self, *, file_path: str, purpose: Union[str, _models.FilePurpose], sleep_interval: float = 1, **kwargs: Any + self, + *, + file_path: str, + purpose: Union[str, _models.FilePurpose], + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.FileInfo: """Uploads a file for use by other operations. @@ -1443,9 +1458,12 @@ async def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @distributed_trace_async @@ -1458,6 +1476,7 @@ async def upload_and_poll( purpose: Union[str, _models.FilePurpose, None] = None, filename: Optional[str] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.FileInfo: """ @@ -1477,12 +1496,17 @@ async def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: _models.FileInfo :raises FileNotFoundError: If the file_path is invalid. :raises IOError: If there are issues with reading the file. :raises: HttpResponseError for HTTP errors. + :raises TimeoutError: If the polling times out. """ + + curr_time = time.monotonic() if body is not None: uploaded_file = await self.upload(body=body, **kwargs) elif file is not None and purpose is not None: @@ -1491,12 +1515,16 @@ async def upload_and_poll( uploaded_file = await self.upload(file_path=file_path, purpose=purpose, **kwargs) else: raise ValueError( - "Invalid parameters for upload_file_and_poll. Please provide either 'body', " + "Invalid parameters for upload_and_poll. Please provide either 'body', " "or both 'file' and 'purpose', or both 'file_path' and 'purpose'." ) while uploaded_file.status in ["uploaded", "pending", "running"]: - time.sleep(sleep_interval) + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + + await asyncio.sleep(sleep_interval) uploaded_file = await self.get(uploaded_file.id) return uploaded_file @@ -1580,7 +1608,13 @@ class VectorStoresOperations(VectorStoresOperationsGenerated): @overload async def create_and_poll( - self, body: JSON, *, content_type: str = "application/json", sleep_interval: float = 1, **kwargs: Any + self, + body: JSON, + *, + content_type: str = "application/json", + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1592,9 +1626,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1609,6 +1646,7 @@ async def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, metadata: Optional[Dict[str, str]] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1636,14 +1674,23 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload async def create_and_poll( - self, body: IO[bytes], *, content_type: str = "application/json", sleep_interval: float = 1, **kwargs: Any + self, + body: IO[bytes], + *, + content_type: str = "application/json", + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1655,9 +1702,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @distributed_trace_async @@ -1673,6 +1723,7 @@ async def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, metadata: Optional[Dict[str, str]] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1702,11 +1753,14 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ - + curr_time = time.monotonic() if body is not _Unset: if isinstance(body, dict): vector_store = await super().create( @@ -1732,7 +1786,11 @@ async def create_and_poll( ) while vector_store.status == "in_progress": - time.sleep(sleep_interval) + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + + await asyncio.sleep(sleep_interval) vector_store = await super().get(vector_store.id) return vector_store @@ -1748,6 +1806,7 @@ async def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1762,9 +1821,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1777,6 +1839,7 @@ async def create_and_poll( content_type: str = "application/json", chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1796,9 +1859,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1809,6 +1875,7 @@ async def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1823,9 +1890,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @distributed_trace_async @@ -1839,6 +1909,7 @@ async def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1859,11 +1930,16 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ + curr_time = time.monotonic() + if body is not _Unset: if isinstance(body, dict): vector_store_file_batch = await super().create( @@ -1891,7 +1967,10 @@ async def create_and_poll( ) while vector_store_file_batch.status == "in_progress": - time.sleep(sleep_interval) + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + + await asyncio.sleep(sleep_interval) vector_store_file_batch = await super().get( vector_store_id=vector_store_id, batch_id=vector_store_file_batch.id ) @@ -1909,6 +1988,7 @@ async def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1923,9 +2003,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1938,6 +2021,7 @@ async def create_and_poll( data_source: Optional[_models.VectorStoreDataSource] = None, chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1957,9 +2041,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @overload @@ -1970,6 +2057,7 @@ async def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1984,9 +2072,12 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ @distributed_trace_async @@ -2000,6 +2091,7 @@ async def create_and_poll( data_source: Optional[_models.VectorStoreDataSource] = None, chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -2020,11 +2112,16 @@ async def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the polling times out. """ + curr_time = time.monotonic() + if body is not _Unset: if isinstance(body, dict): vector_store_file = await super().create( @@ -2052,7 +2149,11 @@ async def create_and_poll( ) while vector_store_file.status == "in_progress": - time.sleep(sleep_interval) + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + + await asyncio.sleep(sleep_interval) vector_store_file = await super().get(vector_store_id=vector_store_id, file_id=vector_store_file.id) return vector_store_file diff --git a/sdk/ai/azure-ai-agents/azure/ai/agents/models/_patch.py b/sdk/ai/azure-ai-agents/azure/ai/agents/models/_patch.py index 2c744b997b74..708adfa86c41 100644 --- a/sdk/ai/azure-ai-agents/azure/ai/agents/models/_patch.py +++ b/sdk/ai/azure-ai-agents/azure/ai/agents/models/_patch.py @@ -877,6 +877,7 @@ class BingGroundingTool(Tool[BingGroundingToolDefinition]): """ A tool that searches for information using Bing. """ + def __init__(self, connection_id: str, market: str = "", set_lang: str = "", count: int = 5, freshness: str = ""): """ Initialize Bing Custom Search with a connection_id. @@ -889,11 +890,7 @@ def __init__(self, connection_id: str, market: str = "", set_lang: str = "", cou """ self.connection_ids = [ BingGroundingSearchConfiguration( - connection_id=connection_id, - market=market, - set_lang=set_lang, - count=count, - freshness=freshness + connection_id=connection_id, market=market, set_lang=set_lang, count=count, freshness=freshness ) ] @@ -935,12 +932,7 @@ def __init__(self, connection_id: str, instance_name: str): :param connection_id: Connection ID used by tool. Bing Custom Search tools allow only one connection. :param instance_name: Config instance name used by tool. """ - self.connection_ids = [ - BingCustomSearchConfiguration( - connection_id=connection_id, - instance_name=instance_name - ) - ] + self.connection_ids = [BingCustomSearchConfiguration(connection_id=connection_id, instance_name=instance_name)] @property def definitions(self) -> List[BingCustomSearchToolDefinition]: @@ -951,9 +943,7 @@ def definitions(self) -> List[BingCustomSearchToolDefinition]: """ return [ BingCustomSearchToolDefinition( - bing_custom_search=BingCustomSearchConfigurationList( - search_configurations=self.connection_ids - ) + bing_custom_search=BingCustomSearchConfigurationList(search_configurations=self.connection_ids) ) ] diff --git a/sdk/ai/azure-ai-agents/azure/ai/agents/operations/_patch.py b/sdk/ai/azure-ai-agents/azure/ai/agents/operations/_patch.py index b7b22f7db9ad..36510c30c5fe 100644 --- a/sdk/ai/azure-ai-agents/azure/ai/agents/operations/_patch.py +++ b/sdk/ai/azure-ai-agents/azure/ai/agents/operations/_patch.py @@ -1402,7 +1402,9 @@ def upload( raise ValueError("Invalid parameters for upload_file. Please provide the necessary arguments.") @overload - def upload_and_poll(self, body: JSON, *, sleep_interval: float = 1, **kwargs: Any) -> _models.FileInfo: + def upload_and_poll( + self, body: JSON, *, sleep_interval: float = 1, timeout: Optional[float] = None, **kwargs: Any + ) -> _models.FileInfo: """Uploads a file for use by other operations. :param body: Required. @@ -1410,9 +1412,12 @@ def upload_and_poll(self, body: JSON, *, sleep_interval: float = 1, **kwargs: An :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1423,6 +1428,7 @@ def upload_and_poll( purpose: Union[str, _models.FilePurpose], filename: Optional[str] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.FileInfo: """Uploads a file for use by other operations. @@ -1437,14 +1443,23 @@ def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload def upload_and_poll( - self, *, file_path: str, purpose: Union[str, _models.FilePurpose], sleep_interval: float = 1, **kwargs: Any + self, + *, + file_path: str, + purpose: Union[str, _models.FilePurpose], + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.FileInfo: """Uploads a file for use by other operations. @@ -1456,9 +1471,12 @@ def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: ~azure.ai.agents.models.FileInfo :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @distributed_trace @@ -1471,6 +1489,7 @@ def upload_and_poll( purpose: Union[str, _models.FilePurpose, None] = None, filename: Optional[str] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.FileInfo: """ @@ -1490,12 +1509,17 @@ def upload_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the uploaded file. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the uploaded file. + :paramtype timeout: float :return: FileInfo. The FileInfo is compatible with MutableMapping :rtype: _models.FileInfo :raises FileNotFoundError: If the file_path is invalid. :raises IOError: If there are issues with reading the file. :raises: HttpResponseError for HTTP errors. + :raises TimeoutError: If the operation times out while polling for status. """ + + curr_time = time.monotonic() if body is not None: uploaded_file = self.upload(body=body, **kwargs) elif file is not None and purpose is not None: @@ -1504,11 +1528,15 @@ def upload_and_poll( uploaded_file = self.upload(file_path=file_path, purpose=purpose, **kwargs) else: raise ValueError( - "Invalid parameters for upload_file_and_poll. Please provide either 'body', " + "Invalid parameters for upload_and_poll. Please provide either 'body', " "or both 'file' and 'purpose', or both 'file_path' and 'purpose'." ) while uploaded_file.status in ["uploaded", "pending", "running"]: + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + time.sleep(sleep_interval) uploaded_file = self.get(uploaded_file.id) @@ -1583,7 +1611,13 @@ class VectorStoresOperations(VectorStoresOperationsGenerated): @overload def create_and_poll( - self, body: JSON, *, content_type: str = "application/json", sleep_interval: float = 1, **kwargs: Any + self, + body: JSON, + *, + content_type: str = "application/json", + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1595,9 +1629,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1612,6 +1649,7 @@ def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, metadata: Optional[Dict[str, str]] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1639,14 +1677,23 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload def create_and_poll( - self, body: IO[bytes], *, content_type: str = "application/json", sleep_interval: float = 1, **kwargs: Any + self, + body: IO[bytes], + *, + content_type: str = "application/json", + sleep_interval: float = 1, + timeout: Optional[float] = None, + **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1658,9 +1705,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @distributed_trace @@ -1676,6 +1726,7 @@ def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, metadata: Optional[Dict[str, str]] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStore: """Creates a vector store and poll. @@ -1705,11 +1756,16 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStore. The VectorStore is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStore :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ + curr_time = time.monotonic() + if body is not _Unset: if isinstance(body, dict): vector_store = super().create(body=body, content_type=content_type or "application/json", **kwargs) @@ -1733,6 +1789,10 @@ def create_and_poll( ) while vector_store.status == "in_progress": + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + time.sleep(sleep_interval) vector_store = super().get(vector_store.id) @@ -1749,6 +1809,7 @@ def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1763,9 +1824,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1778,6 +1842,7 @@ def create_and_poll( content_type: str = "application/json", chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1797,9 +1862,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1810,6 +1878,7 @@ def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1824,9 +1893,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @distributed_trace @@ -1840,6 +1912,7 @@ def create_and_poll( chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFileBatch: """Create a vector store file batch and poll. @@ -1860,11 +1933,16 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFileBatch. The VectorStoreFileBatch is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFileBatch :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ + curr_time = time.monotonic() + if body is not _Unset: if isinstance(body, dict): vector_store_file_batch = super().create( @@ -1892,6 +1970,10 @@ def create_and_poll( ) while vector_store_file_batch.status == "in_progress": + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + time.sleep(sleep_interval) vector_store_file_batch = super().get(vector_store_id=vector_store_id, batch_id=vector_store_file_batch.id) @@ -1908,6 +1990,7 @@ def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1922,9 +2005,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1937,6 +2023,7 @@ def create_and_poll( data_source: Optional[_models.VectorStoreDataSource] = None, chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1956,9 +2043,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @overload @@ -1969,6 +2059,7 @@ def create_and_poll( *, content_type: str = "application/json", sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -1983,9 +2074,12 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raises TimeoutError: If the operation times out while polling for status. """ @distributed_trace @@ -1999,6 +2093,7 @@ def create_and_poll( data_source: Optional[_models.VectorStoreDataSource] = None, chunking_strategy: Optional[_models.VectorStoreChunkingStrategyRequest] = None, sleep_interval: float = 1, + timeout: Optional[float] = None, **kwargs: Any, ) -> _models.VectorStoreFile: """Create a vector store file by attaching a file to a vector store. @@ -2019,11 +2114,15 @@ def create_and_poll( :keyword sleep_interval: Time to wait before polling for the status of the vector store. Default value is 1. :paramtype sleep_interval: float + :keyword timeout: Time to wait before polling for the status of the vector store. + :paramtype timeout: float :return: VectorStoreFile. The VectorStoreFile is compatible with MutableMapping :rtype: ~azure.ai.agents.models.VectorStoreFile :raises ~azure.core.exceptions.HttpResponseError: + :raise TimeoutError: If the operation times out while polling for status. """ + curr_time = time.monotonic() if body is not _Unset: if isinstance(body, dict): vector_store_file = super().create( @@ -2051,6 +2150,10 @@ def create_and_poll( ) while vector_store_file.status == "in_progress": + + if timeout is not None and (time.monotonic() - curr_time - sleep_interval) >= timeout: + raise TimeoutError("Timeout reached. Stopping polling.") + time.sleep(sleep_interval) vector_store_file = super().get(vector_store_id=vector_store_id, file_id=vector_store_file.id)