Improve interface for cachestore

python/packages/autogen-core/docs/src/reference/index.md

@@ -48,7 +48,6 @@ python/autogen_ext.agents.video_surfer.tools
 python/autogen_ext.auth.azure
 python/autogen_ext.teams.magentic_one
 python/autogen_ext.models.openai
-python/autogen_ext.models.replay
 python/autogen_ext.tools.langchain
 python/autogen_ext.tools.code_execution
 python/autogen_ext.code_executors.local

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/model-clients.ipynb

@@ -200,19 +200,19 @@
     "from typing import Any, Dict, Optional\n",
     "\n",
     "from autogen_core import CacheStore\n",
-    "from autogen_core.models import ChatCompletionCache\n",
+    "from autogen_core.models import CHAT_CACHE_VALUE_TYPE, ChatCompletionCache\n",
     "\n",
     "\n",
     "# Simple CacheStore implementation using in-memory dict,\n",
     "# you can also use redis.Redis or diskcache.Cache\n",
-    "class DictStore(CacheStore):\n",
+    "class DictStore(CacheStore[CHAT_CACHE_VALUE_TYPE]):\n",
     "    def __init__(self) -> None:\n",
-    "        self._store: Dict[Any, Any] = {}\n",
+    "        self._store: dict[str, CHAT_CACHE_VALUE_TYPE] = {}\n",
     "\n",
-    "    def get(self, key: Any, default: Optional[Any] = None) -> Optional[Any]:\n",
+    "    def get(self, key: str, default: Optional[CHAT_CACHE_VALUE_TYPE] = None) -> Optional[CHAT_CACHE_VALUE_TYPE]:\n",
     "        return self._store.get(key, default)\n",
     "\n",
-    "    def set(self, key: Any, value: Any) -> None:\n",
+    "    def set(self, key: str, value: CHAT_CACHE_VALUE_TYPE) -> None:\n",
     "        self._store[key] = value\n",
     "\n",
     "\n",

python/packages/autogen-core/src/autogen_core/_cache_store.py

@@ -1,15 +1,17 @@
-from typing import Any, Optional, Protocol
+from typing import Generic, Optional, Protocol, TypeVar
 
+T = TypeVar("T")
 
-class CacheStore(Protocol):
+
+class CacheStore(Protocol, Generic[T]):
     """
     This protocol defines the basic interface for store/cache operations.
 
     Allows duck-typing with any object that implements the get and set methods,
     such as redis or diskcache interfaces.
     """
 
-    def get(self, key: Any, default: Optional[Any] = None) -> Optional[Any]:
+    def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
         """
         Retrieve an item from the store.
 
@@ -23,7 +25,7 @@ def get(self, key: Any, default: Optional[Any] = None) -> Optional[Any]:
         """
         ...
 
-    def set(self, key: Any, value: Any) -> None:
+    def set(self, key: str, value: T) -> None:
         """
         Set an item in the store.
 

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

@@ -1,4 +1,4 @@
-from ._cache import ChatCompletionCache
+from ._cache import CHAT_CACHE_VALUE_TYPE, ChatCompletionCache
 from ._model_client import ChatCompletionClient, ModelCapabilities, ModelFamily, ModelInfo  # type: ignore
 from ._replay_chat_completion_client import ReplayChatCompletionClient
 from ._types import (
@@ -17,6 +17,7 @@
 
 __all__ = [
     "ModelCapabilities",
+    "CHAT_CACHE_VALUE_TYPE",
     "ChatCompletionCache",
     "ChatCompletionClient",
     "SystemMessage",

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

@@ -17,14 +17,16 @@
     RequestUsage,
 )
 
+CHAT_CACHE_VALUE_TYPE = Union[CreateResult, List[Union[str, CreateResult]]]
+
 
 class ChatCompletionCache(ChatCompletionClient):
     """
     A wrapper around a ChatCompletionClient that caches creation results from an underlying client.
     Cache hits do not contribute to token usage of the original client.
     """
 
-    def __init__(self, client: ChatCompletionClient, store: CacheStore):
+    def __init__(self, client: ChatCompletionClient, store: CacheStore[CHAT_CACHE_VALUE_TYPE]):
         """
         Initialize a new ChatCompletionCache.
 

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

@@ -5,7 +5,7 @@
 
 def test_set_and_get_object_key_value() -> None:
     mock_store = Mock(spec=CacheStore)
-    test_key = object()
+    test_key = "test_key"
     test_value = object()
     mock_store.set(test_key, test_value)
     mock_store.get.return_value = test_value

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

@@ -1,9 +1,10 @@
 import copy
-from typing import Any, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import pytest
 from autogen_core import CacheStore
 from autogen_core.models import (
+    CHAT_CACHE_VALUE_TYPE,
     ChatCompletionCache,
     ChatCompletionClient,
     CreateResult,
@@ -14,14 +15,14 @@
 )
 
 
-class DictStore(CacheStore):
+class DictStore(CacheStore[CHAT_CACHE_VALUE_TYPE]):
     def __init__(self) -> None:
-        self._store: dict[Any, Any] = {}
+        self._store: dict[str, CHAT_CACHE_VALUE_TYPE] = {}
 
-    def get(self, key: Any, default: Optional[Any] = None) -> Optional[Any]:
+    def get(self, key: str, default: Optional[CHAT_CACHE_VALUE_TYPE] = None) -> Optional[CHAT_CACHE_VALUE_TYPE]:
         return self._store.get(key, default)
 
-    def set(self, key: Any, value: Any) -> None:
+    def set(self, key: str, value: CHAT_CACHE_VALUE_TYPE) -> None:
         self._store[key] = value