mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 09:11:03 +02:00
Prompt manager integrated and working with 6 tests
This commit is contained in:
Cyber MacGeddon
parent
51aef6c730
commit
3da63a38ce
13 changed files with 472 additions and 63 deletions
25
trustgraph-flow/trustgraph/model/prompt/template/README.md
Normal file
25
trustgraph-flow/trustgraph/model/prompt/template/README.md
Normal file
|
|
@ -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" ] } }'
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue