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 3d157d240b88..84f008d7990e 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 @@ -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: ... diff --git a/python/packages/autogen-core/tests/test_tool_agent.py b/python/packages/autogen-core/tests/test_tool_agent.py index 5f366420880b..715d5519c08a 100644 --- a/python/packages/autogen-core/tests/test_tool_agent.py +++ b/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) diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/playwright_controller.py b/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/playwright_controller.py index b99fbc49afad..37e1cdda4bdc 100644 --- a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/playwright_controller.py +++ b/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 diff --git a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py b/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py index 3e722e103c22..34dd1616de72 100644 --- a/python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/utils/chat_completion_client_recorder.py +++ b/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() diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py index 3829d8dcfcb0..2ed797abd784 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py +++ b/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. 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 813a6043a0f5..1b22abe429b7 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 @@ -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 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 bf219dc19672..59a2538b942d 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 @@ -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() diff --git a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py index 3fad35d2ab9a..761863eeb5df 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py +++ b/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 diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index a309494ffaed..34279acc2515 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/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 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 ddcf99f82271..b09e13b4b280 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 @@ -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 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 f4d480606249..aa58830a9087 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 @@ -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) diff --git a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py b/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py index 8339b6516285..1b159da4e91d 100644 --- a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py +++ b/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))