apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-16 16:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

- Rename embeddings-vectorize to graph-embeddings

Browse source 
- Added document-embeddings processor (broken, needs fixing)
This commit is contained in:
Cyber MacGeddon 2025-01-04 19:24:10 +00:00
parent 5227e3ac58
commit f184f11041
10 changed files with 148 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

8
templates/components/trustgraph.jsonnet
View file

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

6
trustgraph-flow/scripts/embeddings-vectorize
View file

@ -1,6 +0,0 @@
#!/usr/bin/env python3
from trustgraph.embeddings.vectorize import run
run()

19
trustgraph-flow/setup.py
View file

@ -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",
]
)

3
trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from . embeddings import *

2
trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py → trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py
View file

@ -1,5 +1,5 @@
from . vectorize import run
from . embeddings import run
if __name__ == '__main__':
run()

5
trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py → trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py
View file

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

3
trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from . embeddings import *

6
trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py Executable file
View file

@ -0,0 +1,6 @@
from . embeddings import run
if __name__ == '__main__':
run()

118
trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py Executable file
View file

@ -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__)

3
trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py
View file

@ -1,3 +0,0 @@
from . vectorize import *