DP_PRACA/retrieval.py
2026-05-16 08:50:22 +02:00

37 lines
683 B
Python

def retrieve_subgraph(G, query):
query = query.lower()
context = []
for node in G.nodes(data=True):
node_name = node[0].lower()
if (
query in node_name
or node_name in query
):
for neighbor in G.neighbors(node[0]):
edge = G[node[0]][neighbor]
context.append(
f"{node[0]} --{edge.get('relation','related_to')}--> {neighbor}"
)
return "\n".join(context)
#test
from graph_builder import load_graph
if __name__ == "__main__":
G = load_graph()
result = retrieve_subgraph(
G,
"Python"
)
print(result)