Add client close (#5871)

Fixes #4821 by adding a `close()` method to all clients.

Additionally:
* The m1 CLI is updated to close the client before exiting.
* The playwrightcontroller is updated to suppress some other unrelated
chatty warnings (e.g,, produced by markitdown when encountering
conversions that require external utilities)

Author: afourney

python/packages/autogen-core/src/autogen_core/models/_model_client.py

@@ -152,6 +152,9 @@ def create_stream(
         cancellation_token: Optional[CancellationToken] = None,
     ) -> AsyncGenerator[Union[str, CreateResult], None]: ...
 
+    @abstractmethod
+    async def close(self) -> None: ...
+
     @abstractmethod
     def actual_usage(self) -> RequestUsage: ...
 

python/packages/autogen-core/tests/test_tool_agent.py

@@ -126,6 +126,9 @@ def create_stream(
         ) -> AsyncGenerator[Union[str, CreateResult], None]:
             raise NotImplementedError()
 
+        async def close(self) -> None:
+            pass
+
         def actual_usage(self) -> RequestUsage:
             return RequestUsage(prompt_tokens=0, completion_tokens=0)
 

python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/playwright_controller.py

@@ -3,10 +3,13 @@
 import io
 import os
 import random
+import warnings
 from typing import Any, Callable, Dict, Optional, Tuple, Union, cast
 
 # TODO: Fix unfollowed import
 try:
+    # Suppress warnings from markitdown -- which is pretty chatty
+    warnings.filterwarnings(action="ignore", module="markitdown")
     from markitdown import MarkItDown  # type: ignore
 except ImportError:
     MarkItDown = None

python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py

@@ -166,6 +166,9 @@ def create_stream(
             cancellation_token=cancellation_token,
         )
 
+    async def close(self) -> None:
+        await self.base_client.close()
+
     def actual_usage(self) -> RequestUsage:
         # Calls base_client.actual_usage() and returns the result.
         return self.base_client.actual_usage()

python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py

@@ -775,6 +775,9 @@ async def create_stream(
 
         yield result
 
+    async def close(self) -> None:
+        await self._client.close()
+
     def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | ToolSchema] = []) -> int:
         """
         Estimate the number of tokens used by messages and tools.

python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py

@@ -490,6 +490,9 @@ async def create_stream(
 
         yield result
 
+    async def close(self) -> None:
+        await self._client.close()
+
     def actual_usage(self) -> RequestUsage:
         return self._actual_usage
 

python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py

@@ -206,6 +206,9 @@ async def _generator() -> AsyncGenerator[Union[str, CreateResult], None]:
 
         return _generator()
 
+    async def close(self) -> None:
+        await self.client.close()
+
     def actual_usage(self) -> RequestUsage:
         return self.client.actual_usage()
 

python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py

@@ -772,6 +772,9 @@ async def create_stream(
 
         yield result
 
+    async def close(self) -> None:
+        pass  # ollama has no close method?
+
     def actual_usage(self) -> RequestUsage:
         return self._actual_usage
 

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -944,6 +944,9 @@ async def _create_stream_chunks_beta_client(
                 except StopAsyncIteration:
                     break
 
+    async def close(self) -> None:
+        await self._client.close()
+
     def actual_usage(self) -> RequestUsage:
         return self._actual_usage
 

python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py

@@ -3,7 +3,6 @@
 import logging
 import warnings
 from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Union
-from typing_extensions import Self
 
 from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component
 from autogen_core.models import (
@@ -18,6 +17,7 @@
 )
 from autogen_core.tools import Tool, ToolSchema
 from pydantic import BaseModel
+from typing_extensions import Self
 
 logger = logging.getLogger(EVENT_LOGGER_NAME)
 
@@ -229,6 +229,9 @@ async def create_stream(
 
         self._current_index += 1
 
+    async def close(self) -> None:
+        pass
+
     def actual_usage(self) -> RequestUsage:
         return self._cur_usage
 

python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py

@@ -654,6 +654,9 @@ async def create_stream(
             thought=thought,
         )
 
+    async def close(self) -> None:
+        pass  # No explicit close method in SK client?
+
     def actual_usage(self) -> RequestUsage:
         return RequestUsage(prompt_tokens=self._total_prompt_tokens, completion_tokens=self._total_completion_tokens)
 

python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py

@@ -2,7 +2,6 @@
 import asyncio
 import os
 import sys
-import warnings
 from typing import Any, Dict, Optional
 
 import yaml
@@ -13,9 +12,6 @@
 from autogen_ext.teams.magentic_one import MagenticOne
 from autogen_ext.ui import RichConsole
 
-# Suppress warnings about the requests.Session() not being closed
-warnings.filterwarnings(action="ignore", message="unclosed", category=ResourceWarning)
-
 DEFAULT_CONFIG_FILE = "config.yaml"
 DEFAULT_CONFIG_CONTENTS = """# config.yaml
 #
@@ -109,10 +105,9 @@ def main() -> None:
         with open(args.config if isinstance(args.config, str) else args.config[0], "r") as f:
             config = yaml.safe_load(f)
 
-    client = ChatCompletionClient.load_component(config["client"])
-
     # Run the task
     async def run_task(task: str, hil_mode: bool, use_rich_console: bool) -> None:
+        client = ChatCompletionClient.load_component(config["client"])
         input_manager = UserInputManager(callback=cancellable_input)
 
         async with DockerCommandLineCodeExecutor(work_dir=os.getcwd()) as code_executor:
@@ -128,6 +123,8 @@ async def run_task(task: str, hil_mode: bool, use_rich_console: bool) -> None:
             else:
                 await Console(m1.run_stream(task=task), output_stats=False, user_input_manager=input_manager)
 
+        await client.close()
+
     task = args.task if isinstance(args.task, str) else args.task[0]
     asyncio.run(run_task(task, not args.no_hil, args.rich))