From d60c91c7524aa328ccebcd7ff2b86340263e3cb8 Mon Sep 17 00:00:00 2001 From: Daria Korenieva Date: Tue, 12 May 2026 05:34:43 -0700 Subject: [PATCH] feat(ext): add Valkey cache integration (#3057) Signed-off-by: Daria Korenieva --- .../manager_modelcacheparameters_152367.mdx | 2 +- .../manager_modelcacheparameters_152367.mdx | 2 +- .../src/dbgpt/storage/cache/manager.py | 21 +- .../src/dbgpt_ext/storage/cache/__init__.py | 1 + .../dbgpt_ext/storage/cache/tests/__init__.py | 1 + .../storage/cache/tests/test_valkey_cache.py | 458 ++++++++++++++++++ .../tests/test_valkey_cache_integration.py | 227 +++++++++ .../dbgpt_ext/storage/cache/valkey_cache.py | 327 +++++++++++++ 8 files changed, 1036 insertions(+), 3 deletions(-) create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/cache/__init__.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/__init__.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache_integration.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/cache/valkey_cache.py diff --git a/docs/docs/config-reference/utils/manager_modelcacheparameters_152367.mdx b/docs/docs/config-reference/utils/manager_modelcacheparameters_152367.mdx index 339ffb473..8a1bb11a7 100644 --- a/docs/docs/config-reference/utils/manager_modelcacheparameters_152367.mdx +++ b/docs/docs/config-reference/utils/manager_modelcacheparameters_152367.mdx @@ -21,7 +21,7 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; "name": "storage_type", "type": "string", "required": false, - "description": "The storage type, default is memory", + "description": "The storage type. Supported values: memory, disk, valkey. For valkey, configure via environment variables VALKEY_HOST, VALKEY_PORT, VALKEY_PASSWORD.", "defaultValue": "memory" }, { diff --git a/docs/i18n/zh-CN/docusaurus-plugin-content-docs/current/config-reference/utils/manager_modelcacheparameters_152367.mdx b/docs/i18n/zh-CN/docusaurus-plugin-content-docs/current/config-reference/utils/manager_modelcacheparameters_152367.mdx index 339ffb473..8a1bb11a7 100644 --- a/docs/i18n/zh-CN/docusaurus-plugin-content-docs/current/config-reference/utils/manager_modelcacheparameters_152367.mdx +++ b/docs/i18n/zh-CN/docusaurus-plugin-content-docs/current/config-reference/utils/manager_modelcacheparameters_152367.mdx @@ -21,7 +21,7 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; "name": "storage_type", "type": "string", "required": false, - "description": "The storage type, default is memory", + "description": "The storage type. Supported values: memory, disk, valkey. For valkey, configure via environment variables VALKEY_HOST, VALKEY_PORT, VALKEY_PASSWORD.", "defaultValue": "memory" }, { diff --git a/packages/dbgpt-core/src/dbgpt/storage/cache/manager.py b/packages/dbgpt-core/src/dbgpt/storage/cache/manager.py index 164ccf73e..e760529e2 100644 --- a/packages/dbgpt-core/src/dbgpt/storage/cache/manager.py +++ b/packages/dbgpt-core/src/dbgpt/storage/cache/manager.py @@ -33,7 +33,11 @@ class ModelCacheParameters(BaseParameters): storage_type: str = field( default="memory", metadata={ - "help": _("The storage type, default is memory"), + "help": _( + "The storage type. Supported values: memory, disk, valkey. " + "For valkey, configure via env vars: VALKEY_HOST, VALKEY_PORT, " + "VALKEY_PASSWORD." + ), }, ) max_memory_mb: int = field( @@ -172,6 +176,21 @@ def initialize_cache( f"message: {str(e)}" ) cache_storage = MemoryCacheStorage(max_memory_mb=max_memory_mb) + elif storage_type == "valkey": + try: + from dbgpt_ext.storage.cache.valkey_cache import ValkeyCacheStorage + + # Configuration is read from environment variables: + # VALKEY_HOST, VALKEY_PORT, VALKEY_PASSWORD + cache_storage = ValkeyCacheStorage() + except ImportError as e: + logger.warning( + "Can't import ValkeyCacheStorage (valkey-glide not installed). " + "Falling back to MemoryCacheStorage — cache will NOT be shared " + "across nodes. Import error: %s", + e, + ) + cache_storage = MemoryCacheStorage(max_memory_mb=max_memory_mb) else: cache_storage = MemoryCacheStorage(max_memory_mb=max_memory_mb) system_app.register( diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/__init__.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/__init__.py new file mode 100644 index 000000000..9d9559b8e --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/__init__.py @@ -0,0 +1 @@ +"""Cache storage extensions.""" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/__init__.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/__init__.py new file mode 100644 index 000000000..e711dac6d --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for cache storage extensions.""" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache.py new file mode 100644 index 000000000..8a7b37a7b --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache.py @@ -0,0 +1,458 @@ +"""Unit tests for ValkeyCacheStorage. + +These tests use mocked valkey-glide client and do not require a running Valkey server. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +from dbgpt.core.interface.cache import RetrievalPolicy +from dbgpt.storage.cache.storage.base import StorageItem +from dbgpt_ext.storage.cache.valkey_cache import ValkeyCacheStorage + + +class MockCacheKey: + """Mock cache key for testing.""" + + def __init__(self, key_str: str = "test_key"): + self._key_str = key_str + + def get_hash_bytes(self) -> bytes: + return self._key_str.encode() + + def serialize(self) -> bytes: + return self._key_str.encode() + + def __hash__(self): + return hash(self._key_str) + + +class MockCacheValue: + """Mock cache value for testing.""" + + def __init__(self, value_str: str = "test_value"): + self._value_str = value_str + + def serialize(self) -> bytes: + return self._value_str.encode() + + +class MockCacheConfig: + """Mock cache config for testing.""" + + def __init__(self, retrieval_policy=RetrievalPolicy.EXACT_MATCH): + self.retrieval_policy = retrieval_policy + + +@pytest.fixture +def mock_client(): + """Create a mock Valkey client.""" + client = AsyncMock() + client.get = AsyncMock(return_value=None) + client.set = AsyncMock(return_value=None) + client.exists = AsyncMock(return_value=0) + return client + + +@pytest.fixture +def valkey_cache(mock_client): + """Create a ValkeyCacheStorage with mocked client.""" + with patch.object(ValkeyCacheStorage, "_create_client", return_value=mock_client): + storage = ValkeyCacheStorage( + host="localhost", + port=6379, + key_prefix="test_cache:", + ttl_seconds=None, + ) + storage._client = mock_client + return storage + + +class TestValkeyCacheStorageInit: + """Tests for initialization.""" + + def test_default_config(self, mock_client): + """Test default configuration values.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage() + assert storage._host == "localhost" + assert storage._port == 6379 + assert storage._key_prefix == "dbgpt_cache:" + assert storage._ttl_seconds is None + + def test_custom_config(self, valkey_cache): + """Test custom configuration.""" + assert valkey_cache._host == "localhost" + assert valkey_cache._port == 6379 + assert valkey_cache._key_prefix == "test_cache:" + + def test_support_async(self, valkey_cache): + """Test that async is supported.""" + assert valkey_cache.support_async() is True + + +class TestValkeyCacheStorageCheckConfig: + """Tests for check_config.""" + + def test_valid_config(self, valkey_cache): + """Test valid config passes.""" + config = MockCacheConfig(RetrievalPolicy.EXACT_MATCH) + assert valkey_cache.check_config(config) is True + + def test_none_config(self, valkey_cache): + """Test None config passes.""" + assert valkey_cache.check_config(None) is True + + def test_invalid_policy_raises(self, valkey_cache): + """Test invalid retrieval policy raises ValueError.""" + config = MockCacheConfig(retrieval_policy="SIMILARITY") + with pytest.raises(ValueError, match="EXACT_MATCH"): + valkey_cache.check_config(config) + + def test_invalid_policy_no_raise(self, valkey_cache): + """Test invalid retrieval policy returns False when raise_error=False.""" + config = MockCacheConfig(retrieval_policy="SIMILARITY") + assert valkey_cache.check_config(config, raise_error=False) is False + + +class TestValkeyCacheStorageGet: + """Tests for get operations.""" + + def test_get_miss(self, valkey_cache, mock_client): + """Test cache miss returns None.""" + mock_client.get = AsyncMock(return_value=None) + key = MockCacheKey("missing_key") + result = valkey_cache.get(key) + assert result is None + + def test_get_hit(self, valkey_cache, mock_client): + """Test cache hit returns StorageItem.""" + key = MockCacheKey("hit_key") + value = MockCacheValue("hit_value") + item = StorageItem.build_from_kv(key, value) + serialized = item.serialize() + + mock_client.get = AsyncMock(return_value=serialized) + result = valkey_cache.get(key) + + assert result is not None + assert result.key_data == b"hit_key" + assert result.value_data == b"hit_value" + + def test_get_corrupt_data(self, valkey_cache, mock_client): + """Test corrupt data returns None gracefully.""" + mock_client.get = AsyncMock(return_value=b"not_valid_msgpack") + key = MockCacheKey("corrupt") + result = valkey_cache.get(key) + assert result is None + + +class TestValkeyCacheStorageSet: + """Tests for set operations.""" + + def test_set_without_ttl(self, valkey_cache, mock_client): + """Test setting a value without TTL.""" + key = MockCacheKey("set_key") + value = MockCacheValue("set_value") + + valkey_cache.set(key, value) + + mock_client.set.assert_called_once() + call_args = mock_client.set.call_args[0] + assert call_args[0] == "test_cache:" + b"set_key".hex() + assert isinstance(call_args[1], bytes) + + def test_set_with_ttl(self, mock_client): + """Test setting a value with TTL.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", + port=6379, + key_prefix="test_cache:", + ttl_seconds=3600, + ) + storage._client = mock_client + + key = MockCacheKey("ttl_key") + value = MockCacheValue("ttl_value") + + storage.set(key, value) + mock_client.set.assert_called_once() + call_kwargs = mock_client.set.call_args[1] + assert "expiry" in call_kwargs + + +class TestValkeyCacheStorageExists: + """Tests for exists operation.""" + + def test_exists_true(self, valkey_cache, mock_client): + """Test exists returns True when key is present.""" + mock_client.exists = AsyncMock(return_value=1) + key = MockCacheKey("existing") + assert valkey_cache.exists(key) is True + + def test_exists_false(self, valkey_cache, mock_client): + """Test exists returns False when key is absent.""" + mock_client.exists = AsyncMock(return_value=0) + key = MockCacheKey("missing") + assert valkey_cache.exists(key) is False + + +class TestValkeyCacheStorageKeyGeneration: + """Tests for key generation.""" + + def test_make_key(self, valkey_cache): + """Test key generation includes prefix and hex hash.""" + key = MockCacheKey("my_key") + result = valkey_cache._make_key(key) + assert result == "test_cache:" + b"my_key".hex() + + def test_different_keys_produce_different_valkey_keys(self, valkey_cache): + """Test that different cache keys produce different Valkey keys.""" + key1 = MockCacheKey("key_a") + key2 = MockCacheKey("key_b") + assert valkey_cache._make_key(key1) != valkey_cache._make_key(key2) + + +class TestValkeyCacheStorageAsync: + """Tests for async operations.""" + + @pytest.mark.asyncio + async def test_aget_miss(self, valkey_cache, mock_client): + """Test async cache miss returns None.""" + mock_client.get = AsyncMock(return_value=None) + valkey_cache._async_client = mock_client + valkey_cache._async_loop = asyncio.get_running_loop() + key = MockCacheKey("async_missing") + result = await valkey_cache.aget(key) + assert result is None + + @pytest.mark.asyncio + async def test_aget_hit(self, valkey_cache, mock_client): + """Test async cache hit returns StorageItem.""" + key = MockCacheKey("async_hit") + value = MockCacheValue("async_value") + item = StorageItem.build_from_kv(key, value) + serialized = item.serialize() + + mock_client.get = AsyncMock(return_value=serialized) + valkey_cache._async_client = mock_client + valkey_cache._async_loop = asyncio.get_running_loop() + result = await valkey_cache.aget(key) + + assert result is not None + assert result.key_data == b"async_hit" + assert result.value_data == b"async_value" + + @pytest.mark.asyncio + async def test_aget_corrupt_data(self, valkey_cache, mock_client): + """Test async corrupt data returns None gracefully.""" + mock_client.get = AsyncMock(return_value=b"not_valid_msgpack") + valkey_cache._async_client = mock_client + valkey_cache._async_loop = asyncio.get_running_loop() + key = MockCacheKey("async_corrupt") + result = await valkey_cache.aget(key) + assert result is None + + @pytest.mark.asyncio + async def test_aset_without_ttl(self, valkey_cache, mock_client): + """Test async setting a value without TTL.""" + valkey_cache._async_client = mock_client + valkey_cache._async_loop = asyncio.get_running_loop() + key = MockCacheKey("async_set_key") + value = MockCacheValue("async_set_value") + + await valkey_cache.aset(key, value) + + mock_client.set.assert_called_once() + call_args = mock_client.set.call_args[0] + assert call_args[0] == "test_cache:" + b"async_set_key".hex() + assert isinstance(call_args[1], bytes) + + @pytest.mark.asyncio + async def test_aset_with_ttl(self, mock_client): + """Test async setting a value with TTL.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", + port=6379, + key_prefix="test_cache:", + ttl_seconds=3600, + ) + storage._client = mock_client + storage._async_client = mock_client + storage._async_loop = asyncio.get_running_loop() + + key = MockCacheKey("async_ttl_key") + value = MockCacheValue("async_ttl_value") + + await storage.aset(key, value) + mock_client.set.assert_called_once() + call_kwargs = mock_client.set.call_args[1] + assert "expiry" in call_kwargs + + +class TestValkeyCacheStorageClose: + """Tests for close/cleanup.""" + + def test_close_releases_resources(self, valkey_cache, mock_client): + """Test close releases client and event loop.""" + mock_client.close = AsyncMock() + # Access client to ensure it's set + _ = valkey_cache.client + valkey_cache.close() + + assert valkey_cache._client is None + assert valkey_cache._loop.is_closed() + + def test_close_idempotent(self, valkey_cache, mock_client): + """Test calling close multiple times is safe.""" + mock_client.close = AsyncMock() + valkey_cache.close() + # Second call should not raise + valkey_cache.close() + + def test_close_releases_async_client(self, valkey_cache, mock_client): + """Test close properly closes _async_client (not just sets to None).""" + async_client = AsyncMock() + async_client.close = AsyncMock() + valkey_cache._async_client = async_client + + valkey_cache.close() + + async_client.close.assert_called_once() + assert valkey_cache._async_client is None + + +class TestValkeyCacheStorageTTLZero: + """Tests for TTL=0 edge case (falsy but valid).""" + + def test_set_with_ttl_zero(self, mock_client): + """Test that ttl_seconds=0 still applies TTL (not skipped as falsy).""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", + port=6379, + key_prefix="test_cache:", + ttl_seconds=0, + ) + storage._client = mock_client + + key = MockCacheKey("zero_ttl_key") + value = MockCacheValue("zero_ttl_value") + + storage.set(key, value) + mock_client.set.assert_called_once() + call_kwargs = mock_client.set.call_args[1] + assert "expiry" in call_kwargs + + @pytest.mark.asyncio + async def test_aset_with_ttl_zero(self, mock_client): + """Test that async ttl_seconds=0 still applies TTL.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", + port=6379, + key_prefix="test_cache:", + ttl_seconds=0, + ) + storage._client = mock_client + storage._async_client = mock_client + storage._async_loop = asyncio.get_running_loop() + + key = MockCacheKey("async_zero_ttl_key") + value = MockCacheValue("async_zero_ttl_value") + + await storage.aset(key, value) + mock_client.set.assert_called_once() + call_kwargs = mock_client.set.call_args[1] + assert "expiry" in call_kwargs + + +class TestValkeyCacheStorageRequestTimeout: + """Tests for request_timeout configuration.""" + + def test_default_request_timeout(self, mock_client): + """Test default request_timeout is 5000ms.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage(host="localhost", port=6379) + assert storage._request_timeout == 5000 + + def test_custom_request_timeout(self, mock_client): + """Test custom request_timeout is stored.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", port=6379, request_timeout=10000 + ) + assert storage._request_timeout == 10000 + + def test_none_request_timeout(self, mock_client): + """Test None request_timeout disables timeout.""" + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage( + host="localhost", port=6379, request_timeout=None + ) + assert storage._request_timeout is None + + +class TestValkeyCacheStorageContextManager: + """Tests for context manager support.""" + + def test_sync_context_manager(self, mock_client): + """Test usage as a sync context manager.""" + mock_client.close = AsyncMock() + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + with ValkeyCacheStorage(host="localhost", port=6379) as storage: + storage._client = mock_client + assert storage is not None + # After exiting, resources should be cleaned up + assert storage._client is None + + def test_client_access_after_close_raises(self, mock_client): + """Test that accessing client after close raises RuntimeError.""" + mock_client.close = AsyncMock() + with patch.object( + ValkeyCacheStorage, "_create_client", return_value=mock_client + ): + storage = ValkeyCacheStorage(host="localhost", port=6379) + storage._client = mock_client + storage.close() + with pytest.raises(RuntimeError, match="has been closed"): + _ = storage.client + + +class TestValkeyCacheStorageEventLoopSafety: + """Tests for event loop safety in async client.""" + + @pytest.mark.asyncio + async def test_async_client_tracks_loop(self, valkey_cache, mock_client): + """Test that _get_async_client tracks the current event loop.""" + valkey_cache._async_client = None + mock_create = AsyncMock(return_value=mock_client) + valkey_cache._create_client_async = mock_create + + client = await valkey_cache._get_async_client() + assert client is mock_client + assert valkey_cache._async_loop == asyncio.get_running_loop() diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache_integration.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache_integration.py new file mode 100644 index 000000000..75d5d8933 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/tests/test_valkey_cache_integration.py @@ -0,0 +1,227 @@ +"""Integration tests for ValkeyCacheStorage. + +These tests require a running Valkey server. +Skip if Valkey is not available. + +To run locally:: + + docker run -d --name valkey -p 6379:6379 valkey/valkey:latest + + pytest -v -k test_valkey_cache_integration +""" + +from __future__ import annotations + +import os +import uuid + +import pytest + +pytestmark = pytest.mark.integration + +VALKEY_HOST = os.environ.get("VALKEY_HOST", "localhost") +VALKEY_PORT = int(os.environ.get("VALKEY_PORT", "6379")) +VALKEY_PASSWORD = os.environ.get("VALKEY_PASSWORD", None) + + +def _valkey_available() -> bool: + """Check if Valkey is available.""" + try: + import asyncio + + from glide import GlideClient, GlideClientConfiguration, NodeAddress + + async def _check(): + config = GlideClientConfiguration( + addresses=[NodeAddress(host=VALKEY_HOST, port=VALKEY_PORT)] + ) + client = await GlideClient.create(config) + await client.ping() + await client.close() + return True + + return asyncio.run(_check()) + except Exception: + return False + + +if not _valkey_available(): + pytest.skip("Valkey server not available", allow_module_level=True) + +from dbgpt_ext.storage.cache.valkey_cache import ValkeyCacheStorage # noqa: E402 + + +class MockCacheKey: + """Mock cache key for integration testing.""" + + def __init__(self, key_str: str): + self._key_str = key_str + + def get_hash_bytes(self) -> bytes: + return self._key_str.encode() + + def serialize(self) -> bytes: + return f"key_data:{self._key_str}".encode() + + def __hash__(self): + return hash(self._key_str) + + +class MockCacheValue: + """Mock cache value for integration testing.""" + + def __init__(self, value_str: str): + self._value_str = value_str + + def serialize(self) -> bytes: + return f"value_data:{self._value_str}".encode() + + +@pytest.fixture +def storage(): + """Create a ValkeyCacheStorage for integration testing.""" + prefix = f"inttest_cache_{uuid.uuid4().hex[:8]}:" + s = ValkeyCacheStorage( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + key_prefix=prefix, + ttl_seconds=60, + ) + yield s + # Cleanup + try: + cursor = "0" + while True: + result = s._run_async( + s.client.custom_command( + ["SCAN", cursor, "MATCH", f"{prefix}*", "COUNT", "100"] + ) + ) + if isinstance(result, (list, tuple)) and len(result) == 2: + cursor = result[0] + if isinstance(cursor, bytes): + cursor = cursor.decode() + keys = result[1] + if keys: + key_list = [k.decode() if isinstance(k, bytes) else k for k in keys] + s._run_async(s.client.delete(key_list)) + else: + break + if cursor == "0": + break + except Exception: + pass + + +class TestValkeyCacheStorageIntegration: + """Integration tests for ValkeyCacheStorage.""" + + def test_set_and_get(self, storage): + """Test basic set and get.""" + key = MockCacheKey("int_test_1") + value = MockCacheValue("hello world") + + storage.set(key, value) + result = storage.get(key) + + assert result is not None + assert result.key_data == b"key_data:int_test_1" + assert result.value_data == b"value_data:hello world" + + def test_get_miss(self, storage): + """Test get on non-existent key returns None.""" + key = MockCacheKey("nonexistent_key") + result = storage.get(key) + assert result is None + + def test_exists(self, storage): + """Test exists check.""" + key = MockCacheKey("exists_test") + value = MockCacheValue("data") + + assert storage.exists(key) is False + storage.set(key, value) + assert storage.exists(key) is True + + def test_overwrite(self, storage): + """Test overwriting an existing key.""" + key = MockCacheKey("overwrite_test") + value1 = MockCacheValue("first") + value2 = MockCacheValue("second") + + storage.set(key, value1) + storage.set(key, value2) + + result = storage.get(key) + assert result is not None + assert result.value_data == b"value_data:second" + + def test_multiple_keys(self, storage): + """Test storing and retrieving multiple keys.""" + keys_values = [ + (MockCacheKey(f"multi_{i}"), MockCacheValue(f"val_{i}")) for i in range(5) + ] + + for key, value in keys_values: + storage.set(key, value) + + for i, (key, _) in enumerate(keys_values): + result = storage.get(key) + assert result is not None + assert result.value_data == f"value_data:val_{i}".encode() + + def test_ttl_is_set(self, storage): + """Test that TTL is applied to keys.""" + key = MockCacheKey("ttl_test") + value = MockCacheValue("expires") + + storage.set(key, value) + + valkey_key = storage._make_key(key) + ttl = storage._run_async(storage.client.ttl(valkey_key)) + assert ttl > 0 + assert ttl <= 60 + + @pytest.mark.asyncio + async def test_aset_and_aget(self, storage): + """Test async set and get.""" + key = MockCacheKey("async_int_test") + value = MockCacheValue("async hello") + + await storage.aset(key, value) + result = await storage.aget(key) + + assert result is not None + assert result.key_data == b"key_data:async_int_test" + assert result.value_data == b"value_data:async hello" + + @pytest.mark.asyncio + async def test_aget_miss(self, storage): + """Test async get on non-existent key returns None.""" + key = MockCacheKey("async_nonexistent") + result = await storage.aget(key) + assert result is None + + @pytest.mark.asyncio + async def test_aset_with_ttl(self, storage): + """Test async set respects TTL.""" + key = MockCacheKey("async_ttl_test") + value = MockCacheValue("async expires") + + await storage.aset(key, value) + + valkey_key = storage._make_key(key) + client = await storage._get_async_client() + ttl = await client.ttl(valkey_key) + assert ttl > 0 + assert ttl <= 60 + + def test_close(self, storage): + """Test close releases resources without error.""" + key = MockCacheKey("close_test") + value = MockCacheValue("data") + storage.set(key, value) + + storage.close() + assert storage._client is None diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/valkey_cache.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/valkey_cache.py new file mode 100644 index 000000000..3eba23e44 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/cache/valkey_cache.py @@ -0,0 +1,327 @@ +"""Valkey cache storage implementation. + +Uses valkey-glide client for high-performance in-memory caching of LLM +responses and embeddings. +""" + +import asyncio +import logging +import os +from typing import Optional + +from dbgpt.core.interface.cache import ( + CacheConfig, + CacheKey, + CacheValue, + K, + RetrievalPolicy, + V, +) +from dbgpt.storage.cache.storage.base import CacheStorage, StorageItem + +logger = logging.getLogger(__name__) + + +class ValkeyCacheStorage(CacheStorage): + """Valkey-based cache storage using valkey-glide client. + + Stores serialized cache items as binary values in Valkey with optional TTL. + Supports both sync and async operations. + """ + + def __init__( + self, + host: Optional[str] = None, + port: Optional[int] = None, + password: Optional[str] = None, + use_ssl: bool = False, + key_prefix: str = "dbgpt_cache:", + ttl_seconds: Optional[int] = None, + request_timeout: Optional[int] = 5000, + ): + """Initialize ValkeyCacheStorage. + + Args: + host: Valkey server host. Defaults to VALKEY_HOST env or localhost. + port: Valkey server port. Defaults to VALKEY_PORT env or 6379. + password: Valkey password. Defaults to VALKEY_PASSWORD env. + use_ssl: Whether to use TLS. + key_prefix: Prefix for all cache keys. + ttl_seconds: Optional TTL in seconds for cache entries. + If None, entries don't expire. + request_timeout: Request timeout in milliseconds. Defaults to 5000ms. + If None, no timeout is set. + """ + try: + import glide # noqa: F401 + except ImportError: + raise ImportError( + "Please install valkey-glide: pip install 'valkey-glide>=2.3.0'" + ) + + self._host = host or os.getenv("VALKEY_HOST", "localhost") + self._port = port or int(os.getenv("VALKEY_PORT", "6379")) + self._password = password or os.getenv("VALKEY_PASSWORD") + self._use_ssl = use_ssl + self._key_prefix = key_prefix + self._ttl_seconds = ttl_seconds + self._request_timeout = request_timeout + self._client = None + self._async_client = None + self._async_loop = None + self._loop = asyncio.new_event_loop() + + @property + def client(self): + """Get or create the Valkey client for sync operations (lazy initialization).""" + if self._loop.is_closed(): + raise RuntimeError("ValkeyCacheStorage has been closed") + if self._client is None: + self._client = self._loop.run_until_complete(self._create_client_async()) + return self._client + + def _get_client_config(self): + """Build the GlideClientConfiguration.""" + from glide import GlideClientConfiguration, NodeAddress + + node = NodeAddress(host=self._host, port=self._port) + + kwargs = { + "addresses": [node], + "use_tls": self._use_ssl, + } + if self._request_timeout is not None: + kwargs["request_timeout"] = self._request_timeout + + if self._password: + from glide import ServerCredentials + + kwargs["credentials"] = ServerCredentials(password=self._password) + + return GlideClientConfiguration(**kwargs) + + async def _create_client_async(self): + """Create a Valkey-glide client on the current event loop.""" + from glide import GlideClient + + return await GlideClient.create(self._get_client_config()) + + def _create_client(self): + """Create a Valkey-glide client (sync helper, used by tests).""" + return self._loop.run_until_complete(self._create_client_async()) + + async def _get_async_client(self): + """Get or create a client bound to the current running event loop. + + If called from a different event loop than self._loop, creates a fresh + client on the current loop to avoid cross-loop issues. + """ + current_loop = asyncio.get_running_loop() + if self._async_client is None or self._async_loop != current_loop: + if self._async_client is not None: + try: + await self._async_client.close() + except Exception: + pass + self._async_client = await self._create_client_async() + self._async_loop = current_loop + return self._async_client + + def _run_async(self, coro): + """Run an async coroutine synchronously using the dedicated event loop.""" + return self._loop.run_until_complete(coro) + + def close(self): + """Close the client connection and event loop.""" + if self._client is not None: + try: + self._loop.run_until_complete(self._client.close()) + except Exception: + pass + self._client = None + if self._async_client is not None: + try: + self._loop.run_until_complete(self._async_client.close()) + except Exception: + pass + self._async_client = None + if self._loop and not self._loop.is_closed(): + self._loop.close() + + def __enter__(self): + """Support usage as a context manager.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Close resources on context manager exit.""" + self.close() + return False + + async def __aenter__(self): + """Support usage as an async context manager.""" + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Close resources on async context manager exit.""" + if self._async_client is not None: + try: + await self._async_client.close() + except Exception: + pass + self._async_client = None + self.close() + return False + + def _make_key(self, key: CacheKey[K]) -> str: + """Build the full Valkey key from a CacheKey.""" + key_hash = key.get_hash_bytes() + return self._key_prefix + key_hash.hex() + + def check_config( + self, + cache_config: Optional[CacheConfig] = None, + raise_error: Optional[bool] = True, + ) -> bool: + """Check whether the CacheConfig is legal.""" + if ( + cache_config + and cache_config.retrieval_policy != RetrievalPolicy.EXACT_MATCH + ): + if raise_error: + raise ValueError( + "ValkeyCacheStorage only supports 'EXACT_MATCH' retrieval policy" + ) + return False + return True + + def support_async(self) -> bool: + """Valkey supports async operations.""" + return True + + def get( + self, key: CacheKey[K], cache_config: Optional[CacheConfig] = None + ) -> Optional[StorageItem]: + """Retrieve a storage item from Valkey. + + Args: + key: The cache key. + cache_config: Optional cache configuration. + + Returns: + The StorageItem if found, None otherwise. + """ + self.check_config(cache_config, raise_error=True) + valkey_key = self._make_key(key) + + data = self._run_async(self.client.get(valkey_key)) + if data is None: + return None + + if isinstance(data, str): + data = data.encode() + + try: + return StorageItem.deserialize(data) + except Exception as e: + logger.warning("Failed to deserialize cache item: %s", e) + return None + + async def aget( + self, key: CacheKey[K], cache_config: Optional[CacheConfig] = None + ) -> Optional[StorageItem]: + """Retrieve a storage item from Valkey asynchronously. + + Args: + key: The cache key. + cache_config: Optional cache configuration. + + Returns: + The StorageItem if found, None otherwise. + """ + self.check_config(cache_config, raise_error=True) + valkey_key = self._make_key(key) + + client = await self._get_async_client() + data = await client.get(valkey_key) + if data is None: + return None + + if isinstance(data, str): + data = data.encode() + + try: + return StorageItem.deserialize(data) + except Exception as e: + logger.warning("Failed to deserialize cache item: %s", e) + return None + + def set( + self, + key: CacheKey[K], + value: CacheValue[V], + cache_config: Optional[CacheConfig] = None, + ) -> None: + """Store a cache item in Valkey. + + Args: + key: The cache key. + value: The cache value. + cache_config: Optional cache configuration. + """ + valkey_key = self._make_key(key) + item = StorageItem.build_from_kv(key, value) + data = item.serialize() + + if self._ttl_seconds is not None: + from glide import ExpirySet, ExpiryType + + expiry = ExpirySet(ExpiryType.SEC, self._ttl_seconds) + self._run_async(self.client.set(valkey_key, data, expiry=expiry)) + else: + self._run_async(self.client.set(valkey_key, data)) + + logger.debug("ValkeyCacheStorage set key hash=%s", valkey_key) + + async def aset( + self, + key: CacheKey[K], + value: CacheValue[V], + cache_config: Optional[CacheConfig] = None, + ) -> None: + """Store a cache item in Valkey asynchronously. + + Args: + key: The cache key. + value: The cache value. + cache_config: Optional cache configuration. + """ + valkey_key = self._make_key(key) + item = StorageItem.build_from_kv(key, value) + data = item.serialize() + + client = await self._get_async_client() + if self._ttl_seconds is not None: + from glide import ExpirySet, ExpiryType + + expiry = ExpirySet(ExpiryType.SEC, self._ttl_seconds) + await client.set(valkey_key, data, expiry=expiry) + else: + await client.set(valkey_key, data) + + logger.debug("ValkeyCacheStorage aset key hash=%s", valkey_key) + + def exists( + self, key: CacheKey[K], cache_config: Optional[CacheConfig] = None + ) -> bool: + """Check if a key exists in the cache. + + Args: + key: The cache key. + cache_config: Optional cache configuration. + + Returns: + True if the key exists. + """ + valkey_key = self._make_key(key) + result = self._run_async(self.client.exists([valkey_key])) + return result > 0