Merge branch 'shmily1012/main'
This commit is contained in:
yangdx
@@ -3,13 +3,16 @@ from typing import Any, final
|
||||
from dataclasses import dataclass
|
||||
import pipmaster as pm
|
||||
import configparser
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
if not pm.is_installed("redis"):
|
||||
pm.install("redis")
|
||||
|
||||
# aioredis is a depricated library, replaced with redis
|
||||
from redis.asyncio import Redis # type: ignore
|
||||
from lightrag.utils import logger
|
||||
from redis.asyncio import Redis, ConnectionPool
|
||||
from redis.exceptions import RedisError, ConnectionError
|
||||
from lightrag.utils import logger, compute_mdhash_id
|
||||
|
||||
from lightrag.base import BaseKVStorage
|
||||
import json
|
||||
|
||||
@@ -17,6 +20,11 @@ import json
|
||||
config = configparser.ConfigParser()
|
||||
config.read("config.ini", "utf-8")
|
||||
|
||||
# Constants for Redis connection pool
|
||||
MAX_CONNECTIONS = 50
|
||||
SOCKET_TIMEOUT = 5.0
|
||||
SOCKET_CONNECT_TIMEOUT = 3.0
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
@@ -25,22 +33,71 @@ class RedisKVStorage(BaseKVStorage):
|
||||
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}")
|
||||
# Create a connection pool with limits
|
||||
self._pool = ConnectionPool.from_url(
|
||||
redis_url,
|
||||
max_connections=MAX_CONNECTIONS,
|
||||
decode_responses=True,
|
||||
socket_timeout=SOCKET_TIMEOUT,
|
||||
socket_connect_timeout=SOCKET_CONNECT_TIMEOUT
|
||||
)
|
||||
self._redis = Redis(connection_pool=self._pool)
|
||||
logger.info(f"Initialized Redis connection pool for {self.namespace} with max {MAX_CONNECTIONS} connections")
|
||||
|
||||
@asynccontextmanager
|
||||
async def _get_redis_connection(self):
|
||||
"""Safe context manager for Redis operations."""
|
||||
try:
|
||||
yield self._redis
|
||||
except ConnectionError as e:
|
||||
logger.error(f"Redis connection error in {self.namespace}: {e}")
|
||||
raise
|
||||
except RedisError as e:
|
||||
logger.error(f"Redis operation error in {self.namespace}: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error in Redis operation for {self.namespace}: {e}")
|
||||
raise
|
||||
|
||||
async def close(self):
|
||||
"""Close the Redis connection pool to prevent resource leaks."""
|
||||
if hasattr(self, '_redis') and self._redis:
|
||||
await self._redis.close()
|
||||
await self._pool.disconnect()
|
||||
logger.debug(f"Closed Redis connection pool for {self.namespace}")
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Support for async context manager."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Ensure Redis resources are cleaned up when exiting context."""
|
||||
await self.close()
|
||||
|
||||
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
||||
data = await self._redis.get(f"{self.namespace}:{id}")
|
||||
async with self._get_redis_connection() as redis:
|
||||
try:
|
||||
data = await redis.get(f"{self.namespace}:{id}")
|
||||
return json.loads(data) if data else None
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"JSON decode error for id {id}: {e}")
|
||||
return None
|
||||
|
||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||
pipe = self._redis.pipeline()
|
||||
async with self._get_redis_connection() as redis:
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for id in ids:
|
||||
pipe.get(f"{self.namespace}:{id}")
|
||||
results = await pipe.execute()
|
||||
return [json.loads(result) if result else None for result in results]
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"JSON decode error in batch get: {e}")
|
||||
return [None] * len(ids)
|
||||
|
||||
async def filter_keys(self, keys: set[str]) -> set[str]:
|
||||
pipe = self._redis.pipeline()
|
||||
async with self._get_redis_connection() as redis:
|
||||
pipe = redis.pipeline()
|
||||
for key in keys:
|
||||
pipe.exists(f"{self.namespace}:{key}")
|
||||
results = await pipe.execute()
|
||||
@@ -49,32 +106,34 @@ class RedisKVStorage(BaseKVStorage):
|
||||
return set(keys) - existing_ids
|
||||
|
||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||
logger.info(f"Inserting {len(data)} to {self.namespace}")
|
||||
if not data:
|
||||
return
|
||||
pipe = self._redis.pipeline()
|
||||
|
||||
logger.info(f"Inserting {len(data)} items to {self.namespace}")
|
||||
async with self._get_redis_connection() as redis:
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for k, v in data.items():
|
||||
pipe.set(f"{self.namespace}:{k}", json.dumps(v))
|
||||
await pipe.execute()
|
||||
|
||||
for k in data:
|
||||
data[k]["_id"] = k
|
||||
except json.JSONEncodeError as e:
|
||||
logger.error(f"JSON encode error during upsert: {e}")
|
||||
raise
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
# Redis handles persistence automatically
|
||||
pass
|
||||
|
||||
async def delete(self, ids: list[str]) -> None:
|
||||
"""Delete entries with specified IDs
|
||||
|
||||
Args:
|
||||
ids: List of entry IDs to be deleted
|
||||
"""
|
||||
"""Delete entries with specified IDs"""
|
||||
if not ids:
|
||||
return
|
||||
|
||||
pipe = self._redis.pipeline()
|
||||
async with self._get_redis_connection() as redis:
|
||||
pipe = redis.pipeline()
|
||||
for id in ids:
|
||||
pipe.delete(f"{self.namespace}:{id}")
|
||||
|
||||
@@ -112,11 +171,12 @@ class RedisKVStorage(BaseKVStorage):
|
||||
Returns:
|
||||
dict[str, str]: Status of the operation with keys 'status' and 'message'
|
||||
"""
|
||||
async with self._get_redis_connection() as redis:
|
||||
try:
|
||||
keys = await self._redis.keys(f"{self.namespace}:*")
|
||||
keys = await redis.keys(f"{self.namespace}:*")
|
||||
|
||||
if keys:
|
||||
pipe = self._redis.pipeline()
|
||||
pipe = redis.pipeline()
|
||||
for key in keys:
|
||||
pipe.delete(key)
|
||||
results = await pipe.execute()
|
||||
@@ -131,3 +191,4 @@ class RedisKVStorage(BaseKVStorage):
|
||||
except Exception as e:
|
||||
logger.error(f"Error dropping keys from {self.namespace}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import traceback
|
||||
import json
|
||||
import re
|
||||
import os
|
||||
@@ -1002,6 +1003,7 @@ async def mix_kg_vector_query(
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in get_kg_context: {str(e)}")
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
async def get_vector_context():
|
||||
@@ -1392,9 +1394,16 @@ async def _find_most_related_text_unit_from_entities(
|
||||
all_text_units_lookup[c_id] = index
|
||||
tasks.append((c_id, index, this_edges))
|
||||
|
||||
results = await asyncio.gather(
|
||||
*[text_chunks_db.get_by_id(c_id) for c_id, _, _ in tasks]
|
||||
# Process in batches of 25 tasks at a time to avoid overwhelming resources
|
||||
batch_size = 25
|
||||
results = []
|
||||
|
||||
for i in range(0, len(tasks), batch_size):
|
||||
batch_tasks = tasks[i:i + batch_size]
|
||||
batch_results = await asyncio.gather(
|
||||
*[text_chunks_db.get_by_id(c_id) for c_id, _, _ in batch_tasks]
|
||||
)
|
||||
results.extend(batch_results)
|
||||
|
||||
for (c_id, index, this_edges), data in zip(tasks, results):
|
||||
all_text_units_lookup[c_id] = {
|
||||
|
||||
Reference in New Issue
Block a user