gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix linting

Browse source
This commit is contained in:
yangdx 2025-07-03 12:26:05 +08:00
parent ff1b1c61c7
commit 70e154b0aa
1 changed files with 19 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
lightrag/kg/postgres_impl.py
View file

@ -160,7 +160,9 @@ class PostgreSQLDB:
) )
# Migrate existing data: extract cache_type from flattened keys # Migrate existing data: extract cache_type from flattened keys
logger.info("Migrating existing LLM cache data to populate cache_type field") logger.info(
"Migrating existing LLM cache data to populate cache_type field"
)
update_sql = """ update_sql = """
UPDATE LIGHTRAG_LLM_CACHE UPDATE LIGHTRAG_LLM_CACHE
SET cache_type = CASE SET cache_type = CASE
@ -176,7 +178,9 @@ class PostgreSQLDB:
"cache_type column already exists in LIGHTRAG_LLM_CACHE table" "cache_type column already exists in LIGHTRAG_LLM_CACHE table"
) )
except Exception as e: except Exception as e:
logger.warning(f"Failed to add cache_type column to LIGHTRAG_LLM_CACHE: {e}") logger.warning(
f"Failed to add cache_type column to LIGHTRAG_LLM_CACHE: {e}"
)
async def _migrate_timestamp_columns(self): async def _migrate_timestamp_columns(self):
"""Migrate timestamp columns in tables to timezone-aware types, assuming original data is in UTC time""" """Migrate timestamp columns in tables to timezone-aware types, assuming original data is in UTC time"""
@ -345,7 +349,7 @@ class PostgreSQLDB:
# Determine cache_type based on mode # Determine cache_type based on mode
cache_type = "extract" if record["mode"] == "default" else "unknown" cache_type = "extract" if record["mode"] == "default" else "unknown"
# Generate new flattened key # Generate new flattened key
new_key = f"{record['mode']}:{cache_type}:{new_hash}" new_key = f"{record['mode']}:{cache_type}:{new_hash}"
@ -519,7 +523,9 @@ class PostgreSQLDB:
try: try:
await self._migrate_llm_cache_add_cache_type() await self._migrate_llm_cache_add_cache_type()
except Exception as e: except Exception as e:
logger.error(f"PostgreSQL, Failed to migrate LLM cache cache_type field: {e}") logger.error(
f"PostgreSQL, Failed to migrate LLM cache cache_type field: {e}"
)
# Don't throw an exception, allow the initialization process to continue # Don't throw an exception, allow the initialization process to continue
# Finally, attempt to migrate old doc chunks data if needed # Finally, attempt to migrate old doc chunks data if needed
@ -787,7 +793,9 @@ class PGKVStorage(BaseKVStorage):
response["llm_cache_list"] = llm_cache_list response["llm_cache_list"] = llm_cache_list
# Special handling for LLM cache to ensure compatibility with _get_cached_extraction_results # Special handling for LLM cache to ensure compatibility with _get_cached_extraction_results
if response and is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE): if response and is_namespace(
self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE
):
# Map field names and add cache_type for compatibility # Map field names and add cache_type for compatibility
response = { response = {
**response, **response,
@ -821,7 +829,9 @@ class PGKVStorage(BaseKVStorage):
result["llm_cache_list"] = llm_cache_list result["llm_cache_list"] = llm_cache_list
# Special handling for LLM cache to ensure compatibility with _get_cached_extraction_results # Special handling for LLM cache to ensure compatibility with _get_cached_extraction_results
if results and is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE): if results and is_namespace(
self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE
):
processed_results = [] processed_results = []
for row in results: for row in results:
# Map field names and add cache_type for compatibility # Map field names and add cache_type for compatibility
@ -901,7 +911,9 @@ class PGKVStorage(BaseKVStorage):
"return_value": v["return"], "return_value": v["return"],
"mode": v.get("mode", "default"), # Get mode from data "mode": v.get("mode", "default"), # Get mode from data
"chunk_id": v.get("chunk_id"), "chunk_id": v.get("chunk_id"),
"cache_type": v.get("cache_type", "extract"), # Get cache_type from data "cache_type": v.get(
"cache_type", "extract"
), # Get cache_type from data
} }
await self.db.execute(upsert_sql, _data) await self.db.execute(upsert_sql, _data)