You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
133 lines
4.7 KiB
133 lines
4.7 KiB
STORAGE_IMPLEMENTATIONS = {
|
|
"KV_STORAGE": {
|
|
"implementations": [
|
|
"JsonKVStorage",
|
|
"RedisKVStorage",
|
|
"PGKVStorage",
|
|
"MongoKVStorage",
|
|
# "TiDBKVStorage",
|
|
],
|
|
"required_methods": ["get_by_id", "upsert"],
|
|
},
|
|
"GRAPH_STORAGE": {
|
|
"implementations": [
|
|
"NetworkXStorage",
|
|
"Neo4JStorage",
|
|
"PGGraphStorage",
|
|
# "AGEStorage",
|
|
# "MongoGraphStorage",
|
|
# "TiDBGraphStorage",
|
|
# "GremlinStorage",
|
|
],
|
|
"required_methods": ["upsert_node", "upsert_edge"],
|
|
},
|
|
"VECTOR_STORAGE": {
|
|
"implementations": [
|
|
"NanoVectorDBStorage",
|
|
"MilvusVectorDBStorage",
|
|
"ChromaVectorDBStorage",
|
|
"PGVectorStorage",
|
|
"FaissVectorDBStorage",
|
|
"QdrantVectorDBStorage",
|
|
"MongoVectorDBStorage",
|
|
# "TiDBVectorDBStorage",
|
|
],
|
|
"required_methods": ["query", "upsert"],
|
|
},
|
|
"DOC_STATUS_STORAGE": {
|
|
"implementations": [
|
|
"JsonDocStatusStorage",
|
|
"PGDocStatusStorage",
|
|
"MongoDocStatusStorage",
|
|
],
|
|
"required_methods": ["get_docs_by_status"],
|
|
},
|
|
}
|
|
|
|
# Storage implementation environment variable without default value
|
|
STORAGE_ENV_REQUIREMENTS: dict[str, list[str]] = {
|
|
# KV Storage Implementations
|
|
"JsonKVStorage": [],
|
|
"MongoKVStorage": [],
|
|
"RedisKVStorage": ["REDIS_URI"],
|
|
# "TiDBKVStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
|
"PGKVStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
|
# Graph Storage Implementations
|
|
"NetworkXStorage": [],
|
|
"Neo4JStorage": ["NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD"],
|
|
"MongoGraphStorage": [],
|
|
# "TiDBGraphStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
|
"AGEStorage": [
|
|
"AGE_POSTGRES_DB",
|
|
"AGE_POSTGRES_USER",
|
|
"AGE_POSTGRES_PASSWORD",
|
|
],
|
|
# "GremlinStorage": ["GREMLIN_HOST", "GREMLIN_PORT", "GREMLIN_GRAPH"],
|
|
"PGGraphStorage": [
|
|
"POSTGRES_USER",
|
|
"POSTGRES_PASSWORD",
|
|
"POSTGRES_DATABASE",
|
|
],
|
|
# Vector Storage Implementations
|
|
"NanoVectorDBStorage": [],
|
|
"MilvusVectorDBStorage": [],
|
|
"ChromaVectorDBStorage": [],
|
|
# "TiDBVectorDBStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
|
"PGVectorStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
|
"FaissVectorDBStorage": [],
|
|
"QdrantVectorDBStorage": ["QDRANT_URL"], # QDRANT_API_KEY has default value None
|
|
"MongoVectorDBStorage": [],
|
|
# Document Status Storage Implementations
|
|
"JsonDocStatusStorage": [],
|
|
"PGDocStatusStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
|
"MongoDocStatusStorage": [],
|
|
}
|
|
|
|
# Storage implementation module mapping
|
|
STORAGES = {
|
|
"NetworkXStorage": ".kg.networkx_impl",
|
|
"JsonKVStorage": ".kg.json_kv_impl",
|
|
"NanoVectorDBStorage": ".kg.nano_vector_db_impl",
|
|
"JsonDocStatusStorage": ".kg.json_doc_status_impl",
|
|
"Neo4JStorage": ".kg.neo4j_impl",
|
|
"MilvusVectorDBStorage": ".kg.milvus_impl",
|
|
"MongoKVStorage": ".kg.mongo_impl",
|
|
"MongoDocStatusStorage": ".kg.mongo_impl",
|
|
"MongoGraphStorage": ".kg.mongo_impl",
|
|
"MongoVectorDBStorage": ".kg.mongo_impl",
|
|
"RedisKVStorage": ".kg.redis_impl",
|
|
"ChromaVectorDBStorage": ".kg.chroma_impl",
|
|
# "TiDBKVStorage": ".kg.tidb_impl",
|
|
# "TiDBVectorDBStorage": ".kg.tidb_impl",
|
|
# "TiDBGraphStorage": ".kg.tidb_impl",
|
|
"PGKVStorage": ".kg.postgres_impl",
|
|
"PGVectorStorage": ".kg.postgres_impl",
|
|
"AGEStorage": ".kg.age_impl",
|
|
"PGGraphStorage": ".kg.postgres_impl",
|
|
# "GremlinStorage": ".kg.gremlin_impl",
|
|
"PGDocStatusStorage": ".kg.postgres_impl",
|
|
"FaissVectorDBStorage": ".kg.faiss_impl",
|
|
"QdrantVectorDBStorage": ".kg.qdrant_impl",
|
|
}
|
|
|
|
|
|
def verify_storage_implementation(storage_type: str, storage_name: str) -> None:
|
|
"""Verify if storage implementation is compatible with specified storage type
|
|
|
|
Args:
|
|
storage_type: Storage type (KV_STORAGE, GRAPH_STORAGE etc.)
|
|
storage_name: Storage implementation name
|
|
|
|
Raises:
|
|
ValueError: If storage implementation is incompatible or missing required methods
|
|
"""
|
|
if storage_type not in STORAGE_IMPLEMENTATIONS:
|
|
raise ValueError(f"Unknown storage type: {storage_type}")
|
|
|
|
storage_info = STORAGE_IMPLEMENTATIONS[storage_type]
|
|
if storage_name not in storage_info["implementations"]:
|
|
raise ValueError(
|
|
f"Storage implementation '{storage_name}' is not compatible with {storage_type}. "
|
|
f"Compatible implementations are: {', '.join(storage_info['implementations'])}"
|
|
)
|