mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-18 09:41:03 +02:00
Defs extraction working?
This commit is contained in:
Cyber MacGeddon
parent
8b24fa4209
commit
4bbd050d05
4 changed files with 106 additions and 93 deletions
|
|
@ -5,84 +5,74 @@ get entity definitions which are output as graph edges along with
|
|||
entity/context definitions for embedding.
|
||||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import urllib.parse
|
||||
from pulsar.schema import JsonSchema
|
||||
import uuid
|
||||
|
||||
from .... schema import Chunk, Triple, Triples, Metadata, Value
|
||||
from .... schema import EntityContext, EntityContexts
|
||||
from .... schema import chunk_ingest_queue, triples_store_queue
|
||||
from .... schema import entity_contexts_ingest_queue
|
||||
from .... schema import prompt_request_queue
|
||||
from .... schema import prompt_response_queue
|
||||
from .... schema import PromptRequest, PromptResponse
|
||||
from .... log_level import LogLevel
|
||||
from .... clients.prompt_client import PromptClient
|
||||
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
|
||||
from .... base import ConsumerProducer
|
||||
|
||||
from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
|
||||
from .... base import ProducerSpec
|
||||
|
||||
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
|
||||
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
|
||||
SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
|
||||
|
||||
module = "kg-extract-definitions"
|
||||
default_ident = "kg-extract-definitions"
|
||||
|
||||
default_input_queue = chunk_ingest_queue
|
||||
default_output_queue = triples_store_queue
|
||||
default_entity_context_queue = entity_contexts_ingest_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
class Processor(FlowProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
ec_queue = params.get(
|
||||
"entity_context_queue",
|
||||
default_entity_context_queue
|
||||
)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
)
|
||||
pr_response_queue = params.get(
|
||||
"prompt_response_queue", prompt_response_queue
|
||||
)
|
||||
id = params.get("id")
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": Chunk,
|
||||
"output_schema": Triples,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"id": id,
|
||||
}
|
||||
)
|
||||
|
||||
self.ec_prod = self.client.create_producer(
|
||||
topic=ec_queue,
|
||||
schema=JsonSchema(EntityContexts),
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "input",
|
||||
schema = Chunk,
|
||||
handler = self.on_message
|
||||
)
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.info({
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"entity_context_queue": ec_queue,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": Chunk.__name__,
|
||||
"output_schema": Triples.__name__,
|
||||
"vector_schema": EntityContexts.__name__,
|
||||
})
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "prompt-request",
|
||||
schema = PromptRequest
|
||||
)
|
||||
)
|
||||
|
||||
self.prompt = PromptClient(
|
||||
pulsar_host=self.pulsar_host,
|
||||
pulsar_api_key=self.pulsar_api_key,
|
||||
input_queue=pr_request_queue,
|
||||
output_queue=pr_response_queue,
|
||||
subscriber = module + "-prompt",
|
||||
self.register_specification(
|
||||
SubscriberSpec(
|
||||
name = "prompt-response",
|
||||
schema = PromptResponse,
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "triples",
|
||||
schema = Triples
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "entity-contexts",
|
||||
schema = EntityContexts
|
||||
)
|
||||
)
|
||||
|
||||
def to_uri(self, text):
|
||||
|
|
@ -97,32 +87,71 @@ class Processor(ConsumerProducer):
|
|||
|
||||
return self.prompt.request_definitions(chunk)
|
||||
|
||||
async def emit_edges(self, metadata, triples):
|
||||
async def emit_triples(self, pub, metadata, triples):
|
||||
|
||||
t = Triples(
|
||||
metadata=metadata,
|
||||
triples=triples,
|
||||
)
|
||||
await self.send(t)
|
||||
await pub.send(t)
|
||||
|
||||
async def emit_ecs(self, metadata, entities):
|
||||
async def emit_ecs(self, pub, metadata, entities):
|
||||
|
||||
t = EntityContexts(
|
||||
metadata=metadata,
|
||||
entities=entities,
|
||||
)
|
||||
self.ec_prod.send(t)
|
||||
await pub.send(t)
|
||||
|
||||
async def handle(self, msg):
|
||||
async def on_message(self, msg, consumer, flow):
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
v = msg.value()
|
||||
print(f"Indexing {v.metadata.id}...", flush=True)
|
||||
|
||||
chunk = v.chunk.decode("utf-8")
|
||||
|
||||
print(chunk, flush=True)
|
||||
|
||||
try:
|
||||
|
||||
defs = self.get_definitions(chunk)
|
||||
q = await flow.consumer["prompt-response"].subscribe(id)
|
||||
|
||||
try:
|
||||
|
||||
await flow.producer["prompt-request"].send(
|
||||
PromptRequest(
|
||||
id="extract-definitions",
|
||||
terms={
|
||||
"text": json.dumps(chunk)
|
||||
},
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
|
||||
# FIXME: hard-coded?
|
||||
resp = await asyncio.wait_for(
|
||||
q.get(),
|
||||
timeout=600
|
||||
)
|
||||
|
||||
print("Response", resp, flush=True)
|
||||
|
||||
if resp.error is not None:
|
||||
print("Error:", resp.error.message, flush=True)
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
if resp.object is None:
|
||||
raise RuntimeError("Expecting object in prompt response")
|
||||
|
||||
defs = json.loads(resp.object)
|
||||
|
||||
except Exception as e:
|
||||
print("Prompt exception:", e, flush=True)
|
||||
raise e
|
||||
finally:
|
||||
await flow.consumer["prompt-response"].unsubscribe(id)
|
||||
|
||||
triples = []
|
||||
entities = []
|
||||
|
|
@ -134,8 +163,8 @@ class Processor(ConsumerProducer):
|
|||
|
||||
for defn in defs:
|
||||
|
||||
s = defn.name
|
||||
o = defn.definition
|
||||
s = defn["entity"]
|
||||
o = defn["definition"]
|
||||
|
||||
if s == "": continue
|
||||
if o == "": continue
|
||||
|
|
@ -166,13 +195,13 @@ class Processor(ConsumerProducer):
|
|||
|
||||
ec = EntityContext(
|
||||
entity=s_value,
|
||||
context=defn.definition,
|
||||
context=defn["definition"],
|
||||
)
|
||||
|
||||
entities.append(ec)
|
||||
|
||||
|
||||
await self.emit_edges(
|
||||
await self.emit_triples(
|
||||
flow.producer["triples"],
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
metadata=[],
|
||||
|
|
@ -183,6 +212,7 @@ class Processor(ConsumerProducer):
|
|||
)
|
||||
|
||||
await self.emit_ecs(
|
||||
flow.producer["entity-contexts"],
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
metadata=[],
|
||||
|
|
@ -200,30 +230,9 @@ class Processor(ConsumerProducer):
|
|||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
ConsumerProducer.add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-e', '--entity-context-queue',
|
||||
default=default_entity_context_queue,
|
||||
help=f'Entity context queue (default: {default_entity_context_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
help=f'Prompt request queue (default: {prompt_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-completion-response-queue',
|
||||
default=prompt_response_queue,
|
||||
help=f'Prompt response queue (default: {prompt_response_queue})',
|
||||
)
|
||||
FlowProcessor.add_args(parser)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.launch(module, __doc__)
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
|
|
@ -53,6 +53,8 @@ class PromptManager:
|
|||
|
||||
async def invoke(self, id, input, llm):
|
||||
|
||||
print("Invoke...", flush=True)
|
||||
|
||||
if id not in self.prompts:
|
||||
raise RuntimeError("ID invalid")
|
||||
|
||||
|
|
@ -76,12 +78,13 @@ class PromptManager:
|
|||
try:
|
||||
obj = self.parse_json(resp)
|
||||
except:
|
||||
print("Parse fail:", resp, flush=True)
|
||||
raise RuntimeError("JSON parse fail")
|
||||
|
||||
if self.prompts[id].schema:
|
||||
try:
|
||||
print(self.prompts[id].schema, flush=True)
|
||||
validate(instance=obj, schema=self.prompts[id].schema)
|
||||
print("Validated", flush=True)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Schema validation fail: {e}")
|
||||
|
||||
|
|
|
|||
|
|
@ -11,8 +11,9 @@ from .... schema import Definition, Relationship, Triple
|
|||
from .... schema import Topic
|
||||
from .... schema import PromptRequest, PromptResponse, Error
|
||||
from .... schema import TextCompletionRequest, TextCompletionResponse
|
||||
from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
|
||||
from .... base import ProducerSpec
|
||||
|
||||
from .... base import FlowProcessor
|
||||
from .... base import ProducerSpec, SubscriberSpec, ConsumerSpec
|
||||
|
||||
from . prompt_manager import PromptConfiguration, Prompt, PromptManager
|
||||
|
||||
|
|
@ -172,10 +173,11 @@ class Processor(FlowProcessor):
|
|||
finally:
|
||||
await flow.consumer["text-completion-response"].unsubscribe(id)
|
||||
|
||||
print(resp, flush=True)
|
||||
|
||||
if isinstance(resp, str):
|
||||
|
||||
print("Send text response...", flush=True)
|
||||
print(resp, flush=True)
|
||||
|
||||
r = PromptResponse(
|
||||
text=resp,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue