From 0d45ddc75893624dfeb9efce37c281254f12217f Mon Sep 17 00:00:00 2001 From: Mohammad Mazraeh Date: Sun, 26 Jan 2025 14:06:25 -0800 Subject: [PATCH] update empty chunk tolerance param in create_streams Signed-off-by: Mohammad Mazraeh --- .../autogen-core/src/autogen_core/models/_model_client.py | 1 + python/packages/autogen-core/tests/test_tool_agent.py | 1 + .../src/autogen_ext/models/azure/_azure_ai_client.py | 2 ++ .../src/autogen_ext/models/cache/_chat_completion_cache.py | 2 ++ .../autogen_ext/models/replay/_replay_chat_completion_client.py | 1 + .../models/semantic_kernel/_sk_chat_completion_adapter.py | 1 + 6 files changed, 8 insertions(+) diff --git a/python/packages/autogen-core/src/autogen_core/models/_model_client.py b/python/packages/autogen-core/src/autogen_core/models/_model_client.py index 356fad5487ca..6451af95762d 100644 --- a/python/packages/autogen-core/src/autogen_core/models/_model_client.py +++ b/python/packages/autogen-core/src/autogen_core/models/_model_client.py @@ -74,6 +74,7 @@ def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: ... @abstractmethod diff --git a/python/packages/autogen-core/tests/test_tool_agent.py b/python/packages/autogen-core/tests/test_tool_agent.py index 85fcd3892c97..1826ad4271f1 100644 --- a/python/packages/autogen-core/tests/test_tool_agent.py +++ b/python/packages/autogen-core/tests/test_tool_agent.py @@ -123,6 +123,7 @@ def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: raise NotImplementedError() diff --git a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py index d10f678dc2af..87dae4ec2c82 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py @@ -295,6 +295,7 @@ async def create( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> CreateResult: extra_create_args_keys = set(extra_create_args.keys()) if not create_kwargs.issuperset(extra_create_args_keys): @@ -373,6 +374,7 @@ async def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: extra_create_args_keys = set(extra_create_args.keys()) if not create_kwargs.issuperset(extra_create_args_keys): diff --git a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py index 15ea1cfd2401..e4c76e860565 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py @@ -120,6 +120,7 @@ async def create( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> CreateResult: """ Cached version of ChatCompletionClient.create. @@ -152,6 +153,7 @@ def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: """ Cached version of ChatCompletionClient.create_stream. diff --git a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py index 5ae7b6b665eb..07db637b96a0 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py @@ -169,6 +169,7 @@ async def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: """Return the next completion as a stream.""" if self._current_index >= len(self.chat_completions): diff --git a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py index 9f501fbbffcf..6299b0303c65 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py @@ -384,6 +384,7 @@ async def create_stream( json_output: Optional[bool] = None, extra_create_args: Mapping[str, Any] = {}, cancellation_token: Optional[CancellationToken] = None, + max_consecutive_empty_chunk_tolerance: int = 0, ) -> AsyncGenerator[Union[str, CreateResult], None]: """Create a streaming chat completion using the Semantic Kernel client.