diff --git a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py index b43e061350e..12cee716f23 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py @@ -1,8 +1,8 @@ import asyncio import functools +import warnings from textwrap import dedent from typing import Any, Callable, Sequence -import warnings from pydantic import BaseModel from typing_extensions import Self 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 79ed5f1660a..15ea1cfd240 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 @@ -3,7 +3,7 @@ import warnings from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union, cast -from autogen_core import CacheStore, CancellationToken +from autogen_core import CacheStore, CancellationToken, InMemoryStore from autogen_core.models import ( ChatCompletionClient, CreateResult, @@ -74,11 +74,16 @@ async def main(): client (ChatCompletionClient): The original ChatCompletionClient to wrap. store (CacheStore): A store object that implements get and set methods. The user is responsible for managing the store's lifecycle & clearing it (if needed). + Defaults to using in-memory cache. """ - def __init__(self, client: ChatCompletionClient, store: CacheStore[CHAT_CACHE_VALUE_TYPE]): + def __init__( + self, + client: ChatCompletionClient, + store: Optional[CacheStore[CHAT_CACHE_VALUE_TYPE]] = None, + ): self.client = client - self.store = store + self.store = store or InMemoryStore[CHAT_CACHE_VALUE_TYPE]() def _check_cache( self, diff --git a/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py b/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py index ceb4d9a9f72..5a023b02dff 100644 --- a/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py +++ b/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py @@ -2,7 +2,6 @@ from typing import List, Tuple, Union import pytest -from autogen_core import InMemoryStore from autogen_core.models import ( ChatCompletionClient, CreateResult, @@ -10,7 +9,7 @@ SystemMessage, UserMessage, ) -from autogen_ext.models.cache import CHAT_CACHE_VALUE_TYPE, ChatCompletionCache +from autogen_ext.models.cache import ChatCompletionCache from autogen_ext.models.replay import ReplayChatCompletionClient @@ -21,8 +20,7 @@ def get_test_data() -> Tuple[list[str], list[str], SystemMessage, ChatCompletion system_prompt = SystemMessage(content="This is a system prompt") replay_client = ReplayChatCompletionClient(responses) replay_client.set_cached_bool_value(False) - store = InMemoryStore[CHAT_CACHE_VALUE_TYPE]() - cached_client = ChatCompletionCache(replay_client, store) + cached_client = ChatCompletionCache(replay_client) return responses, prompts, system_prompt, replay_client, cached_client