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 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({"_id": key}, {"$setOnInsert": v}, upsert=True)
 
 
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
- print(f"find one by:{id}")
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(URI, auth=(USERNAME, PASSWORD), max_connection_pool_size=MAX_CONNECTION_POOL_SIZE) as _sync_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: