diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 1006cd25..d8ae20df 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -6,9 +6,8 @@ from . producer import Producer from . publisher import Publisher from . subscriber import Subscriber from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics -from . flow_processor import FlowProcessor, SubscriberSpec, ProducerSpec +from . flow_processor import FlowProcessor from . flow_processor import ConsumerSpec, SettingSpec +from . flow_processor import SubscriberSpec, ProducerSpec from . request_response import RequestResponseService - - diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index fe4c58c6..cbf6a55b 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -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__) diff --git a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py b/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py index 7c765acc..c5c32395 100644 --- a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py @@ -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}") diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index f8df1bc9..ba72fdd3 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py @@ -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,