Skip to content

Commit

Permalink
GremlinStorage: fix linting error, use asyncio.gather in get_node_edg…
Browse files Browse the repository at this point in the history
…es()
  • Loading branch information
alllexx88 committed Dec 20, 2024
1 parent 6f71293 commit 016d9f5
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 10 deletions.
8 changes: 6 additions & 2 deletions examples/lightrag_ollama_gremlin_demo.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
import asyncio
import inspect
import logging
import os

logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.WARN)
# Uncomment these lines below to filter out somewhat verbose INFO level
# logging prints (the default loglevel is INFO).
# This has to go before the lightrag imports to work,
# which triggers linting errors, so we keep it commented out:
# import logging
# logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.WARN)

from lightrag import LightRAG, QueryParam
from lightrag.llm import ollama_embedding, ollama_model_complete
Expand Down
16 changes: 8 additions & 8 deletions lightrag/kg/gremlin_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,13 +306,6 @@ async def get_node_edges(self, source_node_id: str) -> List[Tuple[str, str]]:
.project('connected_label')
.by(__.label())
"""
result1 = await self._query(query1)
edges1 = (
[(node_label, res["connected_label"]) for res in result1[0]]
if result1
else []
)

query2 = f"""
{self.traverse_source_name}
.V().has('graph', '{self.graph_name}')
Expand All @@ -322,7 +315,14 @@ async def get_node_edges(self, source_node_id: str) -> List[Tuple[str, str]]:
.project('connected_label')
.by(__.select('connected').label())
"""
result2 = await self._query(query2)
result1, result2 = await asyncio.gather(
self._query(query1), self._query(query2)
)
edges1 = (
[(node_label, res["connected_label"]) for res in result1[0]]
if result1
else []
)
edges2 = (
[(res["connected_label"], node_label) for res in result2[0]]
if result2
Expand Down

0 comments on commit 016d9f5

Please sign in to comment.