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:
         
     |