diff --git a/trustgraph-base/trustgraph/schema/config.py b/trustgraph-base/trustgraph/schema/config.py index efe49182..3be63aa3 100644 --- a/trustgraph-base/trustgraph/schema/config.py +++ b/trustgraph-base/trustgraph/schema/config.py @@ -2,7 +2,7 @@ from pulsar.schema import Record, Bytes, String, Boolean, Array, Map, Integer from . topic import topic -from . types import Error, RowSchema +from . types import Error ############################################################################ diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index eeb3facd..45aae013 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -7,8 +7,6 @@ entity/context definitions for embedding. import json import urllib.parse -from pulsar.schema import JsonSchema -import uuid from .... schema import Chunk, Triple, Triples, Metadata, Value from .... schema import EntityContext, EntityContexts @@ -75,10 +73,6 @@ class Processor(FlowProcessor): return uri - def get_definitions(self, chunk): - - return self.prompt.request_definitions(chunk) - async def emit_triples(self, pub, metadata, triples): t = Triples( @@ -97,8 +91,6 @@ class Processor(FlowProcessor): async def on_message(self, msg, consumer, flow): - id = str(uuid.uuid4()) - v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index 5d6e093a..7f0cee15 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -5,59 +5,56 @@ relationship analysis to get entity relationship edges which are output as graph edges. """ +import json import urllib.parse from .... schema import Chunk, Triple, Triples from .... schema import Metadata, Value -from .... schema import chunk_ingest_queue, triples_store_queue -from .... schema import prompt_request_queue -from .... schema import prompt_response_queue -from .... log_level import LogLevel -from .... clients.prompt_client import PromptClient +from .... schema import PromptRequest, PromptResponse from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF -from .... base import ConsumerProducer + +from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec +from .... base import ProducerSpec RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True) SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True) -module = "kg-extract-relationships" +default_ident = "kg-extract-relationships" -default_input_queue = chunk_ingest_queue -default_output_queue = triples_store_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) - 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.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( + ConsumerSpec( + name = "input", + schema = Chunk, + handler = self.on_message + ) + ) + + self.register_specification( + RequestResponseSpec( + request_name = "prompt-request", + request_schema = PromptRequest, + response_name = "prompt-response", + response_schema = PromptResponse, + ) + ) + + self.register_specification( + ProducerSpec( + name = "triples", + schema = Triples + ) ) def to_uri(self, text): @@ -68,28 +65,49 @@ class Processor(ConsumerProducer): return uri - def get_relationships(self, chunk): - - return self.prompt.request_relationships(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 handle(self, msg): + async def on_message(self, msg, consumer, flow): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") + print(chunk, flush=True) + try: - rels = self.get_relationships(chunk) + try: + + resp = await flow("prompt-request").request( + PromptRequest( + id="extract-relationships", + terms={ + "text": json.dumps(chunk) + }, + ) + ) + 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") + + rels = json.loads(resp.object) + + except Exception as e: + print("Prompt exception:", e, flush=True) + raise e triples = [] @@ -100,9 +118,9 @@ class Processor(ConsumerProducer): for rel in rels: - s = rel.s - p = rel.p - o = rel.o + s = rel["s"] + p = rel["p"] + o = rel["o"] if s == "": continue if p == "": continue @@ -118,7 +136,7 @@ class Processor(ConsumerProducer): p_uri = self.to_uri(p) p_value = Value(value=str(p_uri), is_uri=True) - if rel.o_entity: + if rel["o_entity"]: o_uri = self.to_uri(o) o_value = Value(value=str(o_uri), is_uri=True) else: @@ -144,7 +162,7 @@ class Processor(ConsumerProducer): o=Value(value=str(p), is_uri=False) )) - if rel.o_entity: + if rel["o_entity"]: # Label for o triples.append(Triple( s=o_value, @@ -159,7 +177,7 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) - if rel.o_entity: + if rel["o_entity"]: # 'Subject of' for o triples.append(Triple( s=o_value, @@ -168,6 +186,7 @@ class Processor(ConsumerProducer): )) await self.emit_edges( + flow("triples"), Metadata( id=v.metadata.id, metadata=[], @@ -185,24 +204,9 @@ class Processor(ConsumerProducer): @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) - - parser.add_argument( - '--prompt-request-queue', - default=prompt_request_queue, - help=f'Prompt request queue (default: {prompt_request_queue})', - ) - - parser.add_argument( - '--prompt-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__)