Defs extraction working?

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
 
-
-

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 PromptRequest, PromptResponse
-from .... schema import entity_contexts_ingest_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 .... 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
+class Processor(FlowProcessor):
-default_output_queue = triples_store_queue
-default_entity_context_queue = entity_contexts_ingest_queue
-default_subscriber = module
-
-class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
-        input_queue = params.get("input_queue", default_input_queue)
+        id = params.get("id")
-        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
-        )
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
+                "id": id,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": Chunk,
-                "output_schema": Triples,
-                "prompt_request_queue": pr_request_queue,
-                "prompt_response_queue": pr_response_queue,
             }
         )
 
-        self.ec_prod = self.client.create_producer(
+        self.register_specification(
-            topic=ec_queue,
+            ConsumerSpec(
-            schema=JsonSchema(EntityContexts),
+                name = "input",
+                schema = Chunk,
+                handler = self.on_message
+            )
         )
 
-        __class__.pubsub_metric.info({
+        self.register_specification(
-            "input_queue": input_queue,
+            ProducerSpec(
-            "output_queue": output_queue,
+                name = "prompt-request",
-            "entity_context_queue": ec_queue,
+                schema = PromptRequest
-            "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.prompt = PromptClient(
+        self.register_specification(
-            pulsar_host=self.pulsar_host,
+            SubscriberSpec(
-            pulsar_api_key=self.pulsar_api_key,
+                name = "prompt-response",
-            input_queue=pr_request_queue,
+                schema = PromptResponse,
-            output_queue=pr_response_queue,
+            )
-            subscriber = module + "-prompt",
+        )
+
+        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
+                s = defn["entity"]
-                o = defn.definition
+                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(
+        FlowProcessor.add_args(parser)
-            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})',
-        )
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)
 

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}")
 

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,