From 5e86a3c3aef6765b00bd66d244fca534b01149c8 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Sat, 26 Oct 2024 15:58:42 +0100
Subject: [PATCH] Definitions is working again

---
 tests/test-lang-definition                    |  8 +++++++-
 .../trustgraph/clients/prompt_client.py       | 19 +++++++++++++++----
 .../model/prompt/template/prompt_manager.py   | 19 ++++++++++++++++---
 .../model/prompt/template/service.py          | 11 -----------
 4 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/tests/test-lang-definition b/tests/test-lang-definition
index c6e593fd..67342779 100755
--- a/tests/test-lang-definition
+++ b/tests/test-lang-definition
@@ -7,7 +7,13 @@ p = PromptClient(pulsar_host="pulsar://localhost:6650")
 
 chunk = """I noticed a cat in my garden.  It is a four-legged animal
 which is a mammal and can be tame or wild.  I wonder if it will be friends
-with me.  I think the cat's name is Fred and it has 4 legs"""
+with me.  I think the cat's name is Fred and it has 4 legs.
+
+A cat is a small mammal.
+
+A grapefruit is a citrus fruit.
+
+"""
 
 resp = p.request_definitions(
     chunk=chunk,
diff --git a/trustgraph-base/trustgraph/clients/prompt_client.py b/trustgraph-base/trustgraph/clients/prompt_client.py
index b64266c1..4fc8958b 100644
--- a/trustgraph-base/trustgraph/clients/prompt_client.py
+++ b/trustgraph-base/trustgraph/clients/prompt_client.py
@@ -1,6 +1,7 @@
 
 import _pulsar
 import json
+import dataclasses
 
 from .. schema import PromptRequest, PromptResponse
 from .. schema import prompt_request_queue
@@ -13,6 +14,11 @@ WARN=_pulsar.LoggerLevel.Warn
 INFO=_pulsar.LoggerLevel.Info
 DEBUG=_pulsar.LoggerLevel.Debug
 
+@dataclasses.dataclass
+class Definition:
+    name: str
+    definition: str
+
 class PromptClient(BaseClient):
 
     def __init__(
@@ -56,20 +62,25 @@ class PromptClient(BaseClient):
 
     def request_definitions(self, chunk, timeout=300):
 
-        return self.request(
+        defs = self.request(
             id="extract-definitions",
             terms={
-                "chunk": chunk
+                "text": chunk
             },
             timeout=timeout
         )
 
+        return [
+            Definition(name=d["entity"], definition=d["definition"])
+            for d in defs
+        ]
+
     def request_relationships(self, chunk, timeout=300):
 
         return self.call(
             id="extract-relationships",
             terms={
-                "chunk": chunk
+                "text": chunk
             },
             timeout=timeout
         )
@@ -79,7 +90,7 @@ class PromptClient(BaseClient):
         return self.call(
             id="extract-topics",
             terms={
-                "chunk": chunk
+                "text": chunk
             },
             timeout=timeout
         )
diff --git a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py b/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py
index 8f2470a5..d8a032ca 100644
--- a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py
+++ b/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py
@@ -2,6 +2,7 @@
 import ibis
 import json
 from jsonschema import validate
+import re
 
 from trustgraph.clients.llm_client import LlmClient
 
@@ -42,6 +43,17 @@ class PromptManager:
             if v.terms is None:
                 v.terms = {}
 
+    def parse_json(self, text):
+        json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
+    
+        if json_match:
+            json_str = json_match.group(1).strip()
+        else:
+            # If no delimiters, assume the entire output is JSON
+            json_str = text.strip()
+
+        return json.loads(json_str)
+
     def invoke(self, id, input):
 
         if id not in self.prompts:
@@ -58,6 +70,8 @@ class PromptManager:
 
         resp = self.llm.request(**prompt)
 
+        print(resp, flush=True)
+
         if resp_type == "text":
             return resp
 
@@ -65,12 +79,11 @@ class PromptManager:
             raise RuntimeError(f"Response type {resp_type} not known")
 
         try:
-            obj = json.loads(resp)
+            obj = self.parse_json(resp)
         except:
             raise RuntimeError("JSON parse fail")
 
-        print(resp)
-        print(obj)
+        print(obj, flush=True)
         if self.prompts[id].schema:
             try:
                 print(self.prompts[id].schema)
diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py
index f3ffe0c2..5576610a 100755
--- a/trustgraph-flow/trustgraph/model/prompt/template/service.py
+++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py
@@ -155,17 +155,6 @@ class Processor(ConsumerProducer):
             config = prompt_configuration,
         )
 
-    def parse_json(self, text):
-        json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
-    
-        if json_match:
-            json_str = json_match.group(1).strip()
-        else:
-            # If no delimiters, assume the entire output is JSON
-            json_str = text.strip()
-
-        return json.loads(json_str)
-
     def handle(self, msg):
 
         v = msg.value()