From 37bf341a6983b6570446e67b3ce65992c14ae7a5 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 29 Jun 2025 14:39:50 +0800 Subject: [PATCH 1/2] Fix LLM cache handling for PGKVStorage to address document deletion scenarios. - Add dynamic cache_type field - Support mode parameter for LLM cache - Maintain backward compatibility --- lightrag/kg/postgres_impl.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index a79ca448..7feae36c 100644 --- a/lightrag/kg/postgres_impl.py +++ b/lightrag/kg/postgres_impl.py @@ -504,14 +504,22 @@ class PGKVStorage(BaseKVStorage): async def get_by_id(self, id: str) -> dict[str, Any] | None: """Get doc_full data by id.""" sql = SQL_TEMPLATES["get_by_id_" + self.namespace] - params = {"workspace": self.db.workspace, "id": id} if is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE): + # For LLM cache, the id parameter actually represents the mode + params = {"workspace": self.db.workspace, "mode": id} array_res = await self.db.query(sql, params, multirows=True) res = {} for row in array_res: - res[row["id"]] = row + # Dynamically add cache_type field based on mode + row_with_cache_type = dict(row) + if id == "default": + row_with_cache_type["cache_type"] = "extract" + else: + row_with_cache_type["cache_type"] = "unknown" + res[row["id"]] = row_with_cache_type return res if res else None else: + params = {"workspace": self.db.workspace, "id": id} response = await self.db.query(sql, params) return response if response else None From e2824b721efd67a97da8c68861c816303d216cef Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 29 Jun 2025 15:03:57 +0800 Subject: [PATCH 2/2] Fix LLM cache handling for MongoKVStorage to address document deletion scenarios. - Support fetching all "default_" prefixed documents - Maintain original behavior for other IDs - Return dictionary of documents for "default" - Keep backward compatibility --- lightrag/kg/mongo_impl.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index 6483f6fd..a8dda1b9 100644 --- a/lightrag/kg/mongo_impl.py +++ b/lightrag/kg/mongo_impl.py @@ -98,7 +98,17 @@ class MongoKVStorage(BaseKVStorage): self._data = None async def get_by_id(self, id: str) -> dict[str, Any] | None: - return await self._data.find_one({"_id": id}) + if id == "default": + # Find all documents with _id starting with "default_" + cursor = self._data.find({"_id": {"$regex": "^default_"}}) + result = {} + async for doc in cursor: + # Use the complete _id as key + result[doc["_id"]] = doc + return result if result else None + else: + # Original behavior for non-"default" ids + return await self._data.find_one({"_id": id}) async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]: cursor = self._data.find({"_id": {"$in": ids}})