alllexx88 commited on
Commit
f39a7c1
·
1 Parent(s): d267c74

GremlinStorage: fix linting error, use asyncio.gather in get_node_edges()

Browse files
examples/lightrag_ollama_gremlin_demo.py CHANGED
@@ -1,9 +1,13 @@
1
  import asyncio
2
  import inspect
3
- import logging
4
  import os
5
 
6
- logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.WARN)
 
 
 
 
 
7
 
8
  from lightrag import LightRAG, QueryParam
9
  from lightrag.llm import ollama_embedding, ollama_model_complete
 
1
  import asyncio
2
  import inspect
 
3
  import os
4
 
5
+ # Uncomment these lines below to filter out somewhat verbose INFO level
6
+ # logging prints (the default loglevel is INFO).
7
+ # This has to go before the lightrag imports to work,
8
+ # which triggers linting errors, so we keep it commented out:
9
+ # import logging
10
+ # logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.WARN)
11
 
12
  from lightrag import LightRAG, QueryParam
13
  from lightrag.llm import ollama_embedding, ollama_model_complete
lightrag/kg/gremlin_impl.py CHANGED
@@ -306,13 +306,6 @@ class GremlinStorage(BaseGraphStorage):
306
  .project('connected_label')
307
  .by(__.label())
308
  """
309
- result1 = await self._query(query1)
310
- edges1 = (
311
- [(node_label, res["connected_label"]) for res in result1[0]]
312
- if result1
313
- else []
314
- )
315
-
316
  query2 = f"""
317
  {self.traverse_source_name}
318
  .V().has('graph', '{self.graph_name}')
@@ -322,7 +315,14 @@ class GremlinStorage(BaseGraphStorage):
322
  .project('connected_label')
323
  .by(__.select('connected').label())
324
  """
325
- result2 = await self._query(query2)
 
 
 
 
 
 
 
326
  edges2 = (
327
  [(res["connected_label"], node_label) for res in result2[0]]
328
  if result2
 
306
  .project('connected_label')
307
  .by(__.label())
308
  """
 
 
 
 
 
 
 
309
  query2 = f"""
310
  {self.traverse_source_name}
311
  .V().has('graph', '{self.graph_name}')
 
315
  .project('connected_label')
316
  .by(__.select('connected').label())
317
  """
318
+ result1, result2 = await asyncio.gather(
319
+ self._query(query1), self._query(query2)
320
+ )
321
+ edges1 = (
322
+ [(node_label, res["connected_label"]) for res in result1[0]]
323
+ if result1
324
+ else []
325
+ )
326
  edges2 = (
327
  [(res["connected_label"], node_label) for res in result2[0]]
328
  if result2