Graph RAG working I think

tests/test-triples

@@ -3,7 +3,9 @@
 import pulsar
 from trustgraph.clients.triples_query_client import TriplesQueryClient
 
-tq = TriplesQueryClient(pulsar_host="pulsar://localhost:6650")
+tq = TriplesQueryClient(
+    pulsar_host="pulsar://localhost:6650",
+)
 
 e = "http://trustgraph.ai/e/shuttle"
 

trustgraph-base/trustgraph/base/prompt_client.py

@@ -29,13 +29,15 @@ class PromptClient(RequestResponse):
     async def extract_definitions(self, text, timeout=600):
         return await self.prompt(
             id = "extract-definitions",
-            variables = { "text": text }
+            variables = { "text": text },
+            timeout = timeout,
         )
 
     async def extract_relationships(self, text, timeout=600):
         return await self.prompt(
             id = "extract-relationships",
-            variables = { "text": text }
+            variables = { "text": text },
+            timeout = timeout,
         )
 
     async def kg_prompt(self, query, kg, timeout=600):
@@ -47,7 +49,8 @@ class PromptClient(RequestResponse):
                     { "s": v[0], "p": v[1], "o": v[2] }
                     for v in kg
                 ]
-            }
+            },
+            timeout = timeout,
         )
 
 class PromptClientSpec(RequestResponseSpec):

trustgraph-base/trustgraph/base/triples_client.py

@@ -1,7 +1,13 @@
 
 from . request_response_spec import RequestResponse, RequestResponseSpec
 from .. schema import TriplesQueryRequest, TriplesQueryResponse, Value
-from .. knowledge import Uri, Literal, Triple
+from .. knowledge import Uri, Literal
+
+class Triple:
+    def __init__(self, s, p, o):
+        self.s = s
+        self.p = p
+        self.o = o
 
 def to_value(x):
     if x.is_uri: return Uri(x.value)
@@ -34,11 +40,13 @@ class TriplesClient(RequestResponse):
         if resp.error:
             raise RuntimeError(resp.error.message)
 
-        return [
+        triples = [
             Triple(to_value(v.s), to_value(v.p), to_value(v.o))
             for v in resp.triples
         ]
 
+        return triples
+
 class TriplesClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,

trustgraph-base/trustgraph/knowledge/defs.py

@@ -23,10 +23,6 @@ URL = 'https://schema.org/url'
 IDENTIFIER = 'https://schema.org/identifier'
 KEYWORD = 'https://schema.org/keywords'
 
-class Triple:
-    def __init__(self, s, p, o):
-        self.s, self.p, self.o = s, p, o
-
 class Uri(str):
     def is_uri(self): return True
     def is_literal(self): return False

trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py

@@ -39,15 +39,11 @@ class Query:
         if self.verbose:
             print("Get entities...", flush=True)
 
-        print("HERE>", query, flush=True)
-
         entities = await self.rag.graph_embeddings_client.query(
             vectors=vectors, limit=self.entity_limit,
             user=self.user, collection=self.collection,
         )
 
-        print("ENT>", entities, flush=True)
-
         entities = [
             str(e)
             for e in entities
@@ -65,9 +61,9 @@ class Query:
         if e in self.rag.label_cache:
             return self.rag.label_cache[e]
 
-        res = await self.rag.triples_client.request(
+        res = await self.rag.triples_client.query(
-            user=self.user, collection=self.collection,
             s=e, p=LABEL, o=None, limit=1,
+            user=self.user, collection=self.collection,
         )
 
         if len(res) == 0:
@@ -100,10 +96,10 @@ class Query:
             if path_length > 1:
                 await self.follow_edges(str(triple.o), subgraph, path_length-1)
 
-        res = await self.rag.triples_client.request(
+        res = await self.rag.triples_client.query(
-            user=self.user, collection=self.collection,
             s=None, p=ent, o=None,
-            limit=self.triple_limit
+            limit=self.triple_limit,
+            user=self.user, collection=self.collection,
         )
 
         for triple in res:
@@ -111,18 +107,20 @@ class Query:
                 (str(triple.s), str(triple.p), str(triple.o))
             )
 
-        res = await self.rag.triples_client.request(
+        res = await self.rag.triples_client.query(
-            user=self.user, collection=self.collection,
             s=None, p=None, o=ent,
             limit=self.triple_limit,
+            user=self.user, collection=self.collection,
         )
 
         for triple in res:
             subgraph.add(
-                (triple.s.value, triple.p.value, triple.o.value)
+                (str(triple.s), str(triple.p), str(triple.o))
             )
             if path_length > 1:
-                await self.follow_edges(triple.s.value, subgraph, path_length-1)
+                await self.follow_edges(
+                    str(triple.s), subgraph, path_length-1
+                )
 
     async def get_subgraph(self, query):
 
@@ -212,7 +210,7 @@ class GraphRag:
             print(kg)
             print(query)
 
-        resp = await self.prompt_client.request_kg_prompt(query, kg)
+        resp = await self.prompt_client.kg_prompt(query, kg)
 
         if self.verbose:
             print("Done", flush=True)