Definitions is working again

This commit is contained in:
Cyber MacGeddon 2024-10-26 15:58:42 +01:00
parent 8f1a44c06e
commit 5e86a3c3ae
4 changed files with 38 additions and 19 deletions

View file

@ -7,7 +7,13 @@ p = PromptClient(pulsar_host="pulsar://localhost:6650")
chunk = """I noticed a cat in my garden. It is a four-legged animal chunk = """I noticed a cat in my garden. It is a four-legged animal
which is a mammal and can be tame or wild. I wonder if it will be friends which is a mammal and can be tame or wild. I wonder if it will be friends
with me. I think the cat's name is Fred and it has 4 legs""" with me. I think the cat's name is Fred and it has 4 legs.
A cat is a small mammal.
A grapefruit is a citrus fruit.
"""
resp = p.request_definitions( resp = p.request_definitions(
chunk=chunk, chunk=chunk,

View file

@ -1,6 +1,7 @@
import _pulsar import _pulsar
import json import json
import dataclasses
from .. schema import PromptRequest, PromptResponse from .. schema import PromptRequest, PromptResponse
from .. schema import prompt_request_queue from .. schema import prompt_request_queue
@ -13,6 +14,11 @@ WARN=_pulsar.LoggerLevel.Warn
INFO=_pulsar.LoggerLevel.Info INFO=_pulsar.LoggerLevel.Info
DEBUG=_pulsar.LoggerLevel.Debug DEBUG=_pulsar.LoggerLevel.Debug
@dataclasses.dataclass
class Definition:
name: str
definition: str
class PromptClient(BaseClient): class PromptClient(BaseClient):
def __init__( def __init__(
@ -56,20 +62,25 @@ class PromptClient(BaseClient):
def request_definitions(self, chunk, timeout=300): def request_definitions(self, chunk, timeout=300):
return self.request( defs = self.request(
id="extract-definitions", id="extract-definitions",
terms={ terms={
"chunk": chunk "text": chunk
}, },
timeout=timeout timeout=timeout
) )
return [
Definition(name=d["entity"], definition=d["definition"])
for d in defs
]
def request_relationships(self, chunk, timeout=300): def request_relationships(self, chunk, timeout=300):
return self.call( return self.call(
id="extract-relationships", id="extract-relationships",
terms={ terms={
"chunk": chunk "text": chunk
}, },
timeout=timeout timeout=timeout
) )
@ -79,7 +90,7 @@ class PromptClient(BaseClient):
return self.call( return self.call(
id="extract-topics", id="extract-topics",
terms={ terms={
"chunk": chunk "text": chunk
}, },
timeout=timeout timeout=timeout
) )

View file

@ -2,6 +2,7 @@
import ibis import ibis
import json import json
from jsonschema import validate from jsonschema import validate
import re
from trustgraph.clients.llm_client import LlmClient from trustgraph.clients.llm_client import LlmClient
@ -42,6 +43,17 @@ class PromptManager:
if v.terms is None: if v.terms is None:
v.terms = {} 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)
def invoke(self, id, input): def invoke(self, id, input):
if id not in self.prompts: if id not in self.prompts:
@ -58,6 +70,8 @@ class PromptManager:
resp = self.llm.request(**prompt) resp = self.llm.request(**prompt)
print(resp, flush=True)
if resp_type == "text": if resp_type == "text":
return resp return resp
@ -65,12 +79,11 @@ class PromptManager:
raise RuntimeError(f"Response type {resp_type} not known") raise RuntimeError(f"Response type {resp_type} not known")
try: try:
obj = json.loads(resp) obj = self.parse_json(resp)
except: except:
raise RuntimeError("JSON parse fail") raise RuntimeError("JSON parse fail")
print(resp) print(obj, flush=True)
print(obj)
if self.prompts[id].schema: if self.prompts[id].schema:
try: try:
print(self.prompts[id].schema) print(self.prompts[id].schema)

View file

@ -155,17 +155,6 @@ class Processor(ConsumerProducer):
config = prompt_configuration, config = prompt_configuration,
) )
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)
def handle(self, msg): def handle(self, msg):
v = msg.value() v = msg.value()