Prompt manager integrated and working with 6 tests

tests/README.prompts

@@ -0,0 +1,27 @@
+
+test-prompt-... is tested with this prompt set...
+
+prompt-template \
+    -p pulsar://localhost:6650 \
+    --system-prompt 'You are a {{attitude}}, you are called {{name}}' \
+    --global-term \
+        'name=Craig' \
+        'attitude=LOUD, SHOUTY ANNOYING BOT' \
+    --prompt \
+        'question={{question}}' \
+        'french-question={{question}}' \
+        "analyze=Find the name and age in this text, and output a JSON structure containing just the name and age fields: {{description}}.  Don't add markup, just output the raw JSON object." \
+        "graph-query=Study the following knowledge graph, and then answer the question.\\n\nGraph:\\n{% for edge in knowledge %}({{edge.0}})-[{{edge.1}}]->({{edge.2}})\\n{%endfor%}\\nQuestion:\\n{{question}}" \
+        "extract-definition=Analyse the text provided, and then return a list of terms and definitions.  The output should be a JSON array, each item in the array is an object with fields 'term' and 'definition'.Don't add markup, just output the raw JSON object.  Here is the text:\\n{{text}}" \
+    --prompt-response-type \
+        'question=text' \
+        'analyze=json' \
+        'graph-query=text' \
+        'extract-definition=json' \
+    --prompt-term \
+        'question=name:Bonny' \
+        'french-question=attitude:French-speaking bot' \
+    --prompt-schema \
+        'analyze={ "type" : "object", "properties" : { "age": { "type" : "number" }, "name": { "type" : "string" } } }' \
+        'extract-definition={ "type": "array", "items": { "type": "object", "properties": { "term": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "term", "definition" ] } }'
+    

tests/test-prompt-analyze

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+import json
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+description = """Fred is a 4-legged cat who is 12 years old"""
+
+resp = p.request(
+    id="analyze",
+    terms = {
+        "description": description,
+    }
+)
+
+print(json.dumps(resp, indent=4))
+

tests/test-prompt-extraction

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+import json
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+chunk="""
+ The Space Shuttle was a reusable spacecraft that transported astronauts and cargo to and from Earth's orbit. It was designed to launch like a rocket, maneuver in orbit like a spacecraft, and land like an airplane. The Space Shuttle was NASA's space transportation system and was used for many purposes, including: 
+
+    Carrying astronauts
+    The Space Shuttle could carry up to seven astronauts at a time. 
+
+Launching, recovering, and repairing satellites
+The Space Shuttle could launch satellites into orbit, recover them, and repair them. 
+Building the International Space Station
+The Space Shuttle carried large parts into space to build the International Space Station. 
+Conducting research
+Astronauts conducted experiments in the Space Shuttle, which was like a science lab in space. 
+
+The Space Shuttle was retired in 2011 after the Columbia accident in 2003. The Columbia Accident Investigation Board report found that the Space Shuttle was unsafe and expensive to make safe. 
+Here are some other facts about the Space Shuttle: 
+
+    The Space Shuttle was 184 ft tall and had a diameter of 29 ft. 
+
+The Space Shuttle had a mass of 4,480,000 lb. 
+The Space Shuttle's first flight was on April 12, 1981. 
+The Space Shuttle's last mission was in 2011. 
+"""
+
+q = "Tell me some facts in the knowledge graph"
+
+resp = p.request(
+    id="extract-definition",
+    terms = {
+        "text": chunk,
+    }
+)
+
+print(resp)
+
+for fact in resp:
+    print(fact["term"], "::")
+    print(fact["definition"])
+    print()
+

tests/test-prompt-french-question

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+import pulsar
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+question = """What is the square root of 16?"""
+
+resp = p.request(
+    id="french-question",
+    terms = {
+        "question": question
+    }
+)
+
+print(resp)
+

tests/test-prompt-knowledge

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+import json
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+knowledge = [
+      ("accident", "evoked", "a wide range of deeply felt public responses"),
+      ("Space Shuttle concept", "had", "genesis"),
+      ("Commission", "had", "a mandate to develop recommendations for corrective or other action based upon the Commission's findings and determinations"),
+      ("Commission", "established", "teams of persons"),
+      ("Space Shuttle Challenger", "http://www.w3.org/2004/02/skos/core#definition", "A space shuttle that was destroyed in an accident during mission 51-L."),
+      ("The mid fuselage", "contains", "the payload bay"),
+      ("Volume I", "contains", "Chapter IX"),
+      ("accident", "resulted in", "firm national resolve that those men and women be forever enshrined in the annals of American heroes"),
+      ("Volume I", "contains", "Chapter VII"),
+      ("Volume I", "contains", "Chapter II"),
+      ("Volume I", "contains", "Chapter V"),
+      ("Commission", "believes", "its investigation and report have been responsive to the request of the President and hopes that they will serve the best interests of the nation in restoring the United States space program to its preeminent position in the world"),
+      ("Commission", "construe", "mandate"),
+      ("accident", "became", "a milestone on the way to achieving the full potential that space offers to mankind"),
+      ("Volume I", "contains", "The Commission"),
+      ("Commission", "http://www.w3.org/2004/02/skos/core#definition", "A group established to investigate the space shuttle accident"),
+      ("Volume I", "contains", "Appendix D"),
+      ("Commission", "had", "a mandate to review the circumstances surrounding the accident to establish the probable cause or causes of the accident"),
+      ("Volume I", "contains", "Recommendations")
+]
+
+q = "Tell me some facts in the knowledge graph"
+
+resp = p.request(
+    id="graph-query",
+    terms = {
+        "name": "Jayney",
+        "knowledge": knowledge,
+        "question": q
+    }
+)
+
+print(resp)
+
+
+

tests/test-prompt-question

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+import pulsar
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+question = """What is the square root of 16?"""
+
+resp = p.request(
+    id="question",
+    terms = {
+        "question": question
+    }
+)
+
+print(resp)
+

tests/test-prompt-spanish-question

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+import pulsar
+from trustgraph.clients.prompt_client import PromptClient
+
+p = PromptClient(pulsar_host="pulsar://localhost:6650")
+
+question = """What is the square root of 16?"""
+
+resp = p.request(
+    id="question",
+    terms = {
+        "question": question,
+        "attitude": "Spanish-speaking bot"
+    }
+)
+
+print(resp)
+

trustgraph-base/trustgraph/clients/prompt_client.py

@@ -1,7 +1,8 @@
 
 import _pulsar
+import json
 
-from .. schema import PromptRequest, PromptResponse, Fact, RowSchema, Field
+from .. schema import PromptRequest, PromptResponse
 from .. schema import prompt_request_queue
 from .. schema import prompt_response_queue
 from . base import BaseClient
@@ -38,12 +39,20 @@ class PromptClient(BaseClient):
             output_schema=PromptResponse,
         )
 
-    def request_definitions(self, chunk, timeout=300):
+    def request(self, id, terms, timeout=300):
 
-        return self.call(
-            kind="extract-definitions", chunk=chunk,
+        resp = self.call(
+            id=id,
+            terms={
+                k: json.dumps(v)
+                for k, v in terms.items()
+            },
             timeout=timeout
-        ).definitions
+        )
+
+        if resp.text: return resp.text
+
+        return json.loads(resp.object)
     
     def request_topics(self, chunk, timeout=300):
 

trustgraph-base/trustgraph/schema/prompt.py

@@ -39,20 +39,21 @@ class Fact(Record):
 #   schema, chunk -> rows
 
 class PromptRequest(Record):
-    kind = String()
-    chunk = String()
-    query = String()
-    kg = Array(Fact())
-    documents = Array(Bytes())
-    row_schema = RowSchema()
+    id = String()
+
+    # JSON encoded values
+    terms = Map(String())
 
 class PromptResponse(Record):
+
+    # Error case
     error = Error()
-    answer = String()
-    definitions = Array(Definition())
-    topics = Array(Topic())
-    relationships = Array(Relationship())
-    rows = Array(Map(String()))
+
+    # Just plain text
+    text = String()
+
+    # JSON encoded
+    object = String()
 
 prompt_request_queue = topic(
     'prompt', kind='non-persistent', namespace='request'

trustgraph-flow/setup.py

@@ -56,6 +56,8 @@ setuptools.setup(
         "neo4j",
         "tiktoken",
         "google-generativeai",
+        "ibis",
+        "jsonschema",
     ],
     scripts=[
         "scripts/chunker-recursive",

trustgraph-flow/trustgraph/model/prompt/template/README.md

@@ -0,0 +1,25 @@
+
+prompt-template \
+    -p pulsar://localhost:6650 \
+    --system-prompt 'You are a {{attitude}}, you are called {{name}}' \
+    --global-term \
+        'name=Craig' \
+        'attitude=LOUD, SHOUTY ANNOYING BOT' \
+    --prompt \
+        'question={{question}}' \
+        'french-question={{question}}' \
+        "analyze=Find the name and age in this text, and output a JSON structure containing just the name and age fields: {{description}}.  Don't add markup, just output the raw JSON object." \
+        "graph-query=Study the following knowledge graph, and then answer the question.\\n\nGraph:\\n{% for edge in knowledge %}({{edge.0}})-[{{edge.1}}]->({{edge.2}})\\n{%endfor%}\\nQuestion:\\n{{question}}" \
+        "extract-definition=Analyse the text provided, and then return a list of terms and definitions.  The output should be a JSON array, each item in the array is an object with fields 'term' and 'definition'.Don't add markup, just output the raw JSON object.  Here is the text:\\n{{text}}" \
+    --prompt-response-type \
+        'question=text' \
+        'analyze=json' \
+        'graph-query=text' \
+        'extract-definition=json' \
+    --prompt-term \
+        'question=name:Bonny' \
+        'french-question=attitude:French-speaking bot' \
+    --prompt-schema \
+        'analyze={ "type" : "object", "properties" : { "age": { "type" : "number" }, "name": { "type" : "string" } } }' \
+        'extract-definition={ "type": "array", "items": { "type": "object", "properties": { "term": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "term", "definition" ] } }'
+    

trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py

@@ -0,0 +1,82 @@
+
+import ibis
+import json
+from jsonschema import validate
+
+from trustgraph.clients.llm_client import LlmClient
+
+class PromptConfiguration:
+    def __init__(self, system_template, global_terms={}, prompts={}):
+        self.system_template = system_template
+        self.global_terms = global_terms
+        self.prompts = prompts
+
+class Prompt:
+    def __init__(self, template, response_type = "text", terms=None, schema=None):
+        self.template = template
+        self.response_type = response_type
+        self.terms = terms
+        self.schema = schema
+
+class PromptManager:
+
+    def __init__(self, llm, config):
+        self.llm = llm
+        self.config = config
+        self.terms = config.global_terms
+
+        self.prompts = config.prompts
+
+        try:
+            self.system_template = ibis.Template(config.system_template)
+        except:
+            raise RuntimeError("Error in system template")
+
+        self.templates = {}
+        for k, v in self.prompts.items():
+            try:
+                self.templates[k] = ibis.Template(v.template)
+            except:
+                raise RuntimeError(f"Error in template: {k}")
+
+            if v.terms is None:
+                v.terms = {}
+
+    def invoke(self, id, input):
+
+        if id not in self.prompts:
+            raise RuntimeError("ID invalid")
+
+        terms = self.terms | self.prompts[id].terms | input
+
+        resp_type = self.prompts[id].response_type
+
+        prompt = {
+            "system": self.system_template.render(terms),
+            "prompt": self.templates[id].render(terms)
+        }
+
+        resp = self.llm.request(**prompt)
+
+        if resp_type == "text":
+            return resp
+
+        if resp_type != "json":
+            raise RuntimeError(f"Response type {resp_type} not known")
+
+        try:
+            obj = json.loads(resp)
+        except:
+            raise RuntimeError("JSON parse fail")
+
+        print(resp)
+        print(obj)
+        if self.prompts[id].schema:
+            try:
+                print(self.prompts[id].schema)
+                validate(instance=obj, schema=self.prompts[id].schema)
+            except Exception as e:
+                raise RuntimeError(f"Schema validation fail: {e}")
+
+        return obj
+

trustgraph-flow/trustgraph/model/prompt/template/service.py

@@ -15,6 +15,7 @@ from .... schema import text_completion_response_queue
 from .... schema import prompt_request_queue, prompt_response_queue
 from .... base import ConsumerProducer
 from .... clients.llm_client import LlmClient
+from . prompt_manager import PromptConfiguration, Prompt, PromptManager
 
 from . prompts import to_definitions, to_relationships, to_rows
 from . prompts import to_kg_query, to_document_query, to_topics
@@ -29,6 +30,82 @@ class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
+        prompt_base = {}
+
+        # Parsing the prompt information to the prompt configuration
+        # structure
+        prompt_arg = params.get("prompt", [])
+        if prompt_arg:
+            for p in prompt_arg:
+                toks = p.split("=", 1)
+                if len(toks) < 2:
+                    raise RuntimeError(f"Prompt string not well-formed: {p}")
+                prompt_base[toks[0]] = {
+                    "template": toks[1]
+                }
+
+        prompt_response_type_arg = params.get("prompt_response_type", [])
+        if prompt_response_type_arg:
+            for p in prompt_response_type_arg:
+                toks = p.split("=", 1)
+                if len(toks) < 2:
+                    raise RuntimeError(f"Response type not well-formed: {p}")
+                if toks[0] not in prompt_base:
+                    raise RuntimeError(f"Response-type, {toks[0]} not known")
+                prompt_base[toks[0]]["response_type"] = toks[1]
+
+        prompt_schema_arg = params.get("prompt_schema", [])
+        if prompt_schema_arg:
+            for p in prompt_schema_arg:
+                toks = p.split("=", 1)
+                if len(toks) < 2:
+                    raise RuntimeError(f"Schema arg not well-formed: {p}")
+                if toks[0] not in prompt_base:
+                    raise RuntimeError(f"Schema, {toks[0]} not known")
+                try:
+                    prompt_base[toks[0]]["schema"] = json.loads(toks[1])
+                except:
+                    raise RuntimeError(f"Failed to parse JSON schema: {p}")
+
+        prompt_term_arg = params.get("prompt_term", [])
+        if prompt_term_arg:
+            for p in prompt_term_arg:
+                toks = p.split("=", 1)
+                if len(toks) < 2:
+                    raise RuntimeError(f"Term arg not well-formed: {p}")
+                if toks[0] not in prompt_base:
+                    raise RuntimeError(f"Term, {toks[0]} not known")
+                kvtoks = toks[1].split(":", 1)
+                if len(kvtoks) < 2:
+                    raise RuntimeError(f"Term not well-formed: {toks[1]}")
+                k, v = kvtoks
+                if "terms" not in prompt_base[toks[0]]:
+                    prompt_base[toks[0]]["terms"] = {}
+                prompt_base[toks[0]]["terms"][k] = v
+
+        global_terms = {}
+
+        global_term_arg = params.get("global_term", [])
+        if global_term_arg:
+            for t in global_term_arg:
+                toks = t.split("=", 1)
+                if len(toks) < 2:
+                    raise RuntimeError(f"Global term arg not well-formed: {t}")
+                global_terms[toks[0]] = toks[1]
+
+        print(global_terms)
+
+        prompts = {
+            k: Prompt(**v)
+            for k, v in prompt_base.items()
+        }
+
+        prompt_configuration = PromptConfiguration(
+            system_template = params.get("system_prompt", ""),
+            global_terms = global_terms,
+            prompts = prompts
+        )
+
         input_queue = params.get("input_queue", default_input_queue)
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
@@ -64,12 +141,21 @@ class Processor(ConsumerProducer):
             pulsar_host = self.pulsar_host
         )
 
-        self.definition_template = definition_template
-        self.topic_template = topic_template
-        self.relationship_template = relationship_template
-        self.rows_template = rows_template
-        self.knowledge_query_template = knowledge_query_template
-        self.document_query_template = document_query_template
+        # System prompt hack
+        class Llm:
+            def __init__(self, llm):
+                self.llm = llm
+            def request(self, system, prompt):
+                print(system)
+                print(prompt, flush=True)
+                return self.llm.request(system + "\n\n" + prompt)
+
+        self.llm = Llm(self.llm)
+
+        self.manager = PromptManager(
+            llm = self.llm,
+            config = prompt_configuration,
+        )
 
     def parse_json(self, text):
         json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
@@ -90,44 +176,64 @@ class Processor(ConsumerProducer):
 
         id = msg.properties()["id"]
 
-        kind = v.kind
+        kind = v.id
 
-        print(f"Handling kind {kind}...", flush=True)
+        try:
+            input = {
+                k: json.loads(v)
+                for k, v in v.terms.items()
+            }
+            
+            print(f"Handling kind {kind}...", flush=True)
+            print(input, flush=True)
 
-        if kind == "extract-definitions":
+            resp = self.manager.invoke(kind, input)
 
-            self.handle_extract_definitions(id, v)
-            return
+            if isinstance(resp, str):
 
-        elif kind == "extract-topics":
+                print("Send text response...", flush=True)
+                print(resp, flush=True)
 
-            self.handle_extract_topics(id, v)
-            return
+                r = PromptResponse(
+                    text=resp,
+                    object=None,
+                    error=None,
+                )
 
-        elif kind == "extract-relationships":
+                self.producer.send(r, properties={"id": id})
 
-            self.handle_extract_relationships(id, v)
-            return
+                return
 
-        elif kind == "extract-rows":
+            else:
 
-            self.handle_extract_rows(id, v)
-            return
+                print("Send object response...", flush=True)
+                print(json.dumps(resp, indent=4), flush=True)
 
-        elif kind == "kg-prompt":
+                r = PromptResponse(
+                    text=None,
+                    object=json.dumps(resp),
+                    error=None,
+                )
 
-            self.handle_kg_prompt(id, v)
-            return
+                self.producer.send(r, properties={"id": id})
 
-        elif kind == "document-prompt":
+                return
+            
+        except Exception as e:
 
-            self.handle_document_prompt(id, v)
-            return
+            print(f"Exception: {e}")
 
-        else:
+            print("Send error response...", flush=True)
 
-            print("Invalid kind.", flush=True)
-            return
+            r = PromptResponse(
+                error=Error(
+                    type = "llm-error",
+                    message = str(e),
+                ),
+                response=None,
+            )
+
+            self.producer.send(r, properties={"id": id})
 
     def handle_extract_definitions(self, id, v):
 
@@ -482,39 +588,33 @@ class Processor(ConsumerProducer):
         )
 
         parser.add_argument(
-            '--definition-template',
-            required=True,
-            help=f'Definition extraction template',
+            '--prompt', nargs='*',
+            help=f'Prompt template form id=template',
         )
 
         parser.add_argument(
-            '--topic-template',
-            required=True,
-            help=f'Topic extraction template',
+            '--prompt-response-type', nargs='*',
+            help=f'Prompt response type, form id=json|text',
         )
 
         parser.add_argument(
-            '--rows-template',
-            required=True,
-            help=f'Rows extraction template',
+            '--prompt-term', nargs='*',
+            help=f'Prompt response type, form id=key:value',
         )
 
         parser.add_argument(
-            '--relationship-template',
-            required=True,
-            help=f'Relationship extraction template',
+            '--prompt-schema', nargs='*',
+            help=f'Prompt response schema, form id=schema',
         )
 
         parser.add_argument(
-            '--knowledge-query-template',
-            required=True,
-            help=f'Knowledge query template',
+            '--system-prompt',
+            help=f'System prompt template',
         )
 
         parser.add_argument(
-            '--document-query-template',
-            required=True,
-            help=f'Document query template',
+            '--global-term', nargs='+',
+            help=f'Global term, form key:value'
         )
 
 def run():