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:
parent
37bf341a69
commit
e2824b721e
1 changed files with 11 additions and 1 deletions
|
|
@ -98,7 +98,17 @@ class MongoKVStorage(BaseKVStorage):
|
||||||
self._data = None
|
self._data = None
|
||||||
|
|
||||||
async def get_by_id(self, id: str) -> dict[str, Any] | 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]]:
|
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||||
cursor = self._data.find({"_id": {"$in": ids}})
|
cursor = self._data.find({"_id": {"$in": ids}})
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue