From f184f11041674c44eb3ab78a090583a9b835965f Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 4 Jan 2025 19:24:10 +0000 Subject: [PATCH] - Rename embeddings-vectorize to graph-embeddings - Added document-embeddings processor (broken, needs fixing) --- templates/components/trustgraph.jsonnet | 8 +- trustgraph-flow/scripts/embeddings-vectorize | 6 - trustgraph-flow/setup.py | 19 +-- .../document_embeddings/__init__.py | 3 + .../__main__.py | 2 +- .../embeddings.py} | 5 +- .../embeddings/graph_embeddings/__init__.py | 3 + .../embeddings/graph_embeddings/__main__.py | 6 + .../embeddings/graph_embeddings/embeddings.py | 118 ++++++++++++++++++ .../embeddings/vectorize/__init__.py | 3 - 10 files changed, 148 insertions(+), 25 deletions(-) delete mode 100755 trustgraph-flow/scripts/embeddings-vectorize create mode 100644 trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py rename trustgraph-flow/trustgraph/embeddings/{vectorize => document_embeddings}/__main__.py (57%) rename trustgraph-flow/trustgraph/embeddings/{vectorize/vectorize.py => document_embeddings/embeddings.py} (95%) create mode 100644 trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py create mode 100755 trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py create mode 100755 trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py delete mode 100644 trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 31ae420e..06bc409b 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -119,15 +119,15 @@ local prompt = import "prompt-template.jsonnet"; }, - "vectorize" +: { + "graph-embeddings" +: { create:: function(engine) local container = - engine.container("vectorize") + engine.container("graph-embeddings") .with_image(images.trustgraph) .with_command([ - "embeddings-vectorize", + "graph-embeddings", "-p", url.pulsar, ]) @@ -135,7 +135,7 @@ local prompt = import "prompt-template.jsonnet"; .with_reservations("0.5", "512M"); local containerSet = engine.containers( - "vectorize", [ container ] + "graph-embeddings", [ container ] ); local service = diff --git a/trustgraph-flow/scripts/embeddings-vectorize b/trustgraph-flow/scripts/embeddings-vectorize deleted file mode 100755 index 3de1e3a9..00000000 --- a/trustgraph-flow/scripts/embeddings-vectorize +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env python3 - -from trustgraph.embeddings.vectorize import run - -run() - diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index b7eab434..83e4f4f7 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -63,29 +63,30 @@ setuptools.setup( "falkordb", ], scripts=[ - "scripts/api-gateway", "scripts/agent-manager-react", + "scripts/api-gateway", "scripts/chunker-recursive", "scripts/chunker-token", "scripts/de-query-milvus", - "scripts/de-query-qdrant", "scripts/de-query-pinecone", + "scripts/de-query-qdrant", "scripts/de-write-milvus", - "scripts/de-write-qdrant", "scripts/de-write-pinecone", + "scripts/de-write-qdrant", + "scripts/document-embeddings", "scripts/document-rag", "scripts/embeddings-ollama", - "scripts/embeddings-vectorize", "scripts/ge-query-milvus", "scripts/ge-query-pinecone", "scripts/ge-query-qdrant", "scripts/ge-write-milvus", "scripts/ge-write-pinecone", "scripts/ge-write-qdrant", + "scripts/graph-embeddings", "scripts/graph-rag", "scripts/kg-extract-definitions", - "scripts/kg-extract-topics", "scripts/kg-extract-relationships", + "scripts/kg-extract-topics", "scripts/metering", "scripts/object-extract-row", "scripts/oe-write-milvus", @@ -103,13 +104,13 @@ setuptools.setup( "scripts/text-completion-ollama", "scripts/text-completion-openai", "scripts/triples-query-cassandra", - "scripts/triples-query-neo4j", - "scripts/triples-query-memgraph", "scripts/triples-query-falkordb", + "scripts/triples-query-memgraph", + "scripts/triples-query-neo4j", "scripts/triples-write-cassandra", - "scripts/triples-write-neo4j", - "scripts/triples-write-memgraph", "scripts/triples-write-falkordb", + "scripts/triples-write-memgraph", + "scripts/triples-write-neo4j", "scripts/wikipedia-lookup", ] ) diff --git a/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py new file mode 100644 index 00000000..40d505a5 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py @@ -0,0 +1,3 @@ + +from . embeddings import * + diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py similarity index 57% rename from trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py rename to trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py index a578de8a..a48cc4d0 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py @@ -1,5 +1,5 @@ -from . vectorize import run +from . embeddings import run if __name__ == '__main__': run() diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py similarity index 95% rename from trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py rename to trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py index 5630a7b5..e16f8ddd 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py @@ -1,7 +1,8 @@ """ -Vectorizer, calls the embeddings service to get embeddings for a chunk. -Input is text chunk, output is chunk and vectors. +Document embeddings, calls the embeddings service to get embeddings for a +chunk of text. Input is chunk of text plus metadata. +Output is chunk plus embedding. """ from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py new file mode 100644 index 00000000..40d505a5 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py @@ -0,0 +1,3 @@ + +from . embeddings import * + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py new file mode 100755 index 00000000..a48cc4d0 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py @@ -0,0 +1,6 @@ + +from . embeddings import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py new file mode 100755 index 00000000..09d1fda4 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py @@ -0,0 +1,118 @@ + +""" +Graph embeddings, calls the embeddings service to get embeddings for a +set of entity contexts. Input is entity plus textual context. +Output is entity plus embedding. +""" + +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 +from ... base import ConsumerProducer + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = entity_contexts_ingest_queue +default_output_queue = graph_embeddings_store_queue +default_subscriber = module + +class Processor(ConsumerProducer): + + 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) + emb_request_queue = params.get( + "embeddings_request_queue", embeddings_request_queue + ) + emb_response_queue = params.get( + "embeddings_response_queue", embeddings_response_queue + ) + + super(Processor, self).__init__( + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "embeddings_request_queue": emb_request_queue, + "embeddings_response_queue": emb_response_queue, + "subscriber": subscriber, + "input_schema": EntityContexts, + "output_schema": GraphEmbeddings, + } + ) + + self.embeddings = EmbeddingsClient( + pulsar_host=self.pulsar_host, + input_queue=emb_request_queue, + output_queue=emb_response_queue, + subscriber=module + "-emb", + ) + + def emit(self, rec, vectors): + + r = GraphEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors) + self.producer.send(r) + + def handle(self, msg): + + v = msg.value() + print(f"Indexing {v.metadata.id}...", flush=True) + + entities = [] + + try: + + for entity in v.entities: + + vectors = self.embeddings.request(entity.context) + + entities.append( + EntityEmbeddings( + entity=entity.entity, + vectors=vectors + ) + ) + + r = GraphEmbeddings( + metadata=v.metadata, + entities=entities, + ) + + self.producer.send(r) + + except Exception as e: + print("Exception:", e, flush=True) + + # Retry + raise e + + print("Done.", flush=True) + + @staticmethod + def add_args(parser): + + ConsumerProducer.add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ) + + parser.add_argument( + '--embeddings-request-queue', + default=embeddings_request_queue, + help=f'Embeddings request queue (default: {embeddings_request_queue})', + ) + + parser.add_argument( + '--embeddings-response-queue', + default=embeddings_response_queue, + help=f'Embeddings request queue (default: {embeddings_response_queue})', + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py b/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py deleted file mode 100644 index 31596b8c..00000000 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ - -from . vectorize import * -