Update hyperrag.py

hyperrag.py

@@ -3,30 +3,19 @@ from vector_store import search
 
 GRAPH = nx.Graph()
 
-
-def add_node(text):
-
-    GRAPH.add_node(text)
-
-
-def link_nodes(a, b):
-
-    GRAPH.add_edge(a, b)
-
-
 def build_graph(docs):
 
     for i in range(len(docs)-1):
 
-        add_node(docs[i])
-        add_node(docs[i+1])
+        GRAPH.add_node(docs[i])
+        GRAPH.add_node(docs[i+1])
 
-        link_nodes(docs[i], docs[i+1])
+        GRAPH.add_edge(docs[i], docs[i+1])
 
 
 def hyperrag(query):
 
-    hits = search(query, 3)
+    hits = search(query)
 
     context = []
 
@@ -38,10 +27,8 @@ def hyperrag(query):
 
         if text in GRAPH:
 
-            neighbors = GRAPH.neighbors(text)
-
-            for n in neighbors:
+            for n in GRAPH.neighbors(text):
 
                 context.append(n)
 
-    return context[:6]
+    return context[:5]