bingo
commited on
Commit
·
7471fe1
1
Parent(s):
2a54ada
add logger.debug for mongo_impl get_by_mode_and_id()
Browse files- lightrag/kg/mongo_impl.py +7 -5
- lightrag/kg/neo4j_impl.py +5 -1
lightrag/kg/mongo_impl.py
CHANGED
@@ -45,7 +45,9 @@ class MongoKVStorage(BaseKVStorage):
|
|
45 |
for mode, items in data.items():
|
46 |
for k, v in tqdm_async(items.items(), desc="Upserting"):
|
47 |
key = f"{mode}_{k}"
|
48 |
-
result = self._data.update_one(
|
|
|
|
|
49 |
if result.upserted_id:
|
50 |
logger.debug(f"\nInserted new document with key: {key}")
|
51 |
data[mode][k]["_id"] = key
|
@@ -54,20 +56,20 @@ class MongoKVStorage(BaseKVStorage):
|
|
54 |
self._data.update_one({"_id": k}, {"$set": v}, upsert=True)
|
55 |
data[k]["_id"] = k
|
56 |
return data
|
57 |
-
|
58 |
async def get_by_mode_and_id(self, mode: str, id: str) -> Union[dict, None]:
|
59 |
if "llm_response_cache" == self.namespace:
|
60 |
res = {}
|
61 |
-
v = self._data.find_one({"_id": mode+"_"+id})
|
62 |
if v:
|
63 |
res[id] = v
|
64 |
-
|
65 |
return res
|
66 |
else:
|
67 |
return None
|
68 |
else:
|
69 |
return None
|
70 |
-
|
71 |
async def drop(self):
|
72 |
""" """
|
73 |
pass
|
|
|
45 |
for mode, items in data.items():
|
46 |
for k, v in tqdm_async(items.items(), desc="Upserting"):
|
47 |
key = f"{mode}_{k}"
|
48 |
+
result = self._data.update_one(
|
49 |
+
{"_id": key}, {"$setOnInsert": v}, upsert=True
|
50 |
+
)
|
51 |
if result.upserted_id:
|
52 |
logger.debug(f"\nInserted new document with key: {key}")
|
53 |
data[mode][k]["_id"] = key
|
|
|
56 |
self._data.update_one({"_id": k}, {"$set": v}, upsert=True)
|
57 |
data[k]["_id"] = k
|
58 |
return data
|
59 |
+
|
60 |
async def get_by_mode_and_id(self, mode: str, id: str) -> Union[dict, None]:
|
61 |
if "llm_response_cache" == self.namespace:
|
62 |
res = {}
|
63 |
+
v = self._data.find_one({"_id": mode + "_" + id})
|
64 |
if v:
|
65 |
res[id] = v
|
66 |
+
logger.debug(f"llm_response_cache find one by:{id}")
|
67 |
return res
|
68 |
else:
|
69 |
return None
|
70 |
else:
|
71 |
return None
|
72 |
+
|
73 |
async def drop(self):
|
74 |
""" """
|
75 |
pass
|
lightrag/kg/neo4j_impl.py
CHANGED
@@ -48,7 +48,11 @@ class Neo4JStorage(BaseGraphStorage):
|
|
48 |
URI, auth=(USERNAME, PASSWORD)
|
49 |
)
|
50 |
_database_name = "home database" if DATABASE is None else f"database {DATABASE}"
|
51 |
-
with GraphDatabase.driver(
|
|
|
|
|
|
|
|
|
52 |
try:
|
53 |
with _sync_driver.session(database=DATABASE) as session:
|
54 |
try:
|
|
|
48 |
URI, auth=(USERNAME, PASSWORD)
|
49 |
)
|
50 |
_database_name = "home database" if DATABASE is None else f"database {DATABASE}"
|
51 |
+
with GraphDatabase.driver(
|
52 |
+
URI,
|
53 |
+
auth=(USERNAME, PASSWORD),
|
54 |
+
max_connection_pool_size=MAX_CONNECTION_POOL_SIZE,
|
55 |
+
) as _sync_driver:
|
56 |
try:
|
57 |
with _sync_driver.session(database=DATABASE) as session:
|
58 |
try:
|