AutoGen cache using Azure Cosmos DB

autogen/cache/cache.py

@@ -26,7 +26,14 @@ class Cache(AbstractCache):
         cache: The cache instance created based on the provided configuration.
     """
 
-    ALLOWED_CONFIG_KEYS = ["cache_seed", "redis_url", "cache_path_root"]
+    ALLOWED_CONFIG_KEYS = [
+        "cache_seed",
+        "redis_url",
+        "cache_path_root",
+        "connection_string",
+        "database_id",
+        "container_id",
+    ]
 
     @staticmethod
     def redis(cache_seed: Union[str, int] = 42, redis_url: str = "redis://localhost:6379/0") -> "Cache":
@@ -56,6 +63,33 @@ def disk(cache_seed: Union[str, int] = 42, cache_path_root: str = ".cache") -> "
         """
         return Cache({"cache_seed": cache_seed, "cache_path_root": cache_path_root})
 
+    @staticmethod
+    def cosmos_db(
+        connection_string: Optional[str] = None,
+        container_id: Optional[str] = None,
+        cache_seed: Union[str, int] = 42,
+    ) -> "Cache":
+        """
+        Create a Cosmos DB cache instance with 'autogen_cache' as database ID.
+
+        Args:
+            connection_string (str, optional): Connection string to the Cosmos DB account. Defaults to None.
+            database_id (str, optional): The database ID for the Cosmos DB account. Defaults to None.
+            container_id (str, optional): The container ID for the Cosmos DB account. Defaults to None.
+            cache_seed (Union[str, int], optional): A seed for the cache. Defaults to 42.
+
+        Returns:
+            Cache: A Cache instance configured for Cosmos DB.
+        """
+        return Cache(
+            {
+                "cache_seed": cache_seed,
+                "connection_string": connection_string,
+                "database_id": "autogen_cache",
+                "container_id": container_id,
+            }
+        )
+
     def __init__(self, config: Dict[str, Any]):
         """
         Initialize the Cache with the given configuration.
@@ -78,6 +112,9 @@ def __init__(self, config: Dict[str, Any]):
             self.config.get("cache_seed", "42"),
             self.config.get("redis_url", None),
             self.config.get("cache_path_root", None),
+            self.config.get("connection_string", None),
+            self.config.get("database_id", None),
+            self.config.get("container_id", None),
         )
 
     def __enter__(self) -> "Cache":

autogen/cache/cache_factory.py

@@ -8,21 +8,31 @@
 class CacheFactory:
     @staticmethod
     def cache_factory(
-        seed: Union[str, int], redis_url: Optional[str] = None, cache_path_root: str = ".cache"
+        seed: Union[str, int],
+        redis_url: Optional[str] = None,
+        cache_path_root: str = ".cache",
+        connection_string: Optional[str] = None,
+        database_id: str = "autogen_cache",
+        container_id: Optional[str] = None,
     ) -> AbstractCache:
         """
         Factory function for creating cache instances.
 
-        Based on the provided redis_url, this function decides whether to create a RedisCache
-        or DiskCache instance. If RedisCache is available and redis_url is provided,
-        a RedisCache instance is created. Otherwise, a DiskCache instance is used.
+        This function decides whether to create a RedisCache, DiskCache, or CosmosDBCache instance
+        based on the provided parameters. If RedisCache is available and a redis_url is provided,
+        a RedisCache instance is created. If connection_string, database_id, and container_id
+        are provided, a CosmosDBCache is created. Otherwise, a DiskCache instance is used.
 
         Args:
             seed (Union[str, int]): A string or int used as a seed or namespace for the cache.
                         This could be useful for creating distinct cache instances
                         or for namespacing keys in the cache.
-            redis_url (str or None): The URL for the Redis server. If this is None
-                                     or if RedisCache is not available, a DiskCache instance is created.
+            redis_url (Optional[str]): The URL for the Redis server. If this is provided and Redis is available,
+                                       a RedisCache instance is created.
+            cache_path_root (str): The root path for the disk cache. Defaults to ".cache".
+            connection_string (Optional[str]): The connection string for the Cosmos DB account. Required for creating a CosmosDBCache.
+            database_id (Optional[str]): The database ID for the Cosmos DB account. Required for creating a CosmosDBCache.
+            container_id (Optional[str]): The container ID for the Cosmos DB account. Required for creating a CosmosDBCache.
 
         Returns:
             An instance of either RedisCache or DiskCache, depending on the availability of RedisCache
@@ -40,14 +50,28 @@ def cache_factory(
         ```python
         disk_cache = cache_factory("myseed", None)
         ```
+
+        Creating a Cosmos DB cache:
+        ```python
+        cosmos_cache = cache_factory("myseed", connection_string="your_connection_string", database_id="your_database_id", container_id="your_container_id")
+        ```
+
         """
-        if redis_url is not None:
+        if redis_url:
             try:
                 from .redis_cache import RedisCache
 
                 return RedisCache(seed, redis_url)
             except ImportError:
-                logging.warning("RedisCache is not available. Creating a DiskCache instance instead.")
-                return DiskCache(f"./{cache_path_root}/{seed}")
-        else:
-            return DiskCache(f"./{cache_path_root}/{seed}")
+                logging.warning("RedisCache is not available. Fallback to DiskCache.")
+
+        if connection_string and database_id and container_id:
+            try:
+                from .cosmos_db_cache import CosmosDBCache
+
+                return CosmosDBCache(seed, connection_string, database_id, container_id)
+            except ImportError:
+                logging.warning("CosmosDBCache is not available. Fallback to DiskCache.")
+
+        # Default to DiskCache if neither Redis nor Cosmos DB configurations are provided
+        return DiskCache(f"./{cache_path_root}/{seed}")

autogen/cache/cosmos_db_cache.py

@@ -0,0 +1,115 @@
+# Install Azure Cosmos DB SDK if not already
+# pip install azure-cosmos
+
+import pickle
+from typing import Any, Optional, Union
+
+from azure.cosmos import CosmosClient, PartitionKey, exceptions
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+from autogen.cache.abstract_cache_base import AbstractCache
+
+
+class CosmosDBCache(AbstractCache):
+    """
+    Synchronous implementation of AbstractCache using Azure Cosmos DB NoSQL API.
+
+    This class provides a concrete implementation of the AbstractCache
+    interface using Azure Cosmos DB for caching data, with synchronous operations.
+
+    Attributes:
+        seed (Union[str, int]): A seed or namespace used as a partition key.
+        client (CosmosClient): The Cosmos DB client used for caching.
+        container: The container instance used for caching.
+    """
+
+    def __init__(
+        self, seed: Union[str, int], client: CosmosClient, container_id: str, database_id: str = "autogen_cache"
+    ):
+        """
+        Initialize the CosmosDBCache instance.
+
+        Args:
+            seed (Union[str, int]): A seed or namespace for the cache, used as a partition key.
+            connection_string (str): The connection string for the Cosmos DB account.
+            container_id (str): The container ID to be used for caching.
+            client (Optional[CosmosClient]): An existing CosmosClient instance to be used for caching.
+        """
+        self.seed = seed
+        self.client = client
+        self.database = self.client.get_database_client(database_id)
+        self.container = self.database.get_container_client(container_id)
+        if not self.container.exists():
+            self.database.create_container(id=container_id, partition_key=PartitionKey(path="/partitionKey"))
+
+    @classmethod
+    def from_connection_string(cls, seed: Union[str, int], connection_string: str, database_id: str, container_id: str):
+        client = CosmosClient.from_connection_string(connection_string)
+        return cls(seed, client, database_id, container_id)
+
+    @classmethod
+    def from_existing_client(cls, seed: Union[str, int], client: CosmosClient, database_id: str, container_id: str):
+        return cls(seed, client, database_id, container_id)
+
+    def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
+        """
+        Retrieve an item from the Cosmos DB cache.
+
+        Args:
+            key (str): The key identifying the item in the cache.
+            default (optional): The default value to return if the key is not found.
+
+        Returns:
+            The deserialized value associated with the key if found, else the default value.
+        """
+        try:
+            response = self.container.read_item(item=key, partition_key=str(self.seed))
+            return pickle.loads(response["data"])
+        except CosmosResourceNotFoundError:
+            return default
+        except Exception as e:
+            # Log the exception or rethrow after logging if needed
+            # Consider logging or handling the error appropriately here
+            raise e
+
+    def set(self, key: str, value: Any) -> None:
+        """
+        Set an item in the Cosmos DB cache.
+
+        Args:
+            key (str): The key under which the item is to be stored.
+            value: The value to be stored in the cache.
+
+        Notes:
+            The value is serialized using pickle before being stored.
+        """
+        serialized_value = pickle.dumps(value)
+        item = {"id": key, "partitionKey": str(self.seed), "data": serialized_value}
+        self.container.upsert_item(item)
+
+    def close(self) -> None:
+        """
+        Close the Cosmos DB client.
+
+        Perform any necessary cleanup, such as closing network connections.
+        """
+        # CosmosClient doesn't require explicit close in the current SDK
+        # If you created the client inside this class, you should close it if necessary
+        pass
+
+    async def __enter__(self):
+        """
+        Context management entry.
+
+        Returns:
+            self: The instance itself.
+        """
+        return self
+
+    def __exit__(self, exc_type: Optional[type], exc_value: Optional[Exception], traceback: Optional[Any]) -> None:
+        """
+        Context management exit.
+
+        Perform cleanup actions such as closing the Cosmos DB client.
+        """
+        self.close()

test/cache/test_cache.py

@@ -8,13 +8,27 @@
 
 class TestCache(unittest.TestCase):
     def setUp(self):
-        self.config = {"cache_seed": "test_seed", "redis_url": "redis://test", "cache_path_root": ".test_cache"}
+        self.config = {
+            "cache_seed": "test_seed",
+            "redis_url": "redis://test",
+            "cache_path_root": ".test_cache",
+            "connection_string": None,
+            "database_id": None,
+            "container_id": None,
+        }
 
     @patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=MagicMock())
     def test_init(self, mock_cache_factory):
         cache = Cache(self.config)
         self.assertIsInstance(cache.cache, MagicMock)
-        mock_cache_factory.assert_called_with("test_seed", "redis://test", ".test_cache")
+        mock_cache_factory.assert_called_with(
+            "test_seed",
+            "redis://test",
+            ".test_cache",
+            None,  # connection_string
+            None,  # database_id
+            None,  # container_id
+        )
 
     @patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=MagicMock())
     def test_context_manager(self, mock_cache_factory):

test/cache/test_cosmos_db_cache.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3 -m pytest
+
+import pickle
+import unittest
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+try:
+    from autogen.cache.cosmos_db_cache import CosmosDBCache
+
+    skip_cosmos_tests = False
+except ImportError:
+    skip_cosmos_tests = True
+
+
+class TestCosmosDBCache(unittest.TestCase):
+    def setUp(self):
+        self.seed = "test_seed"
+        self.client = MagicMock()
+        self.database_id = "test_database"
+        self.container_id = "test_container"
+
+    @pytest.mark.skipif(skip_cosmos_tests, reason="Cosmos DB SDK not installed")
+    @patch("autogen.cache.cosmos_db_cache.CosmosClient.from_connection_string", return_value=MagicMock())
+    def test_init(self, mock_cosmos_from_connection_string):
+        cache = CosmosDBCache(self.seed, self.client, self.database_id, self.container_id)
+        self.assertEqual(cache.seed, self.seed)
+        self.assertEqual(cache.client, self.client)
+
+    @pytest.mark.skipif(skip_cosmos_tests, reason="Cosmos DB SDK not installed")
+    def test_get(self):
+        key = "key"
+        value = "value"
+        serialized_value = pickle.dumps(value)
+        self.client.get_database_client().get_container_client().read_item.return_value = {"data": serialized_value}
+        cache = CosmosDBCache(self.seed, self.client, self.database_id, self.container_id)
+
+        self.assertEqual(cache.get(key), value)
+        self.client.get_database_client().get_container_client().read_item.assert_called_with(
+            item=key, partition_key=str(self.seed)
+        )
+
+        self.client.get_database_client().get_container_client().read_item.side_effect = Exception("not found")
+        self.assertIsNone(cache.get(key))
+
+    @pytest.mark.skipif(skip_cosmos_tests, reason="Cosmos DB SDK not installed")
+    def test_set(self):
+        key = "key"
+        value = "value"
+        serialized_value = pickle.dumps(value)
+        cache = CosmosDBCache(self.seed, self.client, self.database_id, self.container_id)
+        cache.set(key, value)
+        self.client.get_database_client().get_container_client().upsert_item.assert_called_with(
+            {"id": key, "partitionKey": str(self.seed), "data": serialized_value}
+        )
+
+    @pytest.mark.skipif(skip_cosmos_tests, reason="Cosmos DB SDK not installed")
+    def test_context_manager(self):
+        with CosmosDBCache(self.seed, self.client, self.database_id, self.container_id) as cache:
+            self.assertIsInstance(cache, CosmosDBCache)
+        # Note: Testing of close method if needed when CosmosClient requires explicit closing logic
+
+
+if __name__ == "__main__":
+    unittest.main()

website/docs/topics/llm-caching.md

@@ -3,8 +3,7 @@
 AutoGen supports caching API requests so that they can be reused when the same request is issued. This is useful when repeating or continuing experiments for reproducibility and cost saving.
 
 Since version [`0.2.8`](https://github.com/microsoft/autogen/releases/tag/v0.2.8), a configurable context manager allows you to easily
-configure LLM cache, using either [`DiskCache`](/docs/reference/cache/disk_cache#diskcache) or [`RedisCache`](/docs/reference/cache/redis_cache#rediscache). All agents inside the
-context manager will use the same cache.
+configure LLM cache, using either [`DiskCache`](/docs/reference/cache/disk_cache#diskcache), [`RedisCache`](/docs/reference/cache/redis_cache#rediscache), or Cosmos DB Cache. All agents inside the context manager will use the same cache.
 
 ```python
 from autogen import Cache
@@ -16,6 +15,11 @@ with Cache.redis(redis_url="redis://localhost:6379/0") as cache:
 # Use DiskCache as cache
 with Cache.disk() as cache:
     user.initiate_chat(assistant, message=coding_task, cache=cache)
+
+# Use Azure Cosmos DB as cache
+with Cache.cosmos_db(connection_string="your_connection_string", database_id="your_database_id", container_id="your_container_id") as cache:
+    user.initiate_chat(assistant, message=coding_task, cache=cache)
+
 ```
 
 The cache can also be passed directly to the model client's create call.