69358d830d
Format entire codebase with ruff and add type hints across all modules: - Apply ruff formatting to all Python files (121 files, 17K insertions) - Add type hints to function signatures throughout lightrag core and API - Update test suite with improved type annotations and docstrings - Add pyrightconfig.json for static type checking configuration - Create prompt_optimized.py and test_extraction_prompt_ab.py test files - Update ruff.toml and .gitignore for improved linting configuration - Standardize code style across examples, reproduce scripts, and utilities
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
STORAGE_IMPLEMENTATIONS = {
|
|
'KV_STORAGE': {
|
|
'implementations': [
|
|
'JsonKVStorage',
|
|
'RedisKVStorage',
|
|
'PGKVStorage',
|
|
'MongoKVStorage',
|
|
],
|
|
'required_methods': ['get_by_id', 'upsert'],
|
|
},
|
|
'GRAPH_STORAGE': {
|
|
'implementations': [
|
|
'NetworkXStorage',
|
|
'Neo4JStorage',
|
|
'PGGraphStorage',
|
|
'MongoGraphStorage',
|
|
'MemgraphStorage',
|
|
],
|
|
'required_methods': ['upsert_node', 'upsert_edge'],
|
|
},
|
|
'VECTOR_STORAGE': {
|
|
'implementations': [
|
|
'NanoVectorDBStorage',
|
|
'MilvusVectorDBStorage',
|
|
'PGVectorStorage',
|
|
'FaissVectorDBStorage',
|
|
'QdrantVectorDBStorage',
|
|
'MongoVectorDBStorage',
|
|
# "ChromaVectorDBStorage",
|
|
],
|
|
'required_methods': ['query', 'upsert'],
|
|
},
|
|
'DOC_STATUS_STORAGE': {
|
|
'implementations': [
|
|
'JsonDocStatusStorage',
|
|
'RedisDocStatusStorage',
|
|
'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': [
|
|
'MONGO_URI',
|
|
'MONGO_DATABASE',
|
|
],
|
|
'RedisKVStorage': ['REDIS_URI'],
|
|
'PGKVStorage': ['POSTGRES_USER', 'POSTGRES_PASSWORD', 'POSTGRES_DATABASE'],
|
|
# Graph Storage Implementations
|
|
'NetworkXStorage': [],
|
|
'Neo4JStorage': ['NEO4J_URI', 'NEO4J_USERNAME', 'NEO4J_PASSWORD'],
|
|
'MongoGraphStorage': [
|
|
'MONGO_URI',
|
|
'MONGO_DATABASE',
|
|
],
|
|
'MemgraphStorage': ['MEMGRAPH_URI'],
|
|
'AGEStorage': [
|
|
'AGE_POSTGRES_DB',
|
|
'AGE_POSTGRES_USER',
|
|
'AGE_POSTGRES_PASSWORD',
|
|
],
|
|
'PGGraphStorage': [
|
|
'POSTGRES_USER',
|
|
'POSTGRES_PASSWORD',
|
|
'POSTGRES_DATABASE',
|
|
],
|
|
# Vector Storage Implementations
|
|
'NanoVectorDBStorage': [],
|
|
'MilvusVectorDBStorage': [
|
|
'MILVUS_URI',
|
|
'MILVUS_DB_NAME',
|
|
],
|
|
# "ChromaVectorDBStorage": [],
|
|
'PGVectorStorage': ['POSTGRES_USER', 'POSTGRES_PASSWORD', 'POSTGRES_DATABASE'],
|
|
'FaissVectorDBStorage': [],
|
|
'QdrantVectorDBStorage': ['QDRANT_URL'], # QDRANT_API_KEY has default value None
|
|
'MongoVectorDBStorage': [
|
|
'MONGO_URI',
|
|
'MONGO_DATABASE',
|
|
],
|
|
# Document Status Storage Implementations
|
|
'JsonDocStatusStorage': [],
|
|
'RedisDocStatusStorage': ['REDIS_URI'],
|
|
'PGDocStatusStorage': ['POSTGRES_USER', 'POSTGRES_PASSWORD', 'POSTGRES_DATABASE'],
|
|
'MongoDocStatusStorage': [
|
|
'MONGO_URI',
|
|
'MONGO_DATABASE',
|
|
],
|
|
}
|
|
|
|
# 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',
|
|
'RedisDocStatusStorage': '.kg.redis_impl',
|
|
'ChromaVectorDBStorage': '.kg.chroma_impl',
|
|
'PGKVStorage': '.kg.postgres_impl',
|
|
'PGVectorStorage': '.kg.postgres_impl',
|
|
'AGEStorage': '.kg.age_impl',
|
|
'PGGraphStorage': '.kg.postgres_impl',
|
|
'PGDocStatusStorage': '.kg.postgres_impl',
|
|
'FaissVectorDBStorage': '.kg.faiss_impl',
|
|
'QdrantVectorDBStorage': '.kg.qdrant_impl',
|
|
'MemgraphStorage': '.kg.memgraph_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"])}'
|
|
)
|