refactor: Refactor cache packages

dbgpt/app/component_configs.py

@@ -139,7 +139,7 @@ def _load_model(self) -> "Embeddings":
 
 
 def _initialize_model_cache(system_app: SystemApp):
-    from dbgpt.util.cache import initialize_cache
+    from dbgpt.storage.cache import initialize_cache
 
     if not CFG.MODEL_CACHE_ENABLE:
         logger.info("Model cache is not enable")

dbgpt/app/scene/base_chat.py

@@ -544,7 +544,7 @@ def _build_model_operator(
         ModelSaveCacheOperator,
         ModelStreamSaveCacheOperator,
     )
-    from dbgpt.util.cache import CacheManager
+    from dbgpt.storage.cache import CacheManager
 
     # Fetch worker and cache managers from the system configuration
     worker_manager = CFG.SYSTEM_APP.get_component(

dbgpt/core/__init__.py

@@ -12,6 +12,14 @@
 )
 from dbgpt.core.interface.prompt import PromptTemplate, PromptTemplateOperator
 from dbgpt.core.interface.output_parser import BaseOutputParser
+from dbgpt.core.interface.serialization import Serializable, Serializer
+from dbgpt.core.interface.cache import (
+    CacheKey,
+    CacheValue,
+    CacheClient,
+    CachePolicy,
+    CacheConfig,
+)
 
 __ALL__ = [
     "ModelInferenceMetrics",
@@ -20,8 +28,16 @@
     "BaseLLMOperator",
     "RequestBuildOperator",
     "ModelMessage",
-    "ModelMessageRoleType" "OnceConversation",
+    "ModelMessageRoleType",
+    "OnceConversation",
     "PromptTemplate",
     "PromptTemplateOperator",
     "BaseOutputParser",
+    "Serializable",
+    "Serializer",
+    "CacheKey",
+    "CacheValue",
+    "CacheClient",
+    "CachePolicy",
+    "CacheConfig",
 ]

dbgpt/util/cache/base.py → dbgpt/core/interface/cache.py

@@ -1,34 +1,15 @@
-from abc import ABC, abstractmethod, abstractclassmethod
-from typing import Any, TypeVar, Generic, Optional, Type, Dict
+from abc import ABC, abstractmethod
+
+from typing import Any, TypeVar, Generic, Optional
 from dataclasses import dataclass
 from enum import Enum
 
-T = TypeVar("T", bound="Serializable")
+from dbgpt.core.interface.serialization import Serializable
 
 K = TypeVar("K")
 V = TypeVar("V")
 
 
-class Serializable(ABC):
-    @abstractmethod
-    def serialize(self) -> bytes:
-        """Convert the object into bytes for storage or transmission.
-
-        Returns:
-            bytes: The byte array after serialization
-        """
-
-    @abstractmethod
-    def to_dict(self) -> Dict:
-        """Convert the object's state to a dictionary."""
-
-    # @staticmethod
-    # @abstractclassmethod
-    # def from_dict(cls: Type["Serializable"], obj_dict: Dict) -> "Serializable":
-    #     """Deserialize a dictionary to an Serializable object.
-    #     """
-
-
 class RetrievalPolicy(str, Enum):
     EXACT_MATCH = "exact_match"
     SIMILARITY_MATCH = "similarity_match"
@@ -82,30 +63,6 @@ def get_value(self) -> V:
         """Get the underlying real value."""
 
 
-class Serializer(ABC):
-    """The serializer abstract class for serializing cache keys and values."""
-
-    @abstractmethod
-    def serialize(self, obj: Serializable) -> bytes:
-        """Serialize a cache object.
-
-        Args:
-            obj (Serializable): The object to serialize
-        """
-
-    @abstractmethod
-    def deserialize(self, data: bytes, cls: Type[Serializable]) -> Serializable:
-        """Deserialize data back into a cache object of the specified type.
-
-        Args:
-            data (bytes): The byte array to deserialize
-            cls (Type[Serializable]): The type of current object
-
-        Returns:
-            Serializable: The serializable object
-        """
-
-
 class CacheClient(ABC, Generic[K, V]):
     """The cache client interface."""
 

dbgpt/core/interface/serialization.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Type, Dict
+
+
+class Serializable(ABC):
+    @abstractmethod
+    def serialize(self) -> bytes:
+        """Convert the object into bytes for storage or transmission.
+
+        Returns:
+            bytes: The byte array after serialization
+        """
+
+    @abstractmethod
+    def to_dict(self) -> Dict:
+        """Convert the object's state to a dictionary."""
+
+
+class Serializer(ABC):
+    """The serializer abstract class for serializing cache keys and values."""
+
+    @abstractmethod
+    def serialize(self, obj: Serializable) -> bytes:
+        """Serialize a cache object.
+
+        Args:
+            obj (Serializable): The object to serialize
+        """
+
+    @abstractmethod
+    def deserialize(self, data: bytes, cls: Type[Serializable]) -> Serializable:
+        """Deserialize data back into a cache object of the specified type.
+
+        Args:
+            data (bytes): The byte array to deserialize
+            cls (Type[Serializable]): The type of current object
+
+        Returns:
+            Serializable: The serializable object
+        """

dbgpt/model/operator/model_operator.py

@@ -11,7 +11,7 @@
 from dbgpt.core.awel.operator.base import BaseOperator
 from dbgpt.core import ModelOutput
 from dbgpt.model.cluster import WorkerManager, WorkerManagerFactory
-from dbgpt.util.cache import LLMCacheClient, CacheManager, LLMCacheKey, LLMCacheValue
+from dbgpt.storage.cache import LLMCacheClient, CacheManager, LLMCacheKey, LLMCacheValue
 
 logger = logging.getLogger(__name__)
 

dbgpt/storage/cache/__init__.py

@@ -0,0 +1,12 @@
+from dbgpt.storage.cache.manager import CacheManager, initialize_cache
+from dbgpt.storage.cache.storage.base import MemoryCacheStorage
+from dbgpt.storage.cache.llm_cache import LLMCacheKey, LLMCacheValue, LLMCacheClient
+
+__all__ = [
+    "LLMCacheKey",
+    "LLMCacheValue",
+    "LLMCacheClient",
+    "CacheManager",
+    "initialize_cache",
+    "MemoryCacheStorage",
+]

dbgpt/util/cache/llm_cache.py → dbgpt/storage/cache/llm_cache.py

@@ -1,18 +1,15 @@
 from typing import Optional, Dict, Any, Union, List
 from dataclasses import dataclass, asdict
-import json
 import hashlib
 
-from dbgpt.util.cache.base import (
+from dbgpt.core.interface.cache import (
     CacheKey,
     CacheValue,
-    Serializer,
     CacheClient,
     CacheConfig,
 )
-from dbgpt.util.cache.manager import CacheManager
-from dbgpt.util.cache.storage.base import CacheStorage
-from dbgpt.core import ModelOutput
+from dbgpt.storage.cache.manager import CacheManager
+from dbgpt.core import ModelOutput, Serializer
 from dbgpt.model.base import ModelType
 
 

dbgpt/util/cache/manager.py → dbgpt/storage/cache/manager.py

@@ -2,10 +2,9 @@
 from typing import Optional, Type
 import logging
 from concurrent.futures import Executor
-from dbgpt.util.cache.storage.base import CacheStorage, StorageItem
-from dbgpt.util.cache.base import (
-    K,
-    V,
+from dbgpt.storage.cache.storage.base import CacheStorage
+from dbgpt.core.interface.cache import K, V
+from dbgpt.core import (
     CacheKey,
     CacheValue,
     CacheConfig,
@@ -103,13 +102,13 @@ def serializer(self) -> Serializer:
 def initialize_cache(
     system_app: SystemApp, storage_type: str, max_memory_mb: int, persist_dir: str
 ):
-    from dbgpt.util.cache.protocal.json_protocal import JsonSerializer
-    from dbgpt.util.cache.storage.base import MemoryCacheStorage
+    from dbgpt.util.serialization.json_serialization import JsonSerializer
+    from dbgpt.storage.cache.storage.base import MemoryCacheStorage
 
     cache_storage = None
     if storage_type == "disk":
         try:
-            from dbgpt.util.cache.storage.disk.disk_storage import DiskCacheStorage
+            from dbgpt.storage.cache.storage.disk.disk_storage import DiskCacheStorage
 
             cache_storage = DiskCacheStorage(
                 persist_dir, mem_table_buffer_mb=max_memory_mb

dbgpt/util/cache/storage/base.py → dbgpt/storage/cache/storage/base.py

@@ -5,12 +5,11 @@
 import msgpack
 import logging
 
-from dbgpt.util.cache.base import (
+from dbgpt.core.interface.cache import (
     K,
     V,
     CacheKey,
     CacheValue,
-    CacheClient,
     CacheConfig,
     RetrievalPolicy,
     CachePolicy,

dbgpt/util/cache/storage/disk/disk_storage.py → dbgpt/storage/cache/storage/disk/disk_storage.py

@@ -1,18 +1,16 @@
 from typing import Optional
 import logging
-from dbgpt.util.cache.base import (
+from rocksdict import Rdict, Options
+
+from dbgpt.core.interface.cache import (
     K,
     V,
     CacheKey,
     CacheValue,
     CacheConfig,
     RetrievalPolicy,
-    CachePolicy,
 )
-from dbgpt.util.cache.storage.base import StorageItem, CacheStorage
-from rocksdict import Rdict
-from rocksdict import Rdict, Options, SliceTransform, PlainTableFactoryOptions
-
+from dbgpt.storage.cache.storage.base import StorageItem, CacheStorage
 
 logger = logging.getLogger(__name__)
 

dbgpt/util/cache/protocal/json_protocal.py → dbgpt/util/serialization/json_serialization.py

@@ -2,7 +2,7 @@
 from typing import Dict, Type
 import json
 
-from dbgpt.util.cache.base import Serializable, Serializer
+from dbgpt.core.interface.serialization import Serializable, Serializer
 
 JSON_ENCODING = "utf-8"