From c3904b6772c9727d8aeb1846adc4cd3fb8ffe0c3 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 30 Dec 2024 11:53:05 +0000 Subject: [PATCH] Core entity context flow in place --- .../trustgraph/schema/documents.py | 15 +----- trustgraph-base/trustgraph/schema/graph.py | 16 +++++++ .../embeddings/vectorize/vectorize.py | 40 +++++++++++----- .../extract/kg/definitions/extract.py | 43 ++++++++++++++--- .../extract/kg/relationships/extract.py | 47 ++----------------- .../trustgraph/extract/kg/topics/extract.py | 12 ++--- 6 files changed, 92 insertions(+), 81 deletions(-) diff --git a/trustgraph-base/trustgraph/schema/documents.py b/trustgraph-base/trustgraph/schema/documents.py index 79173126..38add83d 100644 --- a/trustgraph-base/trustgraph/schema/documents.py +++ b/trustgraph-base/trustgraph/schema/documents.py @@ -35,20 +35,6 @@ chunk_ingest_queue = topic('chunk-load') ############################################################################ -# Chunk embeddings are an embeddings associated with a text chunk - -class EntityContext(Record): - entity = Value() - context = String() - -class EntityContexts(Record): - metadata = Metadata() - entities = Array(EntityContext()) - -entity_contexts_ingest_queue = topic('entity-contexts-load') - -############################################################################ - # Doc embeddings query class DocumentEmbeddingsRequest(Record): @@ -65,3 +51,4 @@ document_embeddings_request_queue = topic( document_embeddings_response_queue = topic( 'doc-embeddings', kind='non-persistent', namespace='response', ) + diff --git a/trustgraph-base/trustgraph/schema/graph.py b/trustgraph-base/trustgraph/schema/graph.py index ea4867c6..7c304e1d 100644 --- a/trustgraph-base/trustgraph/schema/graph.py +++ b/trustgraph-base/trustgraph/schema/graph.py @@ -7,12 +7,28 @@ from . metadata import Metadata ############################################################################ +# Entity context are an entity associated with textual context + +class EntityContext(Record): + entity = Value() + context = String() + +# This is a 'batching' mechanism for the above data +class EntityContexts(Record): + metadata = Metadata() + entities = Array(EntityContext()) + +entity_contexts_ingest_queue = topic('entity-contexts-load') + +############################################################################ + # Graph embeddings are embeddings associated with a graph entity class EntityEmbeddings(Record): entity = Value() vectors = Array(Array(Double())) +# This is a 'batching' mechanism for the above data class GraphEmbeddings(Record): metadata = Metadata() entities = Array(EntityEmbeddings()) diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py index 4cf2af05..263ec329 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py @@ -4,8 +4,9 @@ Vectorizer, calls the embeddings service to get embeddings for a chunk. Input is text chunk, output is chunk and vectors. """ -from ... schema import Chunk, ChunkEmbeddings -from ... schema import chunk_ingest_queue, chunk_embeddings_ingest_queue +from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings +from ... schema import entity_contexts_ingest_queue +from ... schema import graph_embeddings_store_queue from ... schema import embeddings_request_queue, embeddings_response_queue from ... clients.embeddings_client import EmbeddingsClient from ... log_level import LogLevel @@ -13,8 +14,8 @@ from ... base import ConsumerProducer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_ingest_queue -default_output_queue = chunk_embeddings_ingest_queue +default_input_queue = entity_contexts_ingest_queue +default_output_queue = graph_embeddings_store_queue default_subscriber = module class Processor(ConsumerProducer): @@ -38,8 +39,8 @@ class Processor(ConsumerProducer): "embeddings_request_queue": emb_request_queue, "embeddings_response_queue": emb_response_queue, "subscriber": subscriber, - "input_schema": Chunk, - "output_schema": ChunkEmbeddings, + "input_schema": EntityContexts, + "output_schema": GraphEmbeddings, } ) @@ -50,9 +51,9 @@ class Processor(ConsumerProducer): subscriber=module + "-emb", ) - def emit(self, metadata, chunk, vectors): + def emit(self, rec, vectors): - r = ChunkEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors) + r = GraphEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors) self.producer.send(r) def handle(self, msg): @@ -60,21 +61,34 @@ class Processor(ConsumerProducer): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) - chunk = v.chunk.decode("utf-8") + entities = [] try: - vectors = self.embeddings.request(chunk) + for entity in v.entities: - self.emit( + vectors = self.embeddings.request(entity.context) + + entities.append( + EntityEmbeddings( + entity=entity.entity, + vectors=vectors + ) + ) + + r = GraphEmbeddings( metadata=v.metadata, - chunk=chunk.encode("utf-8"), - vectors=vectors + entities=entiities, ) + self.producer.send(r) + except Exception as e: print("Exception:", e, flush=True) + # Retry + raise e + print("Done.", flush=True) @staticmethod diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index eed34574..dd38029f 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -1,14 +1,17 @@ """ -Simple decoder, accepts embeddings+text chunks input, applies entity analysis to -get entity definitions which are output as graph edges. +Simple decoder, accepts text chunks input, applies entity analysis to +get entity definitions which are output as graph edges along with +entity/context definitions for embedding. """ import urllib.parse import json -from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue +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 .... log_level import LogLevel @@ -22,8 +25,9 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +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): @@ -32,6 +36,10 @@ class Processor(ConsumerProducer): 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 @@ -45,13 +53,30 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "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( + topic=ec_queue, + schema=JsonSchema(EntityContexts), + ) + + __class__.pubsub_metric.info({ + "input_queue": input_queue, + "output_queue": output_queue, + "vector_queue": vector_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.prompt = PromptClient( pulsar_host=self.pulsar_host, input_queue=pr_request_queue, @@ -152,6 +177,12 @@ class Processor(ConsumerProducer): 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, diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index d2dea062..0ed0d1ae 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -1,6 +1,6 @@ """ -Simple decoder, accepts vector+text chunks input, applies entity +Simple decoder, accepts text chunks input, applies entity relationship analysis to get entity relationship edges which are output as graph edges. """ @@ -9,10 +9,9 @@ import urllib.parse import os from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Triple, Triples, GraphEmbeddings +from .... schema import Chunk, Triple, Triples from .... schema import Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue -from .... schema import graph_embeddings_store_queue +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 @@ -25,9 +24,8 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = chunk_ingest_queue default_output_queue = triples_store_queue -default_vector_queue = graph_embeddings_store_queue default_subscriber = module class Processor(ConsumerProducer): @@ -36,7 +34,6 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) - vector_queue = params.get("vector_queue", default_vector_queue) subscriber = params.get("subscriber", default_subscriber) pr_request_queue = params.get( "prompt_request_queue", prompt_request_queue @@ -50,30 +47,13 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": Chunk, "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } ) - self.vec_prod = self.client.create_producer( - topic=vector_queue, - schema=JsonSchema(GraphEmbeddings), - ) - - __class__.pubsub_metric.info({ - "input_queue": input_queue, - "output_queue": output_queue, - "vector_queue": vector_queue, - "prompt_request_queue": pr_request_queue, - "prompt_response_queue": pr_response_queue, - "subscriber": subscriber, - "input_schema": ChunkEmbeddings.__name__, - "output_schema": Triples.__name__, - "vector_schema": GraphEmbeddings.__name__, - }) - self.prompt = PromptClient( pulsar_host=self.pulsar_host, input_queue=pr_request_queue, @@ -101,11 +81,6 @@ class Processor(ConsumerProducer): ) self.producer.send(t) - def emit_vec(self, metadata, ent, vec): - - r = GraphEmbeddings(metadata=metadata, entity=ent, vectors=vec) - self.vec_prod.send(r) - def handle(self, msg): v = msg.value() @@ -193,12 +168,6 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) - self.emit_vec(v.metadata, s_value, v.vectors) - self.emit_vec(v.metadata, p_value, v.vectors) - - if rel.o_entity: - self.emit_vec(v.metadata, o_value, v.vectors) - self.emit_edges( Metadata( id=v.metadata.id, @@ -222,12 +191,6 @@ class Processor(ConsumerProducer): default_output_queue, ) - parser.add_argument( - '-c', '--vector-queue', - default=default_vector_queue, - help=f'Vector output queue (default: {default_vector_queue})' - ) - parser.add_argument( '--prompt-request-queue', default=prompt_request_queue, diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index 8dfc3e6e..9181ae2c 100755 --- a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py @@ -1,14 +1,14 @@ """ -Simple decoder, accepts embeddings+text chunks input, applies entity analysis to -get entity definitions which are output as graph edges. +Simple decoder, accepts text chunks input, applies entity analysis to +get topics which are output as graph edges. """ import urllib.parse import json -from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue +from .... schema import Chunk, Triple, Triples, 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 @@ -20,7 +20,7 @@ DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = chunk_ingest_queue default_output_queue = triples_store_queue default_subscriber = module @@ -43,7 +43,7 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": Chunk, "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue,