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
This commit is contained in:
yangdx 2025-06-29 15:03:57 +08:00
parent 37bf341a69
commit e2824b721e

View file

@ -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}})