Breakout store queries (#8)

- Break out store queries, so not locked into a Milvus/Cassandra backend
- Break out prompting into a separate module, so that prompts can be tailored to other LLMs
- Jsonnet used to generate docker compose templates
- Version to 0.6.0

trustgraph/model/prompt/generic/__init__.py

@@ -0,0 +1,3 @@
+
+from . service import *
+

trustgraph/model/prompt/generic/__main__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+if __name__ == '__main__':
+    run()
+

trustgraph/model/prompt/generic/prompts.py

@@ -0,0 +1,81 @@
+
+def to_relationships(text):
+
+    prompt = f"""<instructions>
+Study the following text and derive entity relationships.  For each
+relationship, derive the subject, predicate and object of the relationship.
+Output relationships in JSON format as an arary of objects with fields:
+- subject: the subject of the relationship
+- predicate: the predicate
+- object: the object of the relationship
+- object-entity: false if the object is a simple data type: name, value or date.  true if it is an entity.
+</instructions>
+
+<text>
+{text}
+</text>
+
+<requirements>
+You will respond only with raw JSON format data. Do not provide
+explanations. Do not use special characters in the abstract text. The
+abstract must be written as plain text.  Do not add markdown formatting
+or headers or prefixes.
+</requirements>"""
+    
+    return prompt
+    
+def to_definitions(text):
+
+    prompt = f"""<instructions>
+Study the following text and derive definitions for any discovered entities.
+Do not provide definitions for entities whose definitions are incomplete
+or unknown.
+Output relationships in JSON format as an arary of objects with fields:
+- entity: the name of the entity
+- definition: English text which defines the entity
+</instructions>
+
+<text>
+{text}
+</text>
+
+<requirements>
+You will respond only with raw JSON format data. Do not provide
+explanations. Do not use special characters in the abstract text. The
+abstract will be written as plain text.  Do not add markdown formatting
+or headers or prefixes.  Do not include null or unknown definitions.
+</requirements>"""
+    
+    return prompt
+
+def get_cypher(kg):
+
+    sg2 = []
+
+    for f in kg:
+
+        print(f)
+
+        sg2.append(f"({f.s})-[{f.p}]->({f.o})")
+
+    print(sg2)
+
+    kg = "\n".join(sg2)
+    kg = kg.replace("\\", "-")
+
+    return kg
+
+def to_kg_query(query, kg):
+
+    cypher =  get_cypher(kg)
+
+    prompt=f"""Study the following set of knowledge statements. The statements are written in Cypher format that has been extracted from a knowledge graph. Use only the provided set of knowledge statements in your response. Do not speculate if the answer is not found in the provided set of knowledge statements.
+
+Here's the knowledge statements:
+{cypher}
+
+Use only the provided knowledge statements to respond to the following:
+{query}
+"""
+
+    return prompt

trustgraph/model/prompt/generic/service.py

@@ -0,0 +1,195 @@
+
+"""
+Language service abstracts prompt engineering from LLM.
+"""
+
+import json
+
+from .... schema import Definition, Relationship, Triple
+from .... schema import PromptRequest, PromptResponse
+from .... schema import TextCompletionRequest, TextCompletionResponse
+from .... schema import text_completion_request_queue
+from .... schema import text_completion_response_queue
+from .... schema import prompt_request_queue, prompt_response_queue
+from .... base import ConsumerProducer
+from .... llm_client import LlmClient
+
+from . prompts import to_definitions, to_relationships, to_kg_query
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = prompt_request_queue
+default_output_queue = prompt_response_queue
+default_subscriber = module
+
+class Processor(ConsumerProducer):
+
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        tc_request_queue = params.get(
+            "text_completion_request_queue", text_completion_request_queue
+        )
+        tc_response_queue = params.get(
+            "text_completion_response_queue", text_completion_response_queue
+        )
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": PromptRequest,
+                "output_schema": PromptResponse,
+                "text_completion_request_queue": tc_request_queue,
+                "text_completion_response_queue": tc_response_queue,
+            }
+        )
+
+        self.llm = LlmClient(
+            subscriber=subscriber,
+            input_queue=tc_request_queue,
+            output_queue=tc_response_queue,
+            pulsar_host = self.pulsar_host
+        )
+
+    def handle(self, msg):
+
+        v = msg.value()
+
+        # Sender-produced ID
+
+        id = msg.properties()["id"]
+
+        kind = v.kind
+
+        print(f"Handling kind {kind}...", flush=True)
+
+        if kind == "extract-definitions":
+
+            self.handle_extract_definitions(id, v)
+            return
+
+        elif kind == "extract-relationships":
+
+            self.handle_extract_relationships(id, v)
+            return
+
+        elif kind == "kg-prompt":
+
+            self.handle_kg_prompt(id, v)
+            return
+
+        else:
+
+            print("Invalid kind.", flush=True)
+            return
+
+    def handle_extract_definitions(self, id, v):
+
+        prompt = to_definitions(v.chunk)
+
+        print(prompt)
+
+        ans = self.llm.request(prompt)
+
+        print(ans)
+
+        defs = json.loads(ans)
+
+        output = []
+
+        for defn in defs:
+
+            try:
+                e = defn["entity"]
+                d = defn["definition"]
+
+                output.append(
+                    Definition(
+                        name=e, definition=d
+                    )
+                )
+
+            except:
+                pass
+
+        print("Send response...", flush=True)
+        r = PromptResponse(definitions=output)
+        self.producer.send(r, properties={"id": id})
+
+        print("Done.", flush=True)
+        
+    def handle_extract_relationships(self, id, v):
+
+        prompt = to_relationships(v.chunk)
+
+        ans = self.llm.request(prompt)
+
+        defs = json.loads(ans)
+
+        output = []
+
+        for defn in defs:
+
+            try:
+                output.append(
+                    Relationship(
+                        s = defn["subject"],
+                        p = defn["predicate"],
+                        o = defn["object"],
+                        o_entity = defn["object-entity"],
+                    )
+                )
+
+            except Exception as e:
+                print(e)
+
+        print("Send response...", flush=True)
+        r = PromptResponse(relationships=output)
+        self.producer.send(r, properties={"id": id})
+
+        print("Done.", flush=True)
+        
+    def handle_kg_prompt(self, id, v):
+
+        prompt = to_kg_query(v.query, v.kg)
+
+        print(prompt)
+
+        ans = self.llm.request(prompt)
+
+        print(ans)
+
+        print("Send response...", flush=True)
+        r = PromptResponse(answer=ans)
+        self.producer.send(r, properties={"id": id})
+
+        print("Done.", flush=True)
+        
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '--text-completion-request-queue',
+            default=text_completion_request_queue,
+            help=f'Text completion request queue (default: {text_completion_request_queue})',
+        )
+
+        parser.add_argument(
+            '--text-completion-response-queue',
+            default=text_completion_response_queue,
+            help=f'Text completion response queue (default: {text_completion_response_queue})',
+        )
+
+def run():
+
+    Processor.start(module, __doc__)
+