From 0f693dedd9364413a371c1f7770718a25bec8f47 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 16 Sep 2025 23:42:11 +0100
Subject: [PATCH] API additions

---
 trustgraph-base/trustgraph/api/flow.py        |  6 ++---
 .../messaging/translators/diagnosis.py        |  2 ++
 .../retrieval/structured_diag/service.py      | 25 +++++++++++++------
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py
index 58cf5c80..d1d5f95e 100644
--- a/trustgraph-base/trustgraph/api/flow.py
+++ b/trustgraph-base/trustgraph/api/flow.py
@@ -528,7 +528,7 @@ class FlowInstance:
             error_message = response["error"].get("message", "Unknown error")
             raise ProtocolException(f"{error_type}: {error_message}")
 
-        return response
+        return response["detected-type"]
 
     def generate_descriptor(self, sample, data_type, schema_name, options=None):
         """
@@ -565,7 +565,7 @@ class FlowInstance:
             error_message = response["error"].get("message", "Unknown error")
             raise ProtocolException(f"{error_type}: {error_message}")
 
-        return response
+        return response["descriptor"]
 
     def diagnose_data(self, sample, schema_name=None, options=None):
         """
@@ -635,5 +635,5 @@ class FlowInstance:
             error_message = response["error"].get("message", "Unknown error")
             raise ProtocolException(f"{error_type}: {error_message}")
 
-        return response
+        return response["schema-matches"]
 
diff --git a/trustgraph-base/trustgraph/messaging/translators/diagnosis.py b/trustgraph-base/trustgraph/messaging/translators/diagnosis.py
index f7b0150d..f4dbdbed 100644
--- a/trustgraph-base/trustgraph/messaging/translators/diagnosis.py
+++ b/trustgraph-base/trustgraph/messaging/translators/diagnosis.py
@@ -57,6 +57,8 @@ class StructuredDataDiagnosisResponseTranslator(MessageTranslator):
                 result["descriptor"] = obj.descriptor
         if obj.metadata:
             result["metadata"] = obj.metadata
+        if obj.schema_matches:
+            result["schema-matches"] = obj.schema_matches
 
         return result
 
diff --git a/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py b/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py
index fc45c71c..d69c8f17 100644
--- a/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py
+++ b/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py
@@ -21,6 +21,7 @@ default_ident = "structured-diag"
 default_csv_prompt = "diagnose-csv"
 default_json_prompt = "diagnose-json"
 default_xml_prompt = "diagnose-xml"
+default_schema_selection_prompt = "schema-selection"
 
 
 class Processor(FlowProcessor):
@@ -36,6 +37,7 @@ class Processor(FlowProcessor):
         self.csv_prompt = params.get("csv_prompt", default_csv_prompt)
         self.json_prompt = params.get("json_prompt", default_json_prompt)
         self.xml_prompt = params.get("xml_prompt", default_xml_prompt)
+        self.schema_selection_prompt = params.get("schema_selection_prompt", default_schema_selection_prompt)
 
         super(Processor, self).__init__(
             **params | {
@@ -336,21 +338,22 @@ class Processor(FlowProcessor):
             all_schemas.append(schema_info)
 
         # Create prompt variables - schemas array contains ALL schemas
+        # Note: The prompt expects 'question' not 'sample'
         variables = {
-            "sample": request.sample,
+            "question": request.sample,  # The prompt template expects 'question'
             "schemas": all_schemas,
             "options": request.options or {}
         }
 
-        # Call prompt service (assuming there's a schema-selection prompt template)
+        # Call prompt service with configurable template
         terms = {k: json.dumps(v) for k, v in variables.items()}
         prompt_request = PromptRequest(
-            id="schema-selection",  # This prompt template needs to exist
+            id=self.schema_selection_prompt,
             terms=terms
         )
 
         try:
-            logger.info("Calling prompt service for schema selection")
+            logger.info(f"Calling prompt service for schema selection with template: {self.schema_selection_prompt}")
             response = await flow("prompt-request").request(prompt_request)
 
             if response.error:
@@ -361,8 +364,16 @@ class Processor(FlowProcessor):
                 )
                 return StructuredDataDiagnosisResponse(error=error, operation=request.operation)
 
-            if not response.text or not response.text.strip():
-                logger.error("Empty response from prompt service")
+            # Check both text and object fields for response
+            response_data = None
+            if response.object and response.object.strip():
+                response_data = response.object.strip()
+                logger.debug(f"Using response from 'object' field: {response_data}")
+            elif response.text and response.text.strip():
+                response_data = response.text.strip()
+                logger.debug(f"Using response from 'text' field: {response_data}")
+            else:
+                logger.error("Empty response from prompt service (checked both text and object fields)")
                 error = Error(
                     type="PromptServiceError",
                     message="Empty response from prompt service"
@@ -371,7 +382,7 @@ class Processor(FlowProcessor):
 
             # Parse the response as JSON array of schema IDs
             try:
-                schema_matches = json.loads(response.text.strip())
+                schema_matches = json.loads(response_data)
                 if not isinstance(schema_matches, list):
                     raise ValueError("Response must be an array")
             except (json.JSONDecodeError, ValueError) as e: