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,
+        )
+