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