From dc0bcba269079c6dad7ec2814b95c1730706bf75 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Sat, 19 Apr 2025 13:41:33 +0100
Subject: [PATCH] Tidying clients

---
 .../trustgraph/base/embedings_client.py            |  5 +++++
 .../trustgraph/model/prompt/template/service.py    | 14 +++++---------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/trustgraph-base/trustgraph/base/embedings_client.py b/trustgraph-base/trustgraph/base/embedings_client.py
index 4016c000..89bdfa86 100644
--- a/trustgraph-base/trustgraph/base/embedings_client.py
+++ b/trustgraph-base/trustgraph/base/embedings_client.py
@@ -4,12 +4,17 @@ from .. schema import EmbeddingsReqeust, EmbeddingsResponse
 
 class EmbeddingsClient(RequestResponse):
     async def embed(self, text, timeout=30):
+
         resp = await self.request(
             EmbeddingsRequest(
                 text = text
             ),
             timeout=timeout
         )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
         return resp.vectors
 
 class EmbeddingsClientSpec(RequestResponseSpec):
diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py
index 70731245..3435b94c 100755
--- a/trustgraph-flow/trustgraph/model/prompt/template/service.py
+++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py
@@ -13,7 +13,7 @@ from .... schema import PromptRequest, PromptResponse, Error
 from .... schema import TextCompletionRequest, TextCompletionResponse
 
 from .... base import FlowProcessor
-from .... base import ProducerSpec, ConsumerSpec, RequestResponseSpec
+from .... base import ProducerSpec, ConsumerSpec, TextCompletionSpec
 
 from . prompt_manager import PromptConfiguration, Prompt, PromptManager
 
@@ -43,11 +43,9 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            RequestResponseSpec(
+            TextCompletionSpec(
                 request_name = "text-completion-request",
-                request_schema = TextCompletionRequest,
                 response_name = "text-completion-response",
-                response_schema = TextCompletionResponse,
             )
         )
 
@@ -139,14 +137,12 @@ class Processor(FlowProcessor):
                 print(system, flush=True)
                 print(prompt, flush=True)
 
-                resp = await flow("text-completion-request").request(
-                    TextCompletionRequest(
-                        system=system, prompt=prompt
-                    ),
+                resp = await flow("text-completion-request").text_completion(
+                    system = system, prompt = prompt,
                 )
 
                 try:
-                    return resp.response
+                    return resp
                 except Exception as e:
                     print("LLM Exception:", e, flush=True)
                     return None