mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-12 06:42:11 +02:00
a9197d11ee
- Keeps processing in different flows separate so that data can go to different stores / collections etc. - Potentially supports different processing flows - Tidies the processing API with common base-classes for e.g. LLMs, and automatic configuration of 'clients' to use the right queue names in a flow
92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
|
|
import ibis
|
|
import json
|
|
from jsonschema import validate
|
|
import re
|
|
|
|
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, config):
|
|
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 parse_json(self, text):
|
|
json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
|
|
|
|
if json_match:
|
|
json_str = json_match.group(1).strip()
|
|
else:
|
|
# If no delimiters, assume the entire output is JSON
|
|
json_str = text.strip()
|
|
|
|
return json.loads(json_str)
|
|
|
|
async def invoke(self, id, input, llm):
|
|
|
|
print("Invoke...", flush=True)
|
|
|
|
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 = await llm(**prompt)
|
|
|
|
if resp_type == "text":
|
|
return resp
|
|
|
|
if resp_type != "json":
|
|
raise RuntimeError(f"Response type {resp_type} not known")
|
|
|
|
try:
|
|
obj = self.parse_json(resp)
|
|
except:
|
|
print("Parse fail:", resp, flush=True)
|
|
raise RuntimeError("JSON parse fail")
|
|
|
|
if self.prompts[id].schema:
|
|
try:
|
|
validate(instance=obj, schema=self.prompts[id].schema)
|
|
print("Validated", flush=True)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Schema validation fail: {e}")
|
|
|
|
return obj
|
|
|