gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat optimize storage configuration and environment variables

Browse Source 
* add storage type compatibility validation table
* add enviroment variables check for storage
* modify storage init to get setting from confing.ini and env
This commit is contained in:
yangdx
2025-02-11 00:55:52 +08:00
parent d0779209d9
commit 56c1792767
9 changed files with 249 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lightrag/kg/gremlin_impl.py
View File

@@ -47,7 +47,9 @@ class GremlinStorage(BaseGraphStorage):
# All vertices will have graph={GRAPH} property, so that we can
# have several logical graphs for one source
GRAPH = GremlinStorage._to_value_map(os.environ["GREMLIN_GRAPH"])
GRAPH = GremlinStorage._to_value_map(
os.environ.get("GREMLIN_GRAPH", "LightRAG")
)
self.graph_name = GRAPH

18
lightrag/kg/milvus_impl.py
View File

@@ -5,14 +5,17 @@ from dataclasses import dataclass
import numpy as np
from lightrag.utils import logger
from ..base import BaseVectorStorage
import pipmaster as pm
import configparser
if not pm.is_installed("pymilvus"):
pm.install("pymilvus")
from pymilvus import MilvusClient
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
@dataclass
class MilvusVectorDBStorge(BaseVectorStorage):
@staticmethod
@@ -27,14 +30,11 @@ class MilvusVectorDBStorge(BaseVectorStorage):
def __post_init__(self):
self._client = MilvusClient(
uri=os.environ.get(
"MILVUS_URI",
os.path.join(self.global_config["working_dir"], "milvus_lite.db"),
),
user=os.environ.get("MILVUS_USER", ""),
password=os.environ.get("MILVUS_PASSWORD", ""),
token=os.environ.get("MILVUS_TOKEN", ""),
db_name=os.environ.get("MILVUS_DB_NAME", ""),
uri = os.environ.get("MILVUS_URI", config.get("milvus", "uri", fallback=os.path.join(self.global_config["working_dir"], "milvus_lite.db"))),
user = os.environ.get("MILVUS_USER", config.get("milvus", "user", fallback=None)),
password = os.environ.get("MILVUS_PASSWORD", config.get("milvus", "password", fallback=None)),
token = os.environ.get("MILVUS_TOKEN", config.get("milvus", "token", fallback=None)),
db_name = os.environ.get("MILVUS_DB_NAME", config.get("milvus", "db_name", fallback=None)),
)
self._max_batch_size = self.global_config["embedding_batch_num"]
MilvusVectorDBStorge.create_collection_if_not_exist(

17
lightrag/kg/mongo_impl.py
View File

@@ -1,8 +1,8 @@
import os
from dataclasses import dataclass
import numpy as np
import pipmaster as pm
import configparser
from tqdm.asyncio import tqdm as tqdm_async
if not pm.is_installed("pymongo"):
@@ -12,22 +12,23 @@ if not pm.is_installed("motor"):
pm.install("motor")
from typing import Any, List, Tuple, Union
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import MongoClient
from ..base import BaseGraphStorage, BaseKVStorage
from ..namespace import NameSpace, is_namespace
from ..utils import logger
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
@dataclass
class MongoKVStorage(BaseKVStorage):
def __post_init__(self):
client = MongoClient(
os.environ.get("MONGO_URI", "mongodb://root:root@localhost:27017/")
os.environ.get("MONGO_URI", config.get("mongodb", "uri", fallback="mongodb://root:root@localhost:27017/"))
)
database = client.get_database(os.environ.get("MONGO_DATABASE", "LightRAG"))
database = client.get_database(os.environ.get("MONGO_DATABASE", mongo_database = config.get("mongodb", "database", fallback="LightRAG")))
self._data = database.get_collection(self.namespace)
logger.info(f"Use MongoDB as KV {self.namespace}")
@@ -90,10 +91,10 @@ class MongoGraphStorage(BaseGraphStorage):
embedding_func=embedding_func,
)
self.client = AsyncIOMotorClient(
os.environ.get("MONGO_URI", "mongodb://root:root@localhost:27017/")
os.environ.get("MONGO_URI", config.get("mongodb", "uri", fallback="mongodb://root:root@localhost:27017/"))
)
self.db = self.client[os.environ.get("MONGO_DATABASE", "LightRAG")]
self.collection = self.db[os.environ.get("MONGO_KG_COLLECTION", "MDB_KG")]
self.db = self.client[os.environ.get("MONGO_DATABASE", mongo_database = config.get("mongodb", "database", fallback="LightRAG"))]
self.collection = self.db[os.environ.get("MONGO_KG_COLLECTION", config.getboolean("mongodb", "kg_collection", fallback="MDB_KG"))]
#
# -------------------------------------------------------------------------

14
lightrag/kg/neo4j_impl.py
View File

@@ -5,6 +5,7 @@ import re
from dataclasses import dataclass
from typing import Any, Union, Tuple, List, Dict
import pipmaster as pm
import configparser
if not pm.is_installed("neo4j"):
pm.install("neo4j")
@@ -27,6 +28,9 @@ from ..utils import logger
from ..base import BaseGraphStorage
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
@dataclass
class Neo4JStorage(BaseGraphStorage):
@staticmethod
@@ -41,13 +45,15 @@ class Neo4JStorage(BaseGraphStorage):
)
self._driver = None
self._driver_lock = asyncio.Lock()
URI = os.environ["NEO4J_URI"]
USERNAME = os.environ["NEO4J_USERNAME"]
PASSWORD = os.environ["NEO4J_PASSWORD"]
MAX_CONNECTION_POOL_SIZE = os.environ.get("NEO4J_MAX_CONNECTION_POOL_SIZE", 800)
URI = os.environ["NEO4J_URI", config.get("neo4j", "uri", fallback=None)]
USERNAME = os.environ["NEO4J_USERNAME", config.get("neo4j", "username", fallback=None)]
PASSWORD = os.environ["NEO4J_PASSWORD", config.get("neo4j", "password", fallback=None)]
MAX_CONNECTION_POOL_SIZE = os.environ.get("NEO4J_MAX_CONNECTION_POOL_SIZE", config.get("neo4j", "connection_pool_size", fallback=800))
DATABASE = os.environ.get(
"NEO4J_DATABASE", re.sub(r"[^a-zA-Z0-9-]", "-", namespace)
)
self._driver: AsyncDriver = AsyncGraphDatabase.driver(
URI, auth=(USERNAME, PASSWORD)
)

138
lightrag/kg/postgres_impl_test.py
View File

@@ -1,138 +0,0 @@
import asyncio
import sys
import os
import pipmaster as pm
if not pm.is_installed("psycopg-pool"):
pm.install("psycopg-pool")
pm.install("psycopg[binary,pool]")
if not pm.is_installed("asyncpg"):
pm.install("asyncpg")
import asyncpg
import psycopg
from psycopg_pool import AsyncConnectionPool
from ..kg.postgres_impl import PostgreSQLDB, PGGraphStorage
from ..namespace import NameSpace
DB = "rag"
USER = "rag"
PASSWORD = "rag"
HOST = "localhost"
PORT = "15432"
os.environ["AGE_GRAPH_NAME"] = "dickens"
if sys.platform.startswith("win"):
import asyncio.windows_events
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
async def get_pool():
return await asyncpg.create_pool(
f"postgres://{USER}:{PASSWORD}@{HOST}:{PORT}/{DB}",
min_size=10,
max_size=10,
max_queries=5000,
max_inactive_connection_lifetime=300.0,
)
async def main1():
connection_string = (
f"dbname='{DB}' user='{USER}' password='{PASSWORD}' host='{HOST}' port={PORT}"
)
pool = AsyncConnectionPool(connection_string, open=False)
await pool.open()
try:
conn = await pool.getconn(timeout=10)
async with conn.cursor() as curs:
try:
await curs.execute('SET search_path = ag_catalog, "$user", public')
await curs.execute("SELECT create_graph('dickens-2')")
await conn.commit()
print("create_graph success")
except (
psycopg.errors.InvalidSchemaName,
psycopg.errors.UniqueViolation,
):
print("create_graph already exists")
await conn.rollback()
finally:
pass
db = PostgreSQLDB(
config={
"host": "localhost",
"port": 15432,
"user": "rag",
"password": "rag",
"database": "r1",
}
)
async def query_with_age():
await db.initdb()
graph = PGGraphStorage(
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
global_config={},
embedding_func=None,
)
graph.db = db
res = await graph.get_node('"A CHRISTMAS CAROL"')
print("Node is: ", res)
res = await graph.get_edge('"A CHRISTMAS CAROL"', "PROJECT GUTENBERG")
print("Edge is: ", res)
res = await graph.get_node_edges('"SCROOGE"')
print("Node Edges are: ", res)
async def create_edge_with_age():
await db.initdb()
graph = PGGraphStorage(
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
global_config={},
embedding_func=None,
)
graph.db = db
await graph.upsert_node('"THE CRATCHITS"', {"hello": "world"})
await graph.upsert_node('"THE GIRLS"', {"world": "hello"})
await graph.upsert_edge(
'"THE CRATCHITS"',
'"THE GIRLS"',
edge_data={
"weight": 7.0,
"description": '"The girls are part of the Cratchit family, contributing to their collective efforts and shared experiences.',
"keywords": '"family, collective effort"',
"source_id": "chunk-1d4b58de5429cd1261370c231c8673e8",
},
)
res = await graph.get_edge("THE CRATCHITS", '"THE GIRLS"')
print("Edge is: ", res)
async def main():
pool = await get_pool()
sql = r"SELECT * FROM ag_catalog.cypher('dickens', $$ MATCH (n:帅哥) RETURN n $$) AS (n ag_catalog.agtype)"
# cypher = "MATCH (n:how_are_you_doing) RETURN n"
async with pool.acquire() as conn:
try:
await conn.execute(
"""SET search_path = ag_catalog, "$user", public;select create_graph('dickens')"""
)
except asyncpg.exceptions.InvalidSchemaNameError:
print("create_graph already exists")
# stmt = await conn.prepare(sql)
row = await conn.fetch(sql)
print("row is: ", row)
row = await conn.fetchrow("select '100'::int + 200 as result")
print(row) # <Record result=300>
if __name__ == "__main__":
asyncio.run(query_with_age())

10
lightrag/kg/qdrant_impl.py
View File

@@ -5,11 +5,10 @@ from dataclasses import dataclass
import numpy as np
import hashlib
import uuid
from ..utils import logger
from ..base import BaseVectorStorage
import pipmaster as pm
import configparser
if not pm.is_installed("qdrant_client"):
pm.install("qdrant_client")
@@ -17,6 +16,9 @@ if not pm.is_installed("qdrant_client"):
from qdrant_client import QdrantClient, models
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
def compute_mdhash_id_for_qdrant(
content: str, prefix: str = "", style: str = "simple"
) -> str:
@@ -57,8 +59,8 @@ class QdrantVectorDBStorage(BaseVectorStorage):
def __post_init__(self):
self._client = QdrantClient(
url=os.environ.get("QDRANT_URL"),
api_key=os.environ.get("QDRANT_API_KEY", None),
url=os.environ.get("QDRANT_URL", config.get("qdrant", "uri", fallback=None)),
api_key=os.environ.get("QDRANT_API_KEY", config.get("qdrant", "apikey", fallback=None)),
)
self._max_batch_size = self.global_config["embedding_batch_num"]
QdrantVectorDBStorage.create_collection_if_not_exist(

6
lightrag/kg/redis_impl.py
View File

@@ -3,6 +3,7 @@ from typing import Any, Union
from tqdm.asyncio import tqdm as tqdm_async
from dataclasses import dataclass
import pipmaster as pm
import configparser
if not pm.is_installed("redis"):
pm.install("redis")
@@ -14,10 +15,13 @@ from lightrag.base import BaseKVStorage
import json
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
@dataclass
class RedisKVStorage(BaseKVStorage):
def __post_init__(self):
redis_url = os.environ.get("REDIS_URI", "redis://localhost:6379")
redis_url = os.environ.get("REDIS_URI", config.get("redis", "uri", fallback="redis://localhost:6379"))
self._redis = Redis.from_url(redis_url, decode_responses=True)
logger.info(f"Use Redis as KV {self.namespace}")