From 9f2e747698c3cd40dba573b1adb62e369fb3bf37 Mon Sep 17 00:00:00 2001 From: Fangyin Cheng Date: Mon, 17 Mar 2025 15:37:02 +0800 Subject: [PATCH] fix: Fix retrieve error --- configs/dbgpt-bm25-rag.toml | 4 ++-- .../dbgpt-core/src/dbgpt/util/configure/manager.py | 4 ++-- .../src/dbgpt_ext/storage/graph_store/neo4j_store.py | 5 +++-- .../dbgpt_ext/storage/graph_store/tugraph_store.py | 2 +- .../storage/knowledge_graph/knowledge_graph.py | 2 +- .../dbgpt_ext/storage/knowledge_graph/open_spg.py | 5 +++-- .../dbgpt_ext/storage/vector_store/chroma_store.py | 12 +----------- .../dbgpt_ext/storage/vector_store/elastic_store.py | 2 +- .../dbgpt_ext/storage/vector_store/milvus_store.py | 2 +- .../storage/vector_store/oceanbase_store.py | 2 +- .../dbgpt_ext/storage/vector_store/pgvector_store.py | 4 +++- .../dbgpt_ext/storage/vector_store/weaviate_store.py | 2 ++ .../src/dbgpt_serve/rag/storage_manager.py | 2 +- 13 files changed, 22 insertions(+), 26 deletions(-) diff --git a/configs/dbgpt-bm25-rag.toml b/configs/dbgpt-bm25-rag.toml index 7ae1388f7..ba48b1b6b 100644 --- a/configs/dbgpt-bm25-rag.toml +++ b/configs/dbgpt-bm25-rag.toml @@ -26,11 +26,11 @@ rerank_top_k=3 [rag.storage] [rag.storage.vector] -type = "Chroma" +type = "chroma" persist_path = "pilot/data" [rag.storage.full_text] -type = "ElasticSearch" +type = "elasticsearch" host="127.0.0.1" port=9200 diff --git a/packages/dbgpt-core/src/dbgpt/util/configure/manager.py b/packages/dbgpt-core/src/dbgpt/util/configure/manager.py index 73288b1f9..ec985f2bb 100644 --- a/packages/dbgpt-core/src/dbgpt/util/configure/manager.py +++ b/packages/dbgpt-core/src/dbgpt/util/configure/manager.py @@ -147,7 +147,7 @@ class RegisterParameters(abc.ABC, metaclass=PolymorphicMeta): """Register a subclass with this base class using a type value.""" if not hasattr(cls, "_type_registry"): cls._type_registry = {} - cls._type_registry[type_value] = subclass + cls._type_registry[type_value.lower()] = subclass @classmethod def get_type_value(cls) -> str: @@ -238,7 +238,7 @@ def _get_concrete_class(base_class: Type[T], data: Dict[str, Any]) -> Type[T]: if not type_value: return base_class type_value = _resolve_env_vars(type_value) - real_cls = base_class.get_subclass(type_value) + real_cls = base_class.get_subclass(type_value.lower()) if not real_cls: raise ValueError( f"Unknown type value: {type_value}, known types: " diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/neo4j_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/neo4j_store.py index c88193efa..07a69d6fa 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/neo4j_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/neo4j_store.py @@ -1,17 +1,18 @@ """Neo4j store.""" import logging +from dataclasses import dataclass -from dbgpt._private.pydantic import ConfigDict from dbgpt.storage.graph_store.base import GraphStoreBase, GraphStoreConfig logger = logging.getLogger(__name__) +@dataclass class Neo4jStoreConfig(GraphStoreConfig): """Neo4j store config.""" - model_config = ConfigDict(arbitrary_types_allowed=True) + __type__ = "neo4j" class Neo4jStore(GraphStoreBase): diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/tugraph_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/tugraph_store.py index 777a58000..5b66edd6d 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/tugraph_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/graph_store/tugraph_store.py @@ -18,7 +18,7 @@ logger = logging.getLogger(__name__) class TuGraphStoreConfig(GraphStoreConfig): """TuGraph store config.""" - __type__ = "TuGraph" + __type__ = "tugraph" host: str = field( default="127.0.0.1", diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py index a49410546..adc1406c1 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py @@ -103,7 +103,7 @@ GRAPH_PARAMETERS = [ class BuiltinKnowledgeGraphConfig(KnowledgeGraphConfig): """Builtin knowledge graph config.""" - __type__ = "TuGraph" + __type__ = "tugraph" llm_model: Optional[str] = field( default=None, metadata={"description": "llm model name."} diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/open_spg.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/open_spg.py index 1a739f8f2..7e5f574a3 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/open_spg.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/open_spg.py @@ -1,17 +1,18 @@ """OpenSPG class.""" import logging +from dataclasses import dataclass -from dbgpt._private.pydantic import ConfigDict from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase, KnowledgeGraphConfig logger = logging.getLogger(__name__) +@dataclass class OpenSPGConfig(KnowledgeGraphConfig): """OpenSPG config.""" - model_config = ConfigDict(arbitrary_types_allowed=True) + __type__ = "openspg" class OpenSPG(KnowledgeGraphBase): diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py index 11b308e2d..37b32e332 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py @@ -40,7 +40,7 @@ logger = logging.getLogger(__name__) class ChromaVectorConfig(VectorStoreConfig): """Chroma vector store config.""" - __type__ = "Chroma" + __type__ = "chroma" persist_path: Optional[str] = field( default=os.getenv("CHROMA_PERSIST_PATH", None), @@ -249,16 +249,6 @@ class ChromaStore(VectorStoreBase): # Delete collection if it exists self._chroma_client.delete_collection(self._collection.name) SharedSystemClient.clear_system_cache() - - # Clean persist folder if it exists - if os.path.exists(self.persist_dir): - try: - self._clean_persist_folder() - except Exception as e: - logger.error(f"Failed to clean persist folder: {e}") - # Even if folder cleanup fails, collection deletion succeeded - return True - return True except Exception as e: diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/elastic_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/elastic_store.py index 6e19ceb08..5ff8f28a8 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/elastic_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/elastic_store.py @@ -77,7 +77,7 @@ logger = logging.getLogger(__name__) class ElasticsearchStoreConfig(VectorStoreConfig): """Elasticsearch vector store config.""" - __type__ = "ElasticSearch" + __type__ = "elasticsearch" uri: str = field( default="localhost", diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/milvus_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/milvus_store.py index 841720113..96ddcaa64 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/milvus_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/milvus_store.py @@ -98,7 +98,7 @@ logger = logging.getLogger(__name__) class MilvusVectorConfig(VectorStoreConfig): """Milvus vector store config.""" - __type__ = "Milvus" + __type__ = "milvus" uri: str = field( default=None, diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/oceanbase_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/oceanbase_store.py index 7848598f6..939d3a5b5 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/oceanbase_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/oceanbase_store.py @@ -126,7 +126,7 @@ def _normalize(vector: List[float]) -> List[float]: class OceanBaseConfig(VectorStoreConfig): """OceanBase vector store config.""" - __type__ = "OceanBase" + __type__ = "oceanbase" ob_host: Optional[str] = field( default=None, diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py index 87c63dd22..affaa76b2 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py @@ -41,6 +41,8 @@ logger = logging.getLogger(__name__) class PGVectorConfig(VectorStoreConfig): """PG vector store config.""" + __type__ = "pgvector" + connection_string: str = field( default=None, metadata={ @@ -50,7 +52,7 @@ class PGVectorConfig(VectorStoreConfig): ) def create_store(self, **kwargs) -> "PGVectorStore": - """Create Milvus Store.""" + """Create a PGVectorStore instance.""" return PGVectorStore(vector_store_config=self, **kwargs) diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/weaviate_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/weaviate_store.py index 1e410c4fb..777adbdd9 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/weaviate_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/weaviate_store.py @@ -49,6 +49,8 @@ logger = logging.getLogger(__name__) class WeaviateVectorConfig(VectorStoreConfig): """Weaviate vector store config.""" + __type__ = "weaviate" + weaviate_url: str = field( default=os.getenv("WEAVIATE_URL", None), metadata={ diff --git a/packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py b/packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py index ffe122676..3f7550c91 100644 --- a/packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py +++ b/packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py @@ -39,7 +39,7 @@ class StorageManager(BaseComponent): """Get storage connector.""" supported_vector_types = self.get_vector_supported_types storage_config = self.storage_config() - if storage_type in supported_vector_types: + if storage_type.lower() in supported_vector_types: return self.create_vector_store(index_name) elif storage_type == "KnowledgeGraph": if not storage_config.graph: