Spec clients

trustgraph-base/trustgraph/base/__init__.py

@@ -14,5 +14,7 @@ from . subscriber_spec import SubscriberSpec
 from . request_response_spec import RequestResponseSpec
 from . llm_service import LlmService, LlmResult
 from . embeddings_service import EmbeddingsService
-
+from . embeddings_client import EmbeddingsClientSpec
+from . text_completion_client import TextCompletionsClientSpec
+from . prompt_client import PromptClientSpec
 

trustgraph-base/trustgraph/base/embedings_client.py

@@ -0,0 +1,26 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import EmbeddingsReqeust, EmbeddingsResponse
+
+class EmbeddingsClient(RequestResponse):
+    async def embed(self, text, timeout=30):
+        resp = await self.request(
+            EmbeddingsRequest(
+                text = text
+            ),
+            timeout=timeout
+        )
+        return resp.vectors
+
+class EmbeddingsClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(EmbeddingsRequestResponseSpec, self).__init__(
+            request_name = request_name,
+            request_schema = EmbeddingsRequest,
+            response_name = response_name,
+            response_schema = EmbeddingsResponse,
+            impl = EmbeddingsClient,
+        )
+

trustgraph-base/trustgraph/base/prompt_client.py

@@ -0,0 +1,52 @@
+
+import json
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import EmbeddingsReqeust, EmbeddingsResponse
+
+class PromptClient(RequestResponse):
+
+    async def prompt(self, id, variables, timeout=600):
+
+        resp = await self.request(
+            PromptRequest(
+                id = id,
+                terms = {
+                    k: json.dumps(v)
+                    for k, v in variables.items()
+                }
+            ),
+            timeout=timeout
+        )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        if resp.text: return resp.text
+
+        return json.loads(resp.object)
+
+    async def extract_definitions(self, text, timeout=600):
+        return await self.prompt(
+            id = "extract-definitions",
+            variables = { "text": text }
+        )
+
+    async def extract_relationships(self, text, timeout=600):
+        return await self.prompt(
+            id = "extract-relationships",
+            variables = { "text": text }
+        )
+
+class PromptClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(EmbeddingsRequestResponseSpec, self).__init__(
+            request_name = request_name,
+            request_schema = EmbeddingsRequest,
+            response_name = response_name,
+            response_schema = EmbeddingsResponse,
+            impl = TextCompletionClient,
+        )
+

trustgraph-base/trustgraph/base/request_response_spec.py

@@ -75,12 +75,14 @@ class RequestResponse(Subscriber):
 # - we receive on the response topic as a subscriber
 class RequestResponseSpec(Spec):
     def __init__(
-            self, request_name, request_schema, response_name, response_schema
+            self, request_name, request_schema, response_name,
+            response_schema, impl=RequestResponse
     ):
         self.request_name = request_name
         self.request_schema = request_schema
         self.response_name = response_name
         self.response_schema = response_schema
+        self.impl = impl
 
     def add(self, flow, processor, definition):
 
@@ -88,7 +90,7 @@ class RequestResponseSpec(Spec):
             flow.id, f"{flow.name}-{self.response_name}"
         )
 
-        rr = RequestResponse(
+        rr = self.impl(
             client = processor.client,
             subscription = flow.id,
             consumer_name = flow.id,
@@ -102,3 +104,4 @@ class RequestResponseSpec(Spec):
 
         flow.consumer[self.request_name] = rr
 
+

trustgraph-base/trustgraph/base/text_completion_client.py

@@ -0,0 +1,30 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import EmbeddingsReqeust, EmbeddingsResponse
+
+class TextCompletionClient(RequestResponse):
+    async def text_completion(self, system, prompt, timeout=600):
+        resp = await self.request(
+            TextCompletionRequest(
+                system = system, prompt = prompt
+            ),
+            timeout=timeout
+        )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp.response
+
+class TextCompletionClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(EmbeddingsRequestResponseSpec, self).__init__(
+            request_name = request_name,
+            request_schema = EmbeddingsRequest,
+            response_name = response_name,
+            response_schema = EmbeddingsResponse,
+            impl = TextCompletionClient,
+        )
+

trustgraph-flow/trustgraph/embeddings/fastembed/processor.py

@@ -49,4 +49,3 @@ def run():
 
     Processor.launch(default_ident, __doc__)
 
-

trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py

@@ -8,7 +8,7 @@ Output is entity plus embedding.
 from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
 from ... schema import EmbeddingsRequest, EmbeddingsResponse
 
-from ... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
+from ... base import FlowProcessor, EmbeddingsClientSpec, ConsumerSpec
 from ... base import ProducerSpec
 
 default_ident = "graph-embeddings"
@@ -34,11 +34,9 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            RequestResponseSpec(
+            EmbeddingsClientSpec(
                 request_name = "embeddings-request",
-                request_schema = EmbeddingsRequest,
                 response_name = "embeddings-response",
-                response_schema = EmbeddingsResponse,
             )
         )
 
@@ -60,10 +58,8 @@ class Processor(FlowProcessor):
 
             for entity in v.entities:
 
-                resp = await flow("embeddings-request").request(
+                resp = await flow("embeddings-request").embed(
-                    EmbeddingsRequest(
+                    text = entity.context
-                        text = entity.context
-                    )
                 )
 
                 vectors = resp.vectors

trustgraph-flow/trustgraph/extract/kg/definitions/extract.py

@@ -13,8 +13,8 @@ from .... schema import EntityContext, EntityContexts
 from .... schema import PromptRequest, PromptResponse
 from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
 
-from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
+from .... base import FlowProcessor, ConsumerSpec,  ProducerSpec
-from .... base import ProducerSpec
+from .... base import PromptClientSpec
 
 DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
 RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
@@ -43,11 +43,9 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            RequestResponseSpec(
+            PromptClientSpec(
                 request_name = "prompt-request",
-                request_schema = PromptRequest,
                 response_name = "prompt-response",
-                response_schema = PromptResponse,
             )
         )
 
@@ -102,24 +100,14 @@ class Processor(FlowProcessor):
 
             try:
 
-                resp = await flow("prompt-request").request(
+                defs = await flow("prompt-request").extract_definitions(
-                    PromptRequest(
+                    text = chunk
-                        id="extract-definitions",
-                        terms={
-                            "text": json.dumps(chunk)
-                        },
-                    )
                 )
-                print("Response", resp, flush=True)
 
-                if resp.error is not None:
+                print("Response", defs, flush=True)
-                    print("Error:", resp.error.message, flush=True)
-                    raise RuntimeError(resp.error.message)
 
-                if resp.object is None:
+                if type(defs) != list:
-                    raise RuntimeError("Expecting object in prompt response")
+                    raise RuntimeError("Expecting array in prompt response")
-
-                defs = json.loads(resp.object)
 
             except Exception as e:
                 print("Prompt exception:", e, flush=True)

trustgraph-flow/trustgraph/extract/kg/relationships/extract.py

@@ -13,8 +13,8 @@ from .... schema import Metadata, Value
 from .... schema import PromptRequest, PromptResponse
 from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF
 
-from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
+from .... base import FlowProcessor, ConsumerSpec,  ProducerSpec
-from .... base import ProducerSpec
+from .... base import PromptClientSpec
 
 RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
 SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
@@ -42,11 +42,9 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            RequestResponseSpec(
+            PromptClientSpec(
                 request_name = "prompt-request",
-                request_schema = PromptRequest,
                 response_name = "prompt-response",
-                response_schema = PromptResponse,
             )
         )
 
@@ -86,24 +84,14 @@ class Processor(FlowProcessor):
 
             try:
 
-                resp = await flow("prompt-request").request(
+                rels = await flow("prompt-request").extract_relationships(
-                    PromptRequest(
+                    text = chunk
-                        id="extract-relationships",
-                        terms={
-                            "text": json.dumps(chunk)
-                        },
-                    )
                 )
-                print("Response", resp, flush=True)
 
-                if resp.error is not None:
+                print("Response", rels, flush=True)
-                    print("Error:", resp.error.message, flush=True)
-                    raise RuntimeError(resp.error.message)
 
-                if resp.object is None:
+                if type(rels) != list:
-                    raise RuntimeError("Expecting object in prompt response")
+                    raise RuntimeError("Expecting array in prompt response")
-
-                rels = json.loads(resp.object)
 
             except Exception as e:
                 print("Prompt exception:", e, flush=True)
@@ -118,9 +106,9 @@ class Processor(FlowProcessor):
 
             for rel in rels:
 
-                s = rel["s"]
+                s = rel["subject"]
-                p = rel["p"]
+                p = rel["predicate"]
-                o = rel["o"]
+                o = rel["object"]
 
                 if s == "": continue
                 if p == "": continue
@@ -136,7 +124,7 @@ class Processor(FlowProcessor):
                 p_uri = self.to_uri(p)
                 p_value = Value(value=str(p_uri), is_uri=True)
 
-                if rel["o_entity"]: 
+                if rel["object-entity"]: 
                     o_uri = self.to_uri(o)
                     o_value = Value(value=str(o_uri), is_uri=True)
                 else: