From a759a1771e3968d8a7cbec0793089ec9d87f8e0f Mon Sep 17 00:00:00 2001
From: cybermaggedon <cybermaggedon@gmail.com>
Date: Fri, 4 Apr 2025 15:43:40 +0100
Subject: [PATCH] Added CLI tg-set-prompt command, can be used to update
 prompts dynamically. (#338)

Added CLI tg-set-prompt command, can be used to update prompts dynamically. (#338)
---
 trustgraph-cli/scripts/tg-set-prompt | 143 +++++++++++++++++++++++++++
 trustgraph-cli/setup.py              |   1 +
 2 files changed, 144 insertions(+)
 create mode 100755 trustgraph-cli/scripts/tg-set-prompt

diff --git a/trustgraph-cli/scripts/tg-set-prompt b/trustgraph-cli/scripts/tg-set-prompt
new file mode 100755
index 00000000..fa4e4c5d
--- /dev/null
+++ b/trustgraph-cli/scripts/tg-set-prompt
@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+
+"""
+Dumps out the current prompts
+"""
+
+import argparse
+import os
+from trustgraph.api import Api, ConfigKey, ConfigValue
+import json
+import tabulate
+import textwrap
+
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+
+def set_system(url, system):
+
+    api = Api(url)
+
+    api.config_put([
+        ConfigValue(type="prompt", key="system", value=json.dumps(system))
+    ])
+
+    print("System prompt set.")
+
+def set_prompt(url, id, prompt, response, schema):
+
+    api = Api(url)
+
+    values = api.config_get([
+        ConfigKey(type="prompt", key="template-index")
+    ])
+
+    ix = json.loads(values[0].value)
+
+    object = {
+        "id": id,
+        "prompt": prompt,
+    }
+
+    if response:
+        object["response-type"] = response
+    else:
+        object["response-type"] = "text"
+
+    if schema:
+        object["schema"] = schema
+
+    if id not in ix:
+        ix.append(id)
+
+    values = api.config_put([
+        ConfigValue(
+            type="prompt", key="template-index", value=json.dumps(ix)
+        ),
+        ConfigValue(
+            type="prompt", key=f"template.{id}", value=json.dumps(object)
+        )
+    ])
+
+    print("Prompt set.")
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='tg-show-prompts',
+        description=__doc__,
+    )
+
+    parser.add_argument(
+        '-u', '--api-url',
+        default=default_url,
+        help=f'API URL (default: {default_url})',
+    )
+
+    parser.add_argument(
+        '--id',
+        help=f'Prompt ID',
+    )
+
+    parser.add_argument(
+        '--response',
+        help=f'Response form, should be one of: text json',
+    )
+
+    parser.add_argument(
+        '--schema',
+        help=f'JSON schema, for JSON response form',
+    )
+
+    parser.add_argument(
+        '--prompt',
+        help=f'Prompt template',
+    )
+
+    parser.add_argument(
+        '--system',
+        help=f'System prompt',
+    )
+
+    args = parser.parse_args()
+
+    try:
+
+        if args.system:
+            if args.id or args.prompt or args.schema or args.response:
+                raise RuntimeError("Can't use --system with other args")
+                
+            set_system(
+                url=args.api_url, system=args.system
+            )
+
+        else:
+
+            if args.id is None:
+                raise RuntimeError("Must specify --id for prompt")
+
+            if args.prompt is None:
+                raise RuntimeError("Must specify --prompt for prompt")
+
+            if args.response:
+                if args.response not in ["text", "json"]:
+                    raise RuntimeError("Response must be one of: text json")
+                    
+            if args.schema:
+                try:
+                    schobj = json.loads(args.schema)
+                except:
+                    raise RuntimeError("JSON schema must be valid JSON")
+            else:
+                schobj = None
+
+            set_prompt(
+                url=args.api_url, id=args.id, prompt=args.prompt,
+                response=args.response, schema=schobj
+            )
+
+    except Exception as e:
+
+        print("Exception:", e, flush=True)
+
+main()
+
diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py
index 0ef523ea..dda401ec 100644
--- a/trustgraph-cli/setup.py
+++ b/trustgraph-cli/setup.py
@@ -63,6 +63,7 @@ setuptools.setup(
         "scripts/tg-save-kg-core",
         "scripts/tg-save-doc-embeds",
         "scripts/tg-show-config",
+        "scripts/tg-set-prompt",
         "scripts/tg-show-tools",
         "scripts/tg-show-prompts",
     ]