From 187b0e65811cd6c13961353f76f6480adde03c35 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sun, 29 Dec 2024 18:08:12 +0000 Subject: [PATCH 01/53] Fix/websocket capacity increase (#230) * Fix invalid variable name invocation * Fix error responses in websockets * Increase websocket limits to 50MB max message. Turn on Pulsar chunking by default. --- .../trustgraph/base/consumer_producer.py | 1 + trustgraph-base/trustgraph/base/producer.py | 1 + .../trustgraph/gateway/document_load.py | 15 +++++++-------- trustgraph-flow/trustgraph/gateway/publisher.py | 2 +- trustgraph-flow/trustgraph/gateway/requestor.py | 10 ++++++++-- trustgraph-flow/trustgraph/gateway/sender.py | 8 +++++++- trustgraph-flow/trustgraph/gateway/socket.py | 5 ++++- 7 files changed, 29 insertions(+), 13 deletions(-) diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py index cabb7525..31441cda 100644 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ b/trustgraph-base/trustgraph/base/consumer_producer.py @@ -66,6 +66,7 @@ class ConsumerProducer(BaseProcessor): self.producer = self.client.create_producer( topic=output_queue, schema=JsonSchema(output_schema), + chunking_enabled=True, ) self.consumer = self.client.subscribe( diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 27d693ee..84d7fc99 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/trustgraph-base/trustgraph/base/producer.py @@ -34,6 +34,7 @@ class Producer(BaseProcessor): self.producer = self.client.create_producer( topic=output_queue, schema=JsonSchema(output_schema), + chunking_enabled=True, ) def send(self, msg, properties={}): diff --git a/trustgraph-flow/trustgraph/gateway/document_load.py b/trustgraph-flow/trustgraph/gateway/document_load.py index 0fd9a0df..4a37ecb6 100644 --- a/trustgraph-flow/trustgraph/gateway/document_load.py +++ b/trustgraph-flow/trustgraph/gateway/document_load.py @@ -1,7 +1,7 @@ import base64 -from .. schema import Document +from .. schema import Document, Metadata from .. schema import document_ingest_queue from . sender import ServiceSender @@ -18,25 +18,24 @@ class DocumentLoadSender(ServiceSender): def to_request(self, body): - if "metadata" in data: - metadata = to_subgraph(data["metadata"]) + if "metadata" in body: + metadata = to_subgraph(body["metadata"]) else: metadata = [] # Doing a base64 decoe/encode here to make sure the # content is valid base64 - doc = base64.b64decode(data["data"]) + doc = base64.b64decode(body["data"]) print("Document received") return Document( metadata=Metadata( - id=data.get("id"), + id=body.get("id"), metadata=metadata, - user=data.get("user", "trustgraph"), - collection=data.get("collection", "default"), + user=body.get("user", "trustgraph"), + collection=body.get("collection", "default"), ), data=base64.b64encode(doc).decode("utf-8") ) - diff --git a/trustgraph-flow/trustgraph/gateway/publisher.py b/trustgraph-flow/trustgraph/gateway/publisher.py index 89c612ce..e3298a47 100644 --- a/trustgraph-flow/trustgraph/gateway/publisher.py +++ b/trustgraph-flow/trustgraph/gateway/publisher.py @@ -7,7 +7,7 @@ import threading class Publisher: def __init__(self, pulsar_host, topic, schema=None, max_size=10, - chunking_enabled=False): + chunking_enabled=True): self.pulsar_host = pulsar_host self.topic = topic self.schema = schema diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index 5f6e2692..5bfedf1e 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -65,7 +65,10 @@ class ServiceRequestor: raise RuntimeError("Timeout") if resp.error: - return { "error": resp.error.message } + err = { "error": resp.error.message } + if responder: + await responder(err, True) + return err resp, fin = self.from_response(resp) @@ -81,7 +84,10 @@ class ServiceRequestor: logging.error(f"Exception: {e}") - return { "error": str(e) } + err = { "error": str(e) } + if responder: + await responder(err, True) + return err finally: self.sub.unsubscribe(id) diff --git a/trustgraph-flow/trustgraph/gateway/sender.py b/trustgraph-flow/trustgraph/gateway/sender.py index 93f1164c..c5bb2e17 100644 --- a/trustgraph-flow/trustgraph/gateway/sender.py +++ b/trustgraph-flow/trustgraph/gateway/sender.py @@ -46,5 +46,11 @@ class ServiceSender: logging.error(f"Exception: {e}") - return { "error": str(e) } + err = { "error": str(e) } + + if responder: + await responder(err, True) + + return err + diff --git a/trustgraph-flow/trustgraph/gateway/socket.py b/trustgraph-flow/trustgraph/gateway/socket.py index fd408d7b..4adc336f 100644 --- a/trustgraph-flow/trustgraph/gateway/socket.py +++ b/trustgraph-flow/trustgraph/gateway/socket.py @@ -44,7 +44,10 @@ class SocketEndpoint: return web.HTTPUnauthorized() running = Running() - ws = web.WebSocketResponse() + + # 50MB max message size + ws = web.WebSocketResponse(max_msg_size=52428800) + await ws.prepare(request) try: From 9a71ff65218a164dd818e4e74556103fb82f6e4e Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sun, 29 Dec 2024 20:02:12 +0000 Subject: [PATCH 02/53] Fix ingested metadata (#231) --- trustgraph-flow/trustgraph/gateway/text_load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trustgraph-flow/trustgraph/gateway/text_load.py b/trustgraph-flow/trustgraph/gateway/text_load.py index ade6b1c1..2499dd01 100644 --- a/trustgraph-flow/trustgraph/gateway/text_load.py +++ b/trustgraph-flow/trustgraph/gateway/text_load.py @@ -36,7 +36,7 @@ class TextLoadSender(ServiceSender): return TextDocument( metadata=Metadata( id=body.get("id"), - metabody=metadata, + metadata=metadata, user=body.get("user", "trustgraph"), collection=body.get("collection", "default"), ), From cb3272bba23a0ee1aab54405a1cad69b373b7000 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sun, 29 Dec 2024 20:21:16 +0000 Subject: [PATCH 03/53] Upgrade workbench (#232) --- templates/values/images.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 40954289..46ae47a6 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -14,5 +14,5 @@ local version = import "version.jsonnet"; memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.6", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.7", } From a0487fcb28b394b035e8d13ea1235ffa8105c201 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 30 Dec 2024 00:39:20 +0000 Subject: [PATCH 04/53] Workbench to 0.1.9 --- templates/values/images.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 46ae47a6..974624bc 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -14,5 +14,5 @@ local version = import "version.jsonnet"; memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.7", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.9", } From cff90cada1c972d5819d26a12396bc778a82951c Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 30 Dec 2024 10:44:33 +0000 Subject: [PATCH 05/53] Prepare for 0.19 --- .github/workflows/release.yaml | 2 +- trustgraph-bedrock/setup.py | 2 +- trustgraph-cli/setup.py | 2 +- trustgraph-embeddings-hf/setup.py | 4 ++-- trustgraph-flow/setup.py | 2 +- trustgraph-vertexai/setup.py | 2 +- trustgraph/setup.py | 12 ++++++------ 7 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 1b6dc177..7e18afb9 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,7 +5,7 @@ on: workflow_dispatch: push: tags: - - v0.18.* + - v0.19.* permissions: contents: read diff --git a/trustgraph-bedrock/setup.py b/trustgraph-bedrock/setup.py index b8dd36bd..09785c16 100644 --- a/trustgraph-bedrock/setup.py +++ b/trustgraph-bedrock/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", "pulsar-client", "prometheus-client", "boto3", diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 8217346f..a810c924 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", "requests", "pulsar-client", "rdflib", diff --git a/trustgraph-embeddings-hf/setup.py b/trustgraph-embeddings-hf/setup.py index 8febd59b..9af55352 100644 --- a/trustgraph-embeddings-hf/setup.py +++ b/trustgraph-embeddings-hf/setup.py @@ -34,8 +34,8 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", - "trustgraph-flow>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", + "trustgraph-flow>=0.19,<0.20", "torch", "urllib3", "transformers", diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 30ec0170..b7eab434 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", "urllib3", "rdflib", "pymilvus", diff --git a/trustgraph-vertexai/setup.py b/trustgraph-vertexai/setup.py index 7f9c2923..4bc2e716 100644 --- a/trustgraph-vertexai/setup.py +++ b/trustgraph-vertexai/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", "pulsar-client", "google-cloud-aiplatform", "prometheus-client", diff --git a/trustgraph/setup.py b/trustgraph/setup.py index a964ff06..0b14cfac 100644 --- a/trustgraph/setup.py +++ b/trustgraph/setup.py @@ -34,12 +34,12 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", - "trustgraph-bedrock>=0.18,<0.19", - "trustgraph-cli>=0.18,<0.19", - "trustgraph-embeddings-hf>=0.18,<0.19", - "trustgraph-flow>=0.18,<0.19", - "trustgraph-vertexai>=0.18,<0.19", + "trustgraph-base>=0.19,<0.20", + "trustgraph-bedrock>=0.19,<0.20", + "trustgraph-cli>=0.19,<0.20", + "trustgraph-embeddings-hf>=0.19,<0.20", + "trustgraph-flow>=0.19,<0.20", + "trustgraph-vertexai>=0.19,<0.20", ], scripts=[ ] From 9942f63773d82be013beab2e97a9d884393debb2 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 30 Dec 2024 10:51:35 +0000 Subject: [PATCH 06/53] tg-dump-msgpack previously output Python 'representation', now does JSON. (#233) --- trustgraph-cli/scripts/tg-dump-msgpack | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/trustgraph-cli/scripts/tg-dump-msgpack b/trustgraph-cli/scripts/tg-dump-msgpack index 18819649..f3b24d73 100755 --- a/trustgraph-cli/scripts/tg-dump-msgpack +++ b/trustgraph-cli/scripts/tg-dump-msgpack @@ -9,6 +9,7 @@ diagnostic utility. import msgpack import sys import argparse +import json def dump(input_file, action): @@ -17,7 +18,7 @@ def dump(input_file, action): unpacker = msgpack.Unpacker(f, raw=False) for unpacked in unpacker: - print(unpacked) + print(json.dumps(unpacked)) def summary(input_file, action): From a458d57af275d4076e66bfaebfe0d46b0ef9a1b3 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 30 Dec 2024 12:53:19 +0000 Subject: [PATCH 07/53] Feature/refactor entity embeddings (#235) * Make schema changes * Core entity context flow in place * extract-def outputs entity contexts * Refactored qdrant write * Refactoring of all vector stores in place --- .../trustgraph/schema/documents.py | 12 +- trustgraph-base/trustgraph/schema/graph.py | 23 +++- .../embeddings/vectorize/vectorize.py | 40 ++++--- .../extract/kg/definitions/extract.py | 72 ++++++++++-- .../extract/kg/relationships/extract.py | 49 +------- .../trustgraph/extract/kg/topics/extract.py | 12 +- .../storage/graph_embeddings/milvus/write.py | 8 +- .../graph_embeddings/pinecone/write.py | 107 ++++++++++-------- .../storage/graph_embeddings/qdrant/write.py | 76 +++++++------ 9 files changed, 230 insertions(+), 169 deletions(-) diff --git a/trustgraph-base/trustgraph/schema/documents.py b/trustgraph-base/trustgraph/schema/documents.py index 2a3d3d0c..38add83d 100644 --- a/trustgraph-base/trustgraph/schema/documents.py +++ b/trustgraph-base/trustgraph/schema/documents.py @@ -35,17 +35,6 @@ chunk_ingest_queue = topic('chunk-load') ############################################################################ -# Chunk embeddings are an embeddings associated with a text chunk - -class ChunkEmbeddings(Record): - metadata = Metadata() - vectors = Array(Array(Double())) - chunk = Bytes() - -chunk_embeddings_ingest_queue = topic('chunk-embeddings-load') - -############################################################################ - # Doc embeddings query class DocumentEmbeddingsRequest(Record): @@ -62,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 78c1a99c..7c304e1d 100644 --- a/trustgraph-base/trustgraph/schema/graph.py +++ b/trustgraph-base/trustgraph/schema/graph.py @@ -7,12 +7,31 @@ 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() - vectors = Array(Array(Double())) - entity = Value() + entities = Array(EntityEmbeddings()) graph_embeddings_store_queue = topic('graph-embeddings-store') diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py index 4cf2af05..5630a7b5 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=entities, ) + 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..dcb1123e 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 pulsar.schema import JsonSchema -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, + "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.prompt = PromptClient( pulsar_host=self.pulsar_host, input_queue=pr_request_queue, @@ -79,6 +104,14 @@ class Processor(ConsumerProducer): ) self.producer.send(t) + def emit_ecs(self, metadata, entities): + + t = EntityContexts( + metadata=metadata, + entities=entities, + ) + self.ec_prod.send(t) + def handle(self, msg): v = msg.value() @@ -91,6 +124,7 @@ class Processor(ConsumerProducer): defs = self.get_definitions(chunk) triples = [] + entities = [] # FIXME: Putting metadata into triples store is duplicated in # relationships extractor too @@ -129,6 +163,14 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) + ec = EntityContext( + entity=s_value, + context=defn.definition, + ) + + entities.append(ec) + + self.emit_edges( Metadata( id=v.metadata.id, @@ -139,6 +181,16 @@ class Processor(ConsumerProducer): triples ) + self.emit_ecs( + Metadata( + id=v.metadata.id, + metadata=[], + user=v.metadata.user, + collection=v.metadata.collection, + ), + entities + ) + except Exception as e: print("Exception: ", e, flush=True) @@ -152,6 +204,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..0fd7b9a8 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -1,18 +1,15 @@ """ -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. """ 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 +22,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 +32,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 +45,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 +79,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 +166,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 +189,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, diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py index 98fe7915..e1379577 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py @@ -38,9 +38,11 @@ class Processor(Consumer): v = msg.value() - if v.entity.value != "": - for vec in v.vectors: - self.vecstore.insert(vec, v.entity.value) + for entity in v.entities: + + if entity.entity.value != "" and entity.entity.value is not None: + for vec in entity.vectors: + self.vecstore.insert(vec, entity.entity.value) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index b918c10b..a32ff627 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -60,76 +60,83 @@ class Processor(Consumer): self.last_index_name = None + def create_index(self, index_name, dim): + + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) + + for i in range(0, 1000): + + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break + + time.sleep(1) + + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) + def handle(self, msg): v = msg.value() id = str(uuid.uuid4()) - if v.entity.value == "" or v.entity.value is None: return + for entity in v.entities: - for vec in v.vectors: + if entity.entity.value == "" or entity.entity.value is None: + continue - dim = len(vec) + for vec in entity.vectors: - index_name = ( - "t-" + v.metadata.user + "-" + str(dim) - ) + dim = len(vec) - if index_name != self.last_index_name: + index_name = ( + "t-" + v.metadata.user + "-" + str(dim) + ) - if not self.pinecone.has_index(index_name): + if index_name != self.last_index_name: - try: + if not self.pinecone.has_index(index_name): - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.create_index(index_name, dim) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + except Exception as e: + print("Pinecone index creation failed") + raise e - time.sleep(1) + print(f"Index {index_name} created", flush=True) - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + self.last_index_name = index_name - except Exception as e: - print("Pinecone index creation failed") - raise e + index = self.pinecone.Index(index_name) - print(f"Index {index_name} created", flush=True) + records = [ + { + "id": id, + "values": vec, + "metadata": { "entity": entity.entity.value }, + } + ] - self.last_index_name = index_name - - index = self.pinecone.Index(index_name) - - records = [ - { - "id": id, - "values": vec, - "metadata": { "entity": v.entity.value }, - } - ] - - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 47b53979..7bc5778c 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -40,49 +40,59 @@ class Processor(Consumer): self.client = QdrantClient(url=store_uri) + def get_collection(self, dim, user, collection): + + cname = ( + "t_" + user + "_" + collection + "_" + str(dim) + ) + + if cname != self.last_collection: + + if not self.client.collection_exists(cname): + + try: + self.client.create_collection( + collection_name=cname, + vectors_config=VectorParams( + size=dim, distance=Distance.COSINE + ), + ) + except Exception as e: + print("Qdrant collection creation failed") + raise e + + self.last_collection = cname + + return cname + def handle(self, msg): v = msg.value() - if v.entity.value == "" or v.entity.value is None: return + for entity in v.entities: - for vec in v.vectors: + if entity.entity.value == "" or entity.entity.value is None: return - dim = len(vec) - collection = ( - "t_" + v.metadata.user + "_" + v.metadata.collection + "_" + - str(dim) - ) + for vec in entity.vectors: - if collection != self.last_collection: + dim = len(vec) - if not self.client.collection_exists(collection): + collection = self.get_collection( + dim, v.metadata.user, v.metadata.collection + ) - try: - self.client.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.COSINE - ), + self.client.upsert( + collection_name=collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "entity": entity.entity.value, + } ) - except Exception as e: - print("Qdrant collection creation failed") - raise e - - self.last_collection = collection - - self.client.upsert( - collection_name=collection, - points=[ - PointStruct( - id=str(uuid.uuid4()), - vector=vec, - payload={ - "entity": v.entity.value, - } - ) - ] - ) + ] + ) @staticmethod def add_args(parser): From 44c0d6f347f348c2337c78a8de00a9e296bb4764 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Tue, 31 Dec 2024 19:58:31 +0000 Subject: [PATCH 08/53] Updated workbench to 0.2.0, search facility (#236) --- templates/values/images.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 974624bc..46c2bffc 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -14,5 +14,5 @@ local version = import "version.jsonnet"; memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.9", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.0", } From 44f8ce8834d1765c39bf3eab9ab85785c6270caa Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Thu, 2 Jan 2025 19:49:22 +0000 Subject: [PATCH 09/53] Migrate cli utils to REST API (#239) * Port a number of commands to use API gateway instead of Pulsar * Ported tg-invoke-agent to websockets API * Rename the 2 RAG commands: tg-query-... to tg-invoke-... --- trustgraph-cli/scripts/tg-graph-show | 29 +++--- trustgraph-cli/scripts/tg-graph-to-turtle | 58 +++++++---- trustgraph-cli/scripts/tg-invoke-agent | 93 ++++++++++++------ ...query-graph-rag => tg-invoke-document-rag} | 37 +++---- ...query-document-rag => tg-invoke-graph-rag} | 37 +++---- trustgraph-cli/scripts/tg-invoke-llm | 18 ++-- trustgraph-cli/scripts/tg-invoke-prompt | 18 ++-- trustgraph-cli/scripts/tg-load-pdf | 97 ++++--------------- trustgraph-cli/scripts/tg-load-text | 96 ++++-------------- trustgraph-cli/setup.py | 17 ++-- 10 files changed, 223 insertions(+), 277 deletions(-) rename trustgraph-cli/scripts/{tg-query-graph-rag => tg-invoke-document-rag} (51%) rename trustgraph-cli/scripts/{tg-query-document-rag => tg-invoke-graph-rag} (50%) diff --git a/trustgraph-cli/scripts/tg-graph-show b/trustgraph-cli/scripts/tg-graph-show index c09266fb..a3d10283 100755 --- a/trustgraph-cli/scripts/tg-graph-show +++ b/trustgraph-cli/scripts/tg-graph-show @@ -6,23 +6,23 @@ Connects to the graph query service and dumps all graph edges. import argparse import os -from trustgraph.clients.triples_query_client import TriplesQueryClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' -def show_graph(pulsar, user, collection): +def show_graph(url, user, collection): - tq = TriplesQueryClient(pulsar_host=pulsar) + api = Api(url) - rows = tq.request( - user=user, collection=collection, - s=None, p=None, o=None, limit=10_000_000 + rows = api.triples_query( +# user=user, collection=collection, + s=None, p=None, o=None, limit=10_000, ) for row in rows: - print(row.s.value, row.p.value, row.o.value) + print(row.s, row.p, row.o) def main(): @@ -32,19 +32,19 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -54,7 +54,8 @@ def main(): try: show_graph( - pulsar=args.pulsar_host, user=args.user, + url=args.api_url, + user=args.user, collection=args.collection, ) diff --git a/trustgraph-cli/scripts/tg-graph-to-turtle b/trustgraph-cli/scripts/tg-graph-to-turtle index 1d75478e..fc17ddd0 100755 --- a/trustgraph-cli/scripts/tg-graph-to-turtle +++ b/trustgraph-cli/scripts/tg-graph-to-turtle @@ -5,37 +5,45 @@ Connects to the graph query service and dumps all graph edges in Turtle format. """ -import argparse -import os -from trustgraph.clients.triples_query_client import TriplesQueryClient import rdflib import io import sys +import argparse +import os -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +from trustgraph.api import Api, Uri -def show_graph(pulsar): +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_user = 'trustgraph' +default_collection = 'default' - tq = TriplesQueryClient(pulsar_host=pulsar) +def show_graph(url, user, collection): - rows = tq.request(None, None, None, limit=10_000_000) + api = Api(url) + + rows = api.triples_query( + s=None, p=None, o=None, + limit=10_000) +# user=user, collection=collection, g = rdflib.Graph() for row in rows: - sv = rdflib.term.URIRef(row.s.value) - pv = rdflib.term.URIRef(row.p.value) + sv = rdflib.term.URIRef(row.s) + pv = rdflib.term.URIRef(row.p) - if row.o.is_uri: + if isinstance(row.o, Uri): # Skip malformed URLs with spaces in - if " " in row.o.value: + if " " in row.o: continue - ov = rdflib.term.URIRef(row.o.value) + ov = rdflib.term.URIRef(row.o) + else: - ov = rdflib.term.Literal(row.o.value) + + ov = rdflib.term.Literal(row.o) g.add((sv, pv, ov)) @@ -56,16 +64,32 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-U', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-C', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' ) args = parser.parse_args() try: - show_graph(args.pulsar_host) + show_graph( + url=args.api_url, + user=args.user, + collection=args.collection + ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-invoke-agent b/trustgraph-cli/scripts/tg-invoke-agent index 3f05071c..bb87c2a6 100755 --- a/trustgraph-cli/scripts/tg-invoke-agent +++ b/trustgraph-cli/scripts/tg-invoke-agent @@ -1,16 +1,18 @@ #!/usr/bin/env python3 """ -Uses the GraphRAG service to answer a query +Uses the GraphRAG service to answer a question """ import argparse import os import textwrap +import uuid +import asyncio +import json +from websockets.asyncio.client import connect -from trustgraph.clients.agent_client import AgentClient - -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -27,15 +29,18 @@ def output(text, prefix="> ", width=78): ) print(out) -def query( - pulsar_host, query, user, collection, +async def question( + url, question, user, collection, plan=None, state=None, verbose=False ): - am = AgentClient(pulsar_host=pulsar_host) + if not url.endswith("/"): + url += "/" + + url = url + "api/v1/socket" if verbose: - output(wrap(query), "\U00002753 ") + output(wrap(question), "\U00002753 ") print() def think(x): @@ -48,11 +53,43 @@ def query( output(wrap(x), "\U0001f4a1 ") print() - resp = am.request( - question=query, think=think, observe=observe, - ) + mid = str(uuid.uuid4()) - print(resp) + async with connect(url) as ws: + + req = json.dumps({ + "id": mid, + "service": "agent", + "request": { + "question": question, + } + + }) + + await ws.send(req) + + while True: + + msg = await ws.recv() + + obj = json.loads(msg) + + if obj["id"] != mid: + print("Ignore message") + continue + + if "thought" in obj["response"]: + think(obj["response"]["thought"]) + + if "observation" in obj["response"]: + observe(obj["response"]["observation"]) + + if "answer" in obj["response"]: + print(obj["response"]["answer"]) + + if obj["complete"]: break + + await ws.close() def main(): @@ -62,25 +99,25 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-q', '--query', + '-q', '--question', required=True, - help=f'Query to execute', + help=f'Question to answer', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -105,14 +142,16 @@ def main(): try: - query( - pulsar_host=args.pulsar_host, - query=args.query, - user=args.user, - collection=args.collection, - plan=args.plan, - state=args.state, - verbose=args.verbose, + asyncio.run( + question( + url=args.url, + question=args.question, + user=args.user, + collection=args.collection, + plan=args.plan, + state=args.state, + verbose=args.verbose, + ) ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-query-graph-rag b/trustgraph-cli/scripts/tg-invoke-document-rag similarity index 51% rename from trustgraph-cli/scripts/tg-query-graph-rag rename to trustgraph-cli/scripts/tg-invoke-document-rag index 8a865eea..3e0a9422 100755 --- a/trustgraph-cli/scripts/tg-query-graph-rag +++ b/trustgraph-cli/scripts/tg-invoke-document-rag @@ -1,50 +1,53 @@ #!/usr/bin/env python3 """ -Uses the GraphRAG service to answer a query +Uses the GraphRAG service to answer a question """ import argparse import os -from trustgraph.clients.graph_rag_client import GraphRagClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' -def query(pulsar_host, query, user, collection): +def question(url, question, user, collection): + + rag = Api(url) + +# user=user, collection=collection, + resp = rag.document_rag(question=question) - rag = GraphRagClient(pulsar_host=pulsar_host) - resp = rag.request(user=user, collection=collection, query=query) print(resp) def main(): parser = argparse.ArgumentParser( - prog='tg-graph-query-rag', + prog='tg-invoke-document-rag', description=__doc__, ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-q', '--query', + '-q', '--question', required=True, - help=f'Query to execute', + help=f'Question to answer', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -53,9 +56,9 @@ def main(): try: - query( - pulsar_host=args.pulsar_host, - query=args.query, + question( + url=args.url, + question=args.question, user=args.user, collection=args.collection, ) diff --git a/trustgraph-cli/scripts/tg-query-document-rag b/trustgraph-cli/scripts/tg-invoke-graph-rag similarity index 50% rename from trustgraph-cli/scripts/tg-query-document-rag rename to trustgraph-cli/scripts/tg-invoke-graph-rag index 8d800629..50de5b74 100755 --- a/trustgraph-cli/scripts/tg-query-document-rag +++ b/trustgraph-cli/scripts/tg-invoke-graph-rag @@ -1,50 +1,53 @@ #!/usr/bin/env python3 """ -Uses the Document RAG service to answer a query +Uses the GraphRAG service to answer a question """ import argparse import os -from trustgraph.clients.document_rag_client import DocumentRagClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' -def query(pulsar_host, query, user, collection): +def question(url, question, user, collection): + + rag = Api(url) + +# user=user, collection=collection, + resp = rag.graph_rag(question=question) - rag = DocumentRagClient(pulsar_host=pulsar) - resp = rag.request(user=user, collection=collection, query=query) print(resp) def main(): parser = argparse.ArgumentParser( - prog='tg-query-document-rag', + prog='tg-invoke-graph-rag', description=__doc__, ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-q', '--query', + '-q', '--question', required=True, - help=f'Query to execute', + help=f'Question to answer', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -53,9 +56,9 @@ def main(): try: - query( - pulsar_host=args.pulsar_host, - query=args.query, + question( + url=args.url, + question=args.question, user=args.user, collection=args.collection, ) diff --git a/trustgraph-cli/scripts/tg-invoke-llm b/trustgraph-cli/scripts/tg-invoke-llm index d7289b5f..d1a49e87 100755 --- a/trustgraph-cli/scripts/tg-invoke-llm +++ b/trustgraph-cli/scripts/tg-invoke-llm @@ -8,15 +8,15 @@ and user prompt. Both arguments are required. import argparse import os import json -from trustgraph.clients.llm_client import LlmClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') -def query(pulsar_host, system, prompt): +def query(url, system, prompt): - cli = LlmClient(pulsar_host=pulsar_host) + api = Api(url) - resp = cli.request(system=system, prompt=prompt) + resp = api.text_completion(system=system, prompt=prompt) print(resp) @@ -28,9 +28,9 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( @@ -50,7 +50,7 @@ def main(): try: query( - pulsar_host=args.pulsar_host, + url=args.url, system=args.system[0], prompt=args.prompt[0], ) diff --git a/trustgraph-cli/scripts/tg-invoke-prompt b/trustgraph-cli/scripts/tg-invoke-prompt index 19f30912..7d54dab0 100755 --- a/trustgraph-cli/scripts/tg-invoke-prompt +++ b/trustgraph-cli/scripts/tg-invoke-prompt @@ -12,15 +12,15 @@ using key=value arguments on the command line, and these replace import argparse import os import json -from trustgraph.clients.prompt_client import PromptClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') -def query(pulsar_host, template_id, variables): +def query(url, template_id, variables): - cli = PromptClient(pulsar_host=pulsar_host) + api = Api(url) - resp = cli.request(id=template_id, variables=variables) + resp = api.prompt(id=template_id, variables=variables) if isinstance(resp, str): print(resp) @@ -35,9 +35,9 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( @@ -70,7 +70,7 @@ specified multiple times''', try: query( - pulsar_host=args.pulsar_host, + url=args.url, template_id=args.id[0], variables=variables, ) diff --git a/trustgraph-cli/scripts/tg-load-pdf b/trustgraph-cli/scripts/tg-load-pdf index a0d2b3bc..f27aaf5a 100755 --- a/trustgraph-cli/scripts/tg-load-pdf +++ b/trustgraph-cli/scripts/tg-load-pdf @@ -6,21 +6,19 @@ Loads a PDF document into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -import base64 import hashlib import argparse import os import time import uuid -from trustgraph.schema import Document, document_ingest_queue -from trustgraph.schema import Metadata, Triple, Value -from trustgraph.log_level import LogLevel -from trustgraph.knowledge import hash, to_uri, Uri +from trustgraph.api import Api +from trustgraph.knowledge import hash, to_uri from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG from trustgraph.knowledge import Organization, PublicationEvent from trustgraph.knowledge import DigitalDocument +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -28,24 +26,13 @@ class Loader: def __init__( self, - pulsar_host, - output_queue, + url, user, collection, - log_level, metadata, ): - self.client = pulsar.Client( - pulsar_host, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(Document), - chunking_enabled=True, - ) + self.api = Api(url) self.user = user self.collection = collection @@ -68,49 +55,18 @@ class Loader: id = to_uri(PREF_DOC, id) - triples = [] - - def emit(t): - triples.append(t) - self.metadata.id = id - self.metadata.emit(emit) - r = Document( - metadata=Metadata( - id=id, - metadata=[ - Triple( - s=Value( - value=t["s"], - is_uri=isinstance(t["s"], Uri) - ), - p=Value( - value=t["p"], - is_uri=isinstance(t["p"], Uri) - ), - o=Value( - value=t["o"], - is_uri=isinstance(t["o"], Uri) - ), - ) - for t in triples - ], - user=self.user, - collection=self.collection, - ), - data=base64.b64encode(data), + self.api.load_document( + document=data, id=id, metadata=self.metadata, +# user=self.user, +# collection=self.collection, ) - self.producer.send(r) - print(f"{file}: Loaded successfully.") except Exception as e: print(f"{file}: Failed: {str(e)}", flush=True) - - def __del__(self): - self.client.close() def main(): @@ -119,29 +75,20 @@ def main(): description=__doc__, ) - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') - default_output_queue = document_ingest_queue - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) - - parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -183,7 +130,7 @@ def main(): ) parser.add_argument( - '--url', help=f'Document URL' + '--document-url', help=f'Document URL' ) parser.add_argument( @@ -194,14 +141,6 @@ def main(): '--identifier', '--id', help=f'Document ID' ) - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.ERROR, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - parser.add_argument( 'files', nargs='+', help=f'File to load' @@ -221,7 +160,7 @@ def main(): copyright_holder=args.copyright_holder, copyright_year=args.copyright_year, license=args.license, - url=args.url, + url=args.document_url, keywords=args.keyword, ) @@ -239,11 +178,9 @@ def main(): ) p = Loader( - pulsar_host=args.pulsar_host, - output_queue=args.output_queue, + url=args.url, user=args.user, collection=args.collection, - log_level=args.log_level, metadata=document, ) diff --git a/trustgraph-cli/scripts/tg-load-text b/trustgraph-cli/scripts/tg-load-text index 51664a1b..634ac3d7 100755 --- a/trustgraph-cli/scripts/tg-load-text +++ b/trustgraph-cli/scripts/tg-load-text @@ -12,14 +12,13 @@ import os import time import uuid -from trustgraph.schema import TextDocument, text_ingest_queue -from trustgraph.schema import Metadata, Triple, Value -from trustgraph.log_level import LogLevel -from trustgraph.knowledge import hash, to_uri, Literal, Uri +from trustgraph.api import Api +from trustgraph.knowledge import hash, to_uri from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG from trustgraph.knowledge import Organization, PublicationEvent from trustgraph.knowledge import DigitalDocument +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -27,24 +26,13 @@ class Loader: def __init__( self, - pulsar_host, - output_queue, + url, user, collection, - log_level, metadata, ): - self.client = pulsar.Client( - pulsar_host, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(TextDocument), - chunking_enabled=True, - ) + self.api = Api(url) self.user = user self.collection = collection @@ -67,49 +55,18 @@ class Loader: id = to_uri(PREF_DOC, id) - triples = [] - - def emit(t): - triples.append(t) - self.metadata.id = id - self.metadata.emit(emit) - r = TextDocument( - metadata=Metadata( - id=id, - metadata=[ - Triple( - s=Value( - value=t["s"], - is_uri=isinstance(t["s"], Uri) - ), - p=Value( - value=t["p"], - is_uri=isinstance(t["p"], Uri) - ), - o=Value( - value=t["o"], - is_uri=isinstance(t["o"], Uri) - ), - ) - for t in triples - ], - user=self.user, - collection=self.collection, - ), - text=data, + self.api.load_text( + text=data, id=id, metadata=self.metadata, +# user=self.user, +# collection=self.collection, ) - self.producer.send(r) - print(f"{file}: Loaded successfully.") except Exception as e: print(f"{file}: Failed: {str(e)}", flush=True) - - def __del__(self): - self.client.close() def main(): @@ -118,29 +75,20 @@ def main(): description=__doc__, ) - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') - default_output_queue = text_ingest_queue - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) - - parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -182,7 +130,7 @@ def main(): ) parser.add_argument( - '--url', help=f'Document URL' + '--document-url', help=f'Document URL' ) parser.add_argument( @@ -193,14 +141,6 @@ def main(): '--identifier', '--id', help=f'Document ID' ) - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.ERROR, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - parser.add_argument( 'files', nargs='+', help=f'File to load' @@ -220,7 +160,7 @@ def main(): copyright_holder=args.copyright_holder, copyright_year=args.copyright_year, license=args.license, - url=args.url, + url=args.document_url, keywords=args.keyword, ) @@ -238,11 +178,9 @@ def main(): ) p = Loader( - pulsar_host=args.pulsar_host, - output_queue=args.output_queue, + url=args.url, user=args.user, collection=args.collection, - log_level=args.log_level, metadata=document, ) diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index a810c924..67c70158 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -40,23 +40,24 @@ setuptools.setup( "rdflib", "tabulate", "msgpack", + "websockets", ], scripts=[ + "scripts/tg-dump-msgpack", "scripts/tg-graph-show", "scripts/tg-graph-to-turtle", + "scripts/tg-init-pulsar", "scripts/tg-init-pulsar-manager", + "scripts/tg-invoke-agent", + "scripts/tg-invoke-document-rag", + "scripts/tg-invoke-graph-rag", + "scripts/tg-invoke-llm", + "scripts/tg-invoke-prompt", + "scripts/tg-load-kg-core", "scripts/tg-load-pdf", "scripts/tg-load-text", "scripts/tg-load-turtle", - "scripts/tg-query-document-rag", - "scripts/tg-query-graph-rag", - "scripts/tg-init-pulsar", "scripts/tg-processor-state", - "scripts/tg-invoke-agent", - "scripts/tg-invoke-prompt", - "scripts/tg-invoke-llm", "scripts/tg-save-kg-core", - "scripts/tg-load-kg-core", - "scripts/tg-dump-msgpack", ] ) From fed7e4ca08221b7301776f2e08d024c98392f3e7 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Thu, 2 Jan 2025 19:50:23 +0000 Subject: [PATCH 10/53] Pulsar resilience refactor (#238) * Add userID control over compose containers * Split Pulsar out, non-standalone * Added bookie & zookeeper to Prometheus stats --- prometheus/prometheus.yml | 12 ++ templates/components/pulsar.jsonnet | 167 ++++++++++++++++++++---- templates/engine/docker-compose.jsonnet | 6 + 3 files changed, 162 insertions(+), 23 deletions(-) diff --git a/prometheus/prometheus.yml b/prometheus/prometheus.yml index 24102a23..785036c9 100644 --- a/prometheus/prometheus.yml +++ b/prometheus/prometheus.yml @@ -20,6 +20,18 @@ scrape_configs: - targets: - 'pulsar:8080' + - job_name: 'bookie' + scrape_interval: 5s + static_configs: + - targets: + - 'bookie:8000' + + - job_name: 'zookeeper' + scrape_interval: 5s + static_configs: + - targets: + - 'zookeeper:8000' + - job_name: 'pdf-decoder' scrape_interval: 5s static_configs: diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 0342b4d5..58170500 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -2,31 +2,112 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; +// This is a Pulsar configuration. Non-standalone mode so we deploy +// individual components: bookkeeper, broker and zookeeper. +// +// This also deploys the TrustGraph 'admin' container which initialises +// TrustGraph-specific namespaces etc. + { "pulsar" +: { create:: function(engine) -// local confVolume = engine.volume("pulsar-conf").with_size("2G"); - local dataVolume = engine.volume("pulsar-data").with_size("20G"); + // Zookeeper volume + local zkVolume = engine.volume("zookeeper").with_size("1G"); - local container = + // Zookeeper container + local zkContainer = + engine.container("zookeeper") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/zookeeper.conf && bin/generate-zookeeper-config.sh conf/zookeeper.conf && exec bin/pulsar zookeeper" + ]) + .with_limits("0.1", "400M") + .with_reservations("0.05", "400M") + .with_volume_mount(zkVolume, "/pulsar/data/zookeeper") + .with_environment({ + "metadataStoreUrl": "zk:zookeeper:2181", + "PULSAR_MEM": "-Xms256m -Xmx256m -XX:MaxDirectMemorySize=256m", + }) + .with_port(2181, 2181, "zookeeper") + .with_port(2888, 2888, "zookeeper2") + .with_port(3888, 3888, "zookeeper3"); + + // Pulsar cluster init container + local initContainer = + engine.container("pulsar-init") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "sleep 10 && bin/pulsar initialize-cluster-metadata --cluster cluster-a --zookeeper zookeeper:2181 --configuration-store zookeeper:2181 --web-service-url http://pulsar:8080 --broker-service-url pulsar://pulsar:6650", + ]) + .with_limits("1", "512M") + .with_reservations("0.05", "512M") + .with_environment({ + "PULSAR_MEM": "-Xms256m -Xmx256m -XX:MaxDirectMemorySize=256m", + }); + + + // Bookkeeper volume + local bookieVolume = engine.volume("bookie").with_size("20G"); + + // Bookkeeper container + local bookieContainer = + engine.container("bookie") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/bookkeeper.conf && exec bin/pulsar bookie" + // false ^ causes this to be a 'failure' exit. + ]) + .with_limits("1", "800M") + .with_reservations("0.1", "800M") + .with_user(0) + .with_volume_mount(bookieVolume, "/pulsar/data/bookkeeper") + .with_environment({ + "clusterName": "cluster-a", + "zkServers": "zookeeper:2181", + "bookieId": "bookie", + "metadataStoreUri": "metadata-store:zk:zookeeper:2181", + "advertisedAddress": "bookie", + "BOOKIE_MEM": "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=256m", + }) + .with_port(3181, 3181, "bookie"); + + // Pulsar broker, stateless (uses ZK and Bookkeeper for state) + local brokerContainer = engine.container("pulsar") .with_image(images.pulsar) - .with_command(["bin/pulsar", "standalone"]) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/broker.conf && exec bin/pulsar broker" + ]) + .with_limits("1", "800M") + .with_reservations("0.1", "800M") .with_environment({ - "PULSAR_MEM": "-Xms600M -Xmx600M" + "metadataStoreUrl": "zk:zookeeper:2181", + "zookeeperServers": "zookeeper:2181", + "clusterName": "cluster-a", + "managedLedgerDefaultEnsembleSize": "1", + "managedLedgerDefaultWriteQuorum": "1", + "managedLedgerDefaultAckQuorum": "1", + "advertisedAddress": "pulsar", + "advertisedListeners": "external:pulsar://pulsar:6650", + "PULSAR_MEM": "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=256m", }) - .with_limits("2.0", "1500M") - .with_reservations("1.0", "1500M") -// .with_volume_mount(confVolume, "/pulsar/conf") - .with_volume_mount(dataVolume, "/pulsar/data") - .with_port(6650, 6650, "bookie") - .with_port(8080, 8080, "http"); + .with_port(6650, 6650, "pulsar") + .with_port(8080, 8080, "admin"); + // Trustgraph Pulsar initialisation local adminContainer = - engine.container("init-pulsar") + engine.container("init-trustgraph") .with_image(images.trustgraph) .with_command([ "tg-init-pulsar", @@ -36,10 +117,32 @@ local url = import "values/url.jsonnet"; .with_limits("1", "128M") .with_reservations("0.1", "128M"); - local containerSet = engine.containers( - "pulsar", + // Container sets + local zkContainerSet = engine.containers( + "zookeeper", [ - container + zkContainer, + ] + ); + + local initContainerSet = engine.containers( + "init-pulsar", + [ + initContainer, + ] + ); + + local bookieContainerSet = engine.containers( + "bookie", + [ + bookieContainer, + ] + ); + + local brokerContainerSet = engine.containers( + "broker", + [ + brokerContainer, ] ); @@ -50,17 +153,35 @@ local url = import "values/url.jsonnet"; ] ); - local service = - engine.service(containerSet) - .with_port(6650, 6650, "bookie") - .with_port(8080, 8080, "http"); + // Zookeeper service + local zkService = + engine.service(zkContainerSet) + .with_port(2181, 2181, "zookeeper") + .with_port(2888, 2888, "zookeeper2") + .with_port(3888, 3888, "zookeeper3"); + + // Bookkeeper service + local bookieService = + engine.service(bookieContainerSet) + .with_port(3181, 3181, "bookie"); + + // Pulsar broker service + local brokerService = + engine.service(brokerContainerSet) + .with_port(6650, 6650, "pulsar") + .with_port(8080, 8080, "admin"); engine.resources([ -// confVolume, - dataVolume, - containerSet, + zkVolume, + bookieVolume, + zkContainerSet, + initContainerSet, + bookieContainerSet, + brokerContainerSet, adminContainerSet, - service, + zkService, + bookieService, + brokerService, ]) } diff --git a/templates/engine/docker-compose.jsonnet b/templates/engine/docker-compose.jsonnet index c37f1df0..0be3c3e3 100644 --- a/templates/engine/docker-compose.jsonnet +++ b/templates/engine/docker-compose.jsonnet @@ -22,6 +22,8 @@ with_image:: function(x) self + { image: x }, + with_user:: function(x) self + { user: x }, + with_command:: function(x) self + { command: x }, with_environment:: function(x) self + { @@ -75,6 +77,10 @@ { command: container.command } else {}) + + (if std.objectHas(container, "user") then + { user: container.user } + else {}) + + (if ! std.isEmpty(container.environment) then { environment: container.environment } else {}) + From e8f7600d74a079c95356f59197bf08e1d5e2b772 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Thu, 2 Jan 2025 19:52:47 +0000 Subject: [PATCH 11/53] Upgrade workbench (#240) --- templates/values/images.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 46c2bffc..a663bdd6 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -14,5 +14,5 @@ local version = import "version.jsonnet"; memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.0", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.2", } From 39ec7b63537073703b6e8695dbacfdfb337cd712 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Sat, 4 Jan 2025 13:27:04 -0800 Subject: [PATCH 12/53] Push Workbench to 0.2.4 --- templates/values/images.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index a663bdd6..9da4e89f 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -14,5 +14,5 @@ local version = import "version.jsonnet"; memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.2", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.4", } From 6aa212061d484d08267bebc0f65c07bf0b9c7ea3 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 4 Jan 2025 21:51:28 +0000 Subject: [PATCH 13/53] Fix/document embeddings (#247) * Update schema for doc embeddings * Rename embeddings-vectorize to graph-embeddings * Added document-embeddings processor (broken, needs fixing) * Added scripts * Fixed DE queue schema * Add missing DE process * Fix doc RAG processing, put graph-rag and doc-rag in appropriate component files. --- templates/components/document-rag.jsonnet | 30 +++++ templates/components/graph-rag.jsonnet | 30 +++++ templates/components/trustgraph.jsonnet | 30 ----- trustgraph-base/trustgraph/api/api.py | 29 +++++ .../clients/document_embeddings_client.py | 6 +- .../trustgraph/schema/documents.py | 17 +++ trustgraph-flow/scripts/document-embeddings | 6 + trustgraph-flow/scripts/embeddings-vectorize | 6 - trustgraph-flow/scripts/graph-embeddings | 6 + trustgraph-flow/setup.py | 19 +-- trustgraph-flow/trustgraph/document_rag.py | 78 +++++++------ .../document_embeddings/__init__.py | 3 + .../__main__.py | 2 +- .../document_embeddings/embeddings.py | 109 ++++++++++++++++++ .../embeddings/graph_embeddings/__init__.py | 3 + .../embeddings/graph_embeddings/__main__.py | 6 + .../embeddings.py} | 10 +- .../embeddings/vectorize/__init__.py | 3 - trustgraph-flow/trustgraph/gateway/service.py | 9 ++ .../storage/doc_embeddings/milvus/write.py | 24 ++-- .../storage/doc_embeddings/pinecone/write.py | 107 ++++++++--------- .../storage/doc_embeddings/qdrant/write.py | 77 +++++++------ 22 files changed, 421 insertions(+), 189 deletions(-) create mode 100755 trustgraph-flow/scripts/document-embeddings delete mode 100755 trustgraph-flow/scripts/embeddings-vectorize create mode 100755 trustgraph-flow/scripts/graph-embeddings create mode 100644 trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py rename trustgraph-flow/trustgraph/embeddings/{vectorize => document_embeddings}/__main__.py (57%) create mode 100755 trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py create mode 100644 trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py create mode 100755 trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py rename trustgraph-flow/trustgraph/embeddings/{vectorize/vectorize.py => graph_embeddings/embeddings.py} (92%) delete mode 100644 trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py diff --git a/templates/components/document-rag.jsonnet b/templates/components/document-rag.jsonnet index 0a68dd52..11dc9c13 100644 --- a/templates/components/document-rag.jsonnet +++ b/templates/components/document-rag.jsonnet @@ -39,5 +39,35 @@ local prompts = import "prompts/mixtral.jsonnet"; }, + "document-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("document-embeddings") + .with_image(images.trustgraph) + .with_command([ + "document-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "document-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/graph-rag.jsonnet b/templates/components/graph-rag.jsonnet index 860152c9..eb72754e 100644 --- a/templates/components/graph-rag.jsonnet +++ b/templates/components/graph-rag.jsonnet @@ -138,5 +138,35 @@ local url = import "values/url.jsonnet"; }, + "graph-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("graph-embeddings") + .with_image(images.trustgraph) + .with_command([ + "graph-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "graph-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 31ae420e..541beeed 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -119,36 +119,6 @@ local prompt = import "prompt-template.jsonnet"; }, - "vectorize" +: { - - create:: function(engine) - - local container = - engine.container("vectorize") - .with_image(images.trustgraph) - .with_command([ - "embeddings-vectorize", - "-p", - url.pulsar, - ]) - .with_limits("1.0", "512M") - .with_reservations("0.5", "512M"); - - local containerSet = engine.containers( - "vectorize", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - containerSet, - service, - ]) - - }, - "metering" +: { create:: function(engine) diff --git a/trustgraph-base/trustgraph/api/api.py b/trustgraph-base/trustgraph/api/api.py index de96499c..24207f32 100644 --- a/trustgraph-base/trustgraph/api/api.py +++ b/trustgraph-base/trustgraph/api/api.py @@ -131,6 +131,35 @@ class Api: except: raise ProtocolException(f"Response not formatted correctly") + def document_rag(self, question): + + # The input consists of a question + input = { + "query": question + } + + url = f"{self.url}document-rag" + + # Invoke the API, input is passed as JSON + resp = requests.post(url, json=input) + + # Should be a 200 status code + if resp.status_code != 200: + raise ProtocolException(f"Status code {resp.status_code}") + + try: + # Parse the response as JSON + object = resp.json() + except: + raise ProtocolException(f"Expected JSON response") + + self.check_error(resp) + + try: + return object["response"] + except: + raise ProtocolException(f"Response not formatted correctly") + def embeddings(self, text): # The input consists of a text block diff --git a/trustgraph-base/trustgraph/clients/document_embeddings_client.py b/trustgraph-base/trustgraph/clients/document_embeddings_client.py index d432991d..5b6d324e 100644 --- a/trustgraph-base/trustgraph/clients/document_embeddings_client.py +++ b/trustgraph-base/trustgraph/clients/document_embeddings_client.py @@ -38,8 +38,12 @@ class DocumentEmbeddingsClient(BaseClient): output_schema=DocumentEmbeddingsResponse, ) - def request(self, vectors, limit=10, timeout=300): + def request( + self, vectors, user="trustgraph", collection="default", + limit=10, timeout=300 + ): return self.call( + user=user, collection=collection, vectors=vectors, limit=limit, timeout=timeout ).documents diff --git a/trustgraph-base/trustgraph/schema/documents.py b/trustgraph-base/trustgraph/schema/documents.py index 38add83d..fd0049ee 100644 --- a/trustgraph-base/trustgraph/schema/documents.py +++ b/trustgraph-base/trustgraph/schema/documents.py @@ -35,11 +35,28 @@ chunk_ingest_queue = topic('chunk-load') ############################################################################ +# Document embeddings are embeddings associated with a chunk + +class ChunkEmbeddings(Record): + chunk = Bytes() + vectors = Array(Array(Double())) + +# This is a 'batching' mechanism for the above data +class DocumentEmbeddings(Record): + metadata = Metadata() + chunks = Array(ChunkEmbeddings()) + +document_embeddings_store_queue = topic('document-embeddings-store') + +############################################################################ + # Doc embeddings query class DocumentEmbeddingsRequest(Record): vectors = Array(Array(Double())) limit = Integer() + user = String() + collection = String() class DocumentEmbeddingsResponse(Record): error = Error() diff --git a/trustgraph-flow/scripts/document-embeddings b/trustgraph-flow/scripts/document-embeddings new file mode 100755 index 00000000..26bb85b0 --- /dev/null +++ b/trustgraph-flow/scripts/document-embeddings @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.document_embeddings import run + +run() + 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/scripts/graph-embeddings b/trustgraph-flow/scripts/graph-embeddings new file mode 100755 index 00000000..29b1fbf4 --- /dev/null +++ b/trustgraph-flow/scripts/graph-embeddings @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.graph_embeddings 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/document_rag.py b/trustgraph-flow/trustgraph/document_rag.py index f3c8b158..86298783 100644 --- a/trustgraph-flow/trustgraph/document_rag.py +++ b/trustgraph-flow/trustgraph/document_rag.py @@ -16,6 +16,44 @@ from . schema import document_embeddings_response_queue LABEL="http://www.w3.org/2000/01/rdf-schema#label" DEFINITION="http://www.w3.org/2004/02/skos/core#definition" +class Query: + + def __init__(self, rag, user, collection, verbose): + self.rag = rag + self.user = user + self.collection = collection + self.verbose = verbose + + def get_vector(self, query): + + if self.verbose: + print("Compute embeddings...", flush=True) + + qembeds = self.rag.embeddings.request(query) + + if self.verbose: + print("Done.", flush=True) + + return qembeds + + def get_docs(self, query): + + vectors = self.get_vector(query) + + if self.verbose: + print("Get entities...", flush=True) + + docs = self.rag.de_client.request( + vectors, limit=self.rag.doc_limit + ) + + if self.verbose: + print("Docs:", flush=True) + for doc in docs: + print(doc, flush=True) + + return docs + class DocumentRag: def __init__( @@ -55,7 +93,7 @@ class DocumentRag: print("Initialising...", flush=True) # FIXME: Configurable - self.entity_limit = 20 + self.doc_limit = 20 self.de_client = DocumentEmbeddingsClient( pulsar_host=pulsar_host, @@ -81,42 +119,16 @@ class DocumentRag: if self.verbose: print("Initialised", flush=True) - def get_vector(self, query): - - if self.verbose: - print("Compute embeddings...", flush=True) - - qembeds = self.embeddings.request(query) - - if self.verbose: - print("Done.", flush=True) - - return qembeds - - def get_docs(self, query): - - vectors = self.get_vector(query) - - if self.verbose: - print("Get entities...", flush=True) - - docs = self.de_client.request( - vectors, self.entity_limit - ) - - if self.verbose: - print("Docs:", flush=True) - for doc in docs: - print(doc, flush=True) - - return docs - - def query(self, query): + def query(self, query, user="trustgraph", collection="default"): if self.verbose: print("Construct prompt...", flush=True) - docs = self.get_docs(query) + q = Query( + rag=self, user=user, collection=collection, verbose=self.verbose + ) + + docs = q.get_docs(query) if self.verbose: print("Invoke LLM...", flush=True) 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/document_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py new file mode 100755 index 00000000..745ab4db --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py @@ -0,0 +1,109 @@ + +""" +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 Chunk, ChunkEmbeddings, DocumentEmbeddings +from ... schema import chunk_ingest_queue +from ... schema import document_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 = chunk_ingest_queue +default_output_queue = document_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": Chunk, + "output_schema": DocumentEmbeddings, + } + ) + + self.embeddings = EmbeddingsClient( + pulsar_host=self.pulsar_host, + input_queue=emb_request_queue, + output_queue=emb_response_queue, + subscriber=module + "-emb", + ) + + def handle(self, msg): + + v = msg.value() + print(f"Indexing {v.metadata.id}...", flush=True) + + try: + + vectors = self.embeddings.request(v.chunk) + + embeds = [ + ChunkEmbeddings( + chunk=v.chunk, + vectors=vectors, + ) + ] + + r = DocumentEmbeddings( + metadata=v.metadata, + chunks=embeds, + ) + + 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/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/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py similarity index 92% rename from trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py rename to trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index 5630a7b5..e4d1646e 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph-flow/trustgraph/embeddings/graph_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. +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 @@ -51,11 +52,6 @@ class Processor(ConsumerProducer): 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() 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 * - diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index a260b631..d6306ac6 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -31,6 +31,7 @@ from . subscriber import Subscriber from . text_completion import TextCompletionRequestor from . prompt import PromptRequestor from . graph_rag import GraphRagRequestor +from . document_rag import DocumentRagRequestor from . triples_query import TriplesQueryRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor from . embeddings import EmbeddingsRequestor @@ -91,6 +92,10 @@ class Api: pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, ), + "document-rag": DocumentRagRequestor( + pulsar_host=self.pulsar_host, timeout=self.timeout, + auth = self.auth, + ), "triples-query": TriplesQueryRequestor( pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, @@ -140,6 +145,10 @@ class Api: endpoint_path = "/api/v1/graph-rag", auth=self.auth, requestor = self.services["graph-rag"], ), + ServiceEndpoint( + endpoint_path = "/api/v1/document-rag", auth=self.auth, + requestor = self.services["document-rag"], + ), ServiceEndpoint( endpoint_path = "/api/v1/triples-query", auth=self.auth, requestor = self.services["triples-query"], diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py index 00f9d5b5..bfa6c123 100755 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py @@ -3,15 +3,16 @@ Accepts entity/vector pairs and writes them to a Milvus store. """ -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue -from .... log_level import LogLevel from .... direct.milvus_doc_embeddings import DocVectors + +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue +from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_store_uri = 'http://localhost:19530' @@ -27,7 +28,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "store_uri": store_uri, } ) @@ -38,11 +39,16 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if v.chunk != "" and v.chunk is not None: - for vec in v.vectors: - self.vecstore.insert(vec, chunk) + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue + + for vec in emb.vectors: + + if chunk != "" and v.chunk is not None: + for vec in v.vectors: + self.vecstore.insert(vec, chunk) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index 24cfcb78..c59ecd7b 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -11,14 +11,14 @@ import time import uuid import os -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_api_key = os.getenv("PINECONE_API_KEY", "not-specified") default_cloud = "aws" @@ -54,7 +54,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "url": self.url, } ) @@ -65,71 +65,74 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d-" + v.metadata.user + "-" + str(dim) - ) + for vec in v.vectors: - if index_name != self.last_index_name: + dim = len(vec) + collection = ( + "d-" + v.metadata.user + "-" + str(dim) + ) - if not self.pinecone.has_index(index_name): + if index_name != self.last_index_name: - try: + if not self.pinecone.has_index(index_name): - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + for i in range(0, 1000): - time.sleep(1) + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + time.sleep(1) - except Exception as e: - print("Pinecone index creation failed") - raise e + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) - print(f"Index {index_name} created", flush=True) + except Exception as e: + print("Pinecone index creation failed") + raise e - self.last_index_name = index_name + print(f"Index {index_name} created", flush=True) - index = self.pinecone.Index(index_name) + self.last_index_name = index_name - records = [ - { - "id": id, - "values": vec, - "metadata": { "doc": chunk }, - } - ] + index = self.pinecone.Index(index_name) - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + records = [ + { + "id": id, + "values": vec, + "metadata": { "doc": chunk }, + } + ] + + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 813c4f29..f852e03b 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -8,14 +8,14 @@ from qdrant_client.models import PointStruct from qdrant_client.models import Distance, VectorParams import uuid -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_store_uri = 'http://localhost:6333' @@ -31,7 +31,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "store_uri": store_uri, } ) @@ -44,47 +44,48 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "": return - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + - str(dim) - ) + dim = len(vec) + collection = ( + "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + + str(dim) + ) - if collection != self.last_collection: + if collection != self.last_collection: - if not self.client.collection_exists(collection): + if not self.client.collection_exists(collection): - try: - self.client.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.DOT - ), + try: + self.client.create_collection( + collection_name=collection, + vectors_config=VectorParams( + size=dim, distance=Distance.COSINE + ), + ) + except Exception as e: + print("Qdrant collection creation failed") + raise e + + self.last_collection = collection + + self.client.upsert( + collection_name=collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "doc": chunk, + } ) - except Exception as e: - print("Qdrant collection creation failed") - raise e - - self.last_collection = collection - - self.client.upsert( - collection_name=collection, - points=[ - PointStruct( - id=str(uuid.uuid4()), - vector=vec, - payload={ - "doc": chunk, - } - ) - ] - ) + ] + ) @staticmethod def add_args(parser): From ae00831e395511dd028c0a1ad9ce4ce4fd540fe0 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 4 Jan 2025 22:08:24 +0000 Subject: [PATCH 14/53] Added missing file (#249) --- .../trustgraph/gateway/document_rag.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 trustgraph-flow/trustgraph/gateway/document_rag.py diff --git a/trustgraph-flow/trustgraph/gateway/document_rag.py b/trustgraph-flow/trustgraph/gateway/document_rag.py new file mode 100644 index 00000000..6665a7dc --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_rag.py @@ -0,0 +1,30 @@ + +from .. schema import DocumentRagQuery, DocumentRagResponse +from .. schema import document_rag_request_queue +from .. schema import document_rag_response_queue + +from . endpoint import ServiceEndpoint +from . requestor import ServiceRequestor + +class DocumentRagRequestor(ServiceRequestor): + def __init__(self, pulsar_host, timeout, auth): + + super(DocumentRagRequestor, self).__init__( + pulsar_host=pulsar_host, + request_queue=document_rag_request_queue, + response_queue=document_rag_response_queue, + request_schema=DocumentRagQuery, + response_schema=DocumentRagResponse, + timeout=timeout, + ) + + def to_request(self, body): + return DocumentRagQuery( + query=body["query"], + user=body.get("user", "trustgraph"), + collection=body.get("collection", "default"), + ) + + def from_response(self, message): + return { "response": message.response }, True + From dbdf65b84dadefbc0511812b15ba266478bf92c2 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sun, 5 Jan 2025 19:24:43 +0000 Subject: [PATCH 15/53] Added a metrics endpoint /api/v1/metrics/X which proxies to Prometheus (#251) --- trustgraph-flow/trustgraph/gateway/metrics.py | 73 +++++++++++++++++++ trustgraph-flow/trustgraph/gateway/service.py | 20 +++++ 2 files changed, 93 insertions(+) create mode 100644 trustgraph-flow/trustgraph/gateway/metrics.py diff --git a/trustgraph-flow/trustgraph/gateway/metrics.py b/trustgraph-flow/trustgraph/gateway/metrics.py new file mode 100644 index 00000000..33c1fe3a --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/metrics.py @@ -0,0 +1,73 @@ + +# +# This provides a Prometheus endpoint on the api-gateway. It proxies +# HTTP GET requests to Prometheus. +# + +import aiohttp +from aiohttp import web +import asyncio +from pulsar.schema import JsonSchema +import uuid +import logging + +logger = logging.getLogger("endpoint") +logger.setLevel(logging.INFO) + +class MetricsEndpoint: + + def __init__(self, prometheus_url, endpoint_path, auth): + + self.prometheus_url = prometheus_url + self.path = endpoint_path + self.auth = auth + self.operation = "service" + + async def start(self): + pass + + def add_routes(self, app): + + app.add_routes([ + web.get(self.path + "/{path:.*}", self.handle), + ]) + + async def handle(self, request): + + print(request.path, "...") + + try: + ht = request.headers["Authorization"] + tokens = ht.split(" ", 2) + if tokens[0] != "Bearer": + return web.HTTPUnauthorized() + token = tokens[1] + except: + token = "" + + if not self.auth.permitted(token, self.operation): + return web.HTTPUnauthorized() + + try: + + path = request.match_info["path"] + + async with aiohttp.ClientSession() as session: + + url = ( + self.prometheus_url + "/api/v1/" + path + "?" + + request.query_string + ) + + async with session.get(url) as resp: + return web.Response( + status=resp.status, + text=await resp.text() + ) + + except Exception as e: + + logging.error(f"Exception: {e}") + + raise web.HTTPInternalServerError() + diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index d6306ac6..b329660f 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -46,6 +46,7 @@ from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint from . mux import MuxEndpoint from . document_load import DocumentLoadSender from . text_load import TextLoadSender +from . metrics import MetricsEndpoint from . endpoint import ServiceEndpoint from . auth import Authenticator @@ -54,6 +55,7 @@ logger = logging.getLogger("api") logger.setLevel(logging.INFO) default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650") +default_prometheus_url = os.getenv("PROMETHEUS_URL", "http://prometheus:9090") default_timeout = 600 default_port = 8088 default_api_token = os.getenv("GATEWAY_SECRET", "") @@ -71,6 +73,13 @@ class Api: self.timeout = int(config.get("timeout", default_timeout)) self.pulsar_host = config.get("pulsar_host", default_pulsar_host) + self.prometheus_url = config.get( + "prometheus_url", default_prometheus_url, + ) + + if not self.prometheus_url.endswith("/"): + self.prometheus_url += "/" + api_token = config.get("api_token", default_api_token) # Token not set, or token equal empty string means no auth @@ -207,6 +216,11 @@ class Api: auth = self.auth, services = self.services, ), + MetricsEndpoint( + endpoint_path = "/api/v1/metrics", + prometheus_url = self.prometheus_url, + auth = self.auth, + ), ] for ep in self.endpoints: @@ -235,6 +249,12 @@ def run(): help=f'Pulsar host (default: {default_pulsar_host})', ) + parser.add_argument( + '-m', '--prometheus-url', + default=default_prometheus_url, + help=f'Prometheus URL (default: {default_prometheus_url})', + ) + parser.add_argument( '--port', type=int, From bb0d48c1b215d4e21e9737794b377023eb8e749f Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Mon, 6 Jan 2025 10:03:22 -0800 Subject: [PATCH 16/53] Hot fix for pulsar user issue --- templates/components/pulsar.jsonnet | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 58170500..745f2cd8 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -28,6 +28,7 @@ local url = import "values/url.jsonnet"; ]) .with_limits("0.1", "400M") .with_reservations("0.05", "400M") + .with_user("0:1000") .with_volume_mount(zkVolume, "/pulsar/data/zookeeper") .with_environment({ "metadataStoreUrl": "zk:zookeeper:2181", @@ -68,7 +69,7 @@ local url = import "values/url.jsonnet"; ]) .with_limits("1", "800M") .with_reservations("0.1", "800M") - .with_user(0) + .with_user("0:1000") .with_volume_mount(bookieVolume, "/pulsar/data/bookkeeper") .with_environment({ "clusterName": "cluster-a", From 970e82458908051b2c44786ad80a89e062a2466e Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Mon, 6 Jan 2025 12:06:51 -0800 Subject: [PATCH 17/53] Fixed Zookeeper startup delay --- templates/components/pulsar.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 745f2cd8..96062989 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -26,7 +26,7 @@ local url = import "values/url.jsonnet"; "-c", "bin/apply-config-from-env.py conf/zookeeper.conf && bin/generate-zookeeper-config.sh conf/zookeeper.conf && exec bin/pulsar zookeeper" ]) - .with_limits("0.1", "400M") + .with_limits("1", "400M") .with_reservations("0.05", "400M") .with_user("0:1000") .with_volume_mount(zkVolume, "/pulsar/data/zookeeper") From 6a27de22db7064ba0332accf151918584d7a476e Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Mon, 6 Jan 2025 12:33:57 -0800 Subject: [PATCH 18/53] Updated docker push action --- .github/workflows/release.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 7e18afb9..b119ec94 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -62,7 +62,7 @@ jobs: - name: Build and push Docker image id: push - uses: docker/build-push-action@3b5e8027fcad23fda98b2e3ac259d8d67585f671 + uses: docker/build-push-action@v6 with: context: . file: ./Containerfile From 5946c47d3dc5c203e20dd6fe15d81807c051a557 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Mon, 6 Jan 2025 12:43:01 -0800 Subject: [PATCH 19/53] Updated docker login and checkout versions --- .github/workflows/release.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b119ec94..b2157edb 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -25,10 +25,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Log in to Docker Hub - uses: docker/login-action@f4ef78c080cd8ba55a85445d5b36e214a81df20a + uses: docker/login-action@v3 with: username: ${{ vars.DOCKER_USERNAME }} password: ${{ secrets.DOCKER_SECRET }} From dc2b599fda861b9b58731fe4a16376e01ef1ddbe Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 6 Jan 2025 21:45:42 +0000 Subject: [PATCH 20/53] Fix/release broken (#257) * Break release into 3 jobs * Replace Github action with podman command --- .github/workflows/release.yaml | 75 +++++++++++++++++++++++++++------- Makefile | 2 +- 2 files changed, 61 insertions(+), 16 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b2157edb..bec60a0d 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,17 +5,17 @@ on: workflow_dispatch: push: tags: - - v0.19.* + - v* permissions: contents: read jobs: - deploy: + python-packages: - name: Build everything - runs-on: ubuntu-latest + name: Release Python packages + runs-on: ubuntu-24.04 permissions: contents: write id-token: write @@ -27,27 +27,44 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - name: Log in to Docker Hub - uses: docker/login-action@v3 - with: - username: ${{ vars.DOCKER_USERNAME }} - password: ${{ secrets.DOCKER_SECRET }} - - - name: Install build dependencies - run: pip3 install jsonnet - - name: Get version id: version run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT - - run: echo ${{ steps.version.outputs.VERSION }} - - name: Build packages run: make packages VERSION=${{ steps.version.outputs.VERSION }} - name: Publish release distributions to PyPI uses: pypa/gh-action-pypi-publish@release/v1 + deploy-container-image: + + name: Release container image + runs-on: ubuntu-24.04 + permissions: + contents: write + id-token: write + environment: + name: release + + steps: + + - name: Checkout + uses: actions/checkout@v4 + + - name: Docker Hub token + run: echo ${{ secrets.DOCKER_SECRET }} > docker-token.txt + + - name: Authenticate with Docker hub + run: make docker-hub-login + + - name: Get version + id: version + run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT + + - name: Put version into package manifests + run: make update-package-versions VERSION=${{ steps.version.outputs.VERSION }} + - name: Extract metadata for container id: meta uses: docker/metadata-action@v4 @@ -60,6 +77,12 @@ jobs: type=semver,pattern={{major}}.{{minor}} type=sha + - name: Build containers + run: make container VERSION=${{ steps.version.outputs.VERSION }} + + - name: Push containers + run: make push VERSION=${{ steps.version.outputs.VERSION }} + - name: Build and push Docker image id: push uses: docker/build-push-action@v6 @@ -70,6 +93,28 @@ jobs: tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} + release-bundle: + + name: Upload release bundle + runs-on: ubuntu-24.04 + permissions: + contents: write + id-token: write + environment: + name: release + + steps: + + - name: Checkout + uses: actions/checkout@v4 + + - name: Install build dependencies + run: pip3 install jsonnet + + - name: Get version + id: version + run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT + - name: Create deploy bundle run: templates/generate-all deploy.zip ${{ steps.version.outputs.VERSION }} diff --git a/Makefile b/Makefile index 67094a90..f408a31f 100644 --- a/Makefile +++ b/Makefile @@ -104,5 +104,5 @@ update-dcs: set-version docker-hub-login: cat docker-token.txt | \ - docker login -u trustgraph --password-stdin registry-1.docker.io + ${DOCKER} login -u trustgraph --password-stdin registry-1.docker.io From ceff3f0e3458dab71a040c318928c223060a442c Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 6 Jan 2025 21:52:25 +0000 Subject: [PATCH 21/53] Remove 2nd push (#258) --- .github/workflows/release.yaml | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index bec60a0d..c182578a 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -65,34 +65,12 @@ jobs: - name: Put version into package manifests run: make update-package-versions VERSION=${{ steps.version.outputs.VERSION }} - - name: Extract metadata for container - id: meta - uses: docker/metadata-action@v4 - with: - images: trustgraph/trustgraph-flow - tags: | - type=ref,event=branch - type=ref,event=pr - type=semver,pattern={{version}} - type=semver,pattern={{major}}.{{minor}} - type=sha - - name: Build containers run: make container VERSION=${{ steps.version.outputs.VERSION }} - name: Push containers run: make push VERSION=${{ steps.version.outputs.VERSION }} - - name: Build and push Docker image - id: push - uses: docker/build-push-action@v6 - with: - context: . - file: ./Containerfile - push: true - tags: ${{ steps.meta.outputs.tags }} - labels: ${{ steps.meta.outputs.labels }} - release-bundle: name: Upload release bundle From ffacf84c7942ba2ff48fe4c5912cf27b009a7346 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Tue, 7 Jan 2025 08:54:35 -0800 Subject: [PATCH 22/53] Fix for Claude key envvar --- templates/components/claude.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/components/claude.jsonnet b/templates/components/claude.jsonnet index 00e4ec79..d1c24392 100644 --- a/templates/components/claude.jsonnet +++ b/templates/components/claude.jsonnet @@ -13,7 +13,7 @@ local prompts = import "prompts/mixtral.jsonnet"; create:: function(engine) local envSecrets = engine.envSecrets("claude-credentials") - .with_env_var("CLAUDE_KEY_TOKEN", "claude-key"); + .with_env_var("CLAUDE_KEY", "claude-key"); local container = engine.container("text-completion") From 08d24bf5aa17b67aa6b5353a153d71698a04a83a Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Tue, 7 Jan 2025 18:50:56 +0000 Subject: [PATCH 23/53] Fix k8s with_user invocation (#261) --- templates/engine/k8s.jsonnet | 2 ++ 1 file changed, 2 insertions(+) diff --git a/templates/engine/k8s.jsonnet b/templates/engine/k8s.jsonnet index 2fec0d1f..dfd8d11f 100644 --- a/templates/engine/k8s.jsonnet +++ b/templates/engine/k8s.jsonnet @@ -14,6 +14,8 @@ with_image:: function(x) self + { image: x }, + with_user:: function(x) self + { user: x }, + with_command:: function(x) self + { command: x }, with_environment:: function(x) self + { From 8f769b1e73e72b1c1bdb2ebcdff6e18e92c0c6fd Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Tue, 7 Jan 2025 12:24:57 -0800 Subject: [PATCH 24/53] Fixed creds for GoogleAIStudio in K8s --- templates/components/googleaistudio.jsonnet | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet index c2a40f2c..85e91520 100644 --- a/templates/components/googleaistudio.jsonnet +++ b/templates/components/googleaistudio.jsonnet @@ -13,7 +13,7 @@ local prompts = import "prompts/mixtral.jsonnet"; create:: function(engine) - local envSecrets = engine.envSecrets("googleaistudio-key") + local envSecrets = engine.envSecrets("googleaistudio-credentials") .with_env_var("GOOGLE_AI_STUDIO_KEY", "googleaistudio-key"); local container = From cec9e292228193e5747eceae0c38b5ac367f7c1c Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 11 Jan 2025 11:40:42 +0000 Subject: [PATCH 25/53] Dual model templates (#263) * Dual-mode templates * Fixed generate-all so that works * Fix ability to specify Claude model --- templates/components.jsonnet | 73 +++++++++++------- templates/components/azure-openai-rag.jsonnet | 61 +++++++++++++++ templates/components/azure-openai.jsonnet | 35 ++------- templates/components/azure-rag.jsonnet | 60 +++++++++++++++ templates/components/azure.jsonnet | 36 ++------- templates/components/bedrock-rag.jsonnet | 66 +++++++++++++++++ templates/components/bedrock.jsonnet | 37 ++-------- templates/components/claude-rag.jsonnet | 63 ++++++++++++++++ templates/components/claude.jsonnet | 38 ++-------- templates/components/cohere-rag.jsonnet | 56 ++++++++++++++ templates/components/cohere.jsonnet | 34 +-------- .../components/googleaistudio-rag.jsonnet | 65 ++++++++++++++++ templates/components/googleaistudio.jsonnet | 32 -------- templates/components/llamafile-rag.jsonnet | 57 ++++++++++++++ templates/components/llamafile.jsonnet | 33 ++------- templates/components/ollama-rag.jsonnet | 57 ++++++++++++++ templates/components/ollama.jsonnet | 33 ++------- templates/components/openai-rag.jsonnet | 63 ++++++++++++++++ templates/components/openai.jsonnet | 37 ++-------- templates/components/prompt-overrides.jsonnet | 4 - templates/components/trustgraph.jsonnet | 3 +- templates/components/vertexai-rag.jsonnet | 74 +++++++++++++++++++ templates/components/vertexai.jsonnet | 54 -------------- templates/generate-all | 14 +++- 24 files changed, 725 insertions(+), 360 deletions(-) create mode 100644 templates/components/azure-openai-rag.jsonnet create mode 100644 templates/components/azure-rag.jsonnet create mode 100644 templates/components/bedrock-rag.jsonnet create mode 100644 templates/components/claude-rag.jsonnet create mode 100644 templates/components/cohere-rag.jsonnet create mode 100644 templates/components/googleaistudio-rag.jsonnet create mode 100644 templates/components/llamafile-rag.jsonnet create mode 100644 templates/components/ollama-rag.jsonnet create mode 100644 templates/components/openai-rag.jsonnet create mode 100644 templates/components/vertexai-rag.jsonnet diff --git a/templates/components.jsonnet b/templates/components.jsonnet index b14665d6..653c48bb 100644 --- a/templates/components.jsonnet +++ b/templates/components.jsonnet @@ -1,47 +1,70 @@ { + + // Essentials + "trustgraph-base": import "components/trustgraph.jsonnet", + "pulsar": import "components/pulsar.jsonnet", + + // LLMs "azure": import "components/azure.jsonnet", "azure-openai": import "components/azure-openai.jsonnet", "bedrock": import "components/bedrock.jsonnet", "claude": import "components/claude.jsonnet", "cohere": import "components/cohere.jsonnet", - "document-rag": import "components/document-rag.jsonnet", - "embeddings-hf": import "components/embeddings-hf.jsonnet", - "embeddings-ollama": import "components/embeddings-ollama.jsonnet", "googleaistudio": import "components/googleaistudio.jsonnet", - "grafana": import "components/grafana.jsonnet", + "ollama": import "components/ollama.jsonnet", + "openai": import "components/openai.jsonnet", + "vertexai": import "components/vertexai.jsonnet", + + // LLMs for RAG + "azure-rag": import "components/azure-rag.jsonnet", + "azure-openai-rag": import "components/azure-openai-rag.jsonnet", + "bedrock-rag": import "components/bedrock-rag.jsonnet", + "claude-rag": import "components/claude-rag.jsonnet", + "cohere-rag": import "components/cohere-rag.jsonnet", + "googleaistudio-rag": import "components/googleaistudio-rag.jsonnet", + "ollama-rag": import "components/ollama-rag.jsonnet", + "openai-rag": import "components/openai-rag.jsonnet", + "vertexai-rag": import "components/vertexai-rag.jsonnet", + + // Embeddings + "embeddings-ollama": import "components/embeddings-ollama.jsonnet", + "embeddings-hf": import "components/embeddings-hf.jsonnet", + + // Processing pipelines "graph-rag": import "components/graph-rag.jsonnet", + "document-rag": import "components/document-rag.jsonnet", + + // Vector stores + "vector-store-milvus": import "components/milvus.jsonnet", + "vector-store-qdrant": import "components/qdrant.jsonnet", + "vector-store-pinecone": import "components/pinecone.jsonnet", + + // Triples stores "triple-store-cassandra": import "components/cassandra.jsonnet", "triple-store-neo4j": import "components/neo4j.jsonnet", "triple-store-falkordb": import "components/falkordb.jsonnet", "triple-store-memgraph": import "components/memgraph.jsonnet", + + // Observability support + "grafana": import "components/grafana.jsonnet", + + // Pulsar manager is a UI for Pulsar. Uses a LOT of memory + "pulsar-manager": import "components/pulsar-manager.jsonnet", + "llamafile": import "components/llamafile.jsonnet", - "ollama": import "components/ollama.jsonnet", - "openai": import "components/openai.jsonnet", "override-recursive-chunker": import "components/chunker-recursive.jsonnet", + // The prompt manager "prompt-template": import "components/prompt-template.jsonnet", "prompt-overrides": import "components/prompt-overrides.jsonnet", - "pulsar": import "components/pulsar.jsonnet", - "pulsar-manager": import "components/pulsar-manager.jsonnet", - "trustgraph-base": import "components/trustgraph.jsonnet", - "vector-store-milvus": import "components/milvus.jsonnet", - "vector-store-qdrant": import "components/qdrant.jsonnet", - "vector-store-pinecone": import "components/pinecone.jsonnet", - "vertexai": import "components/vertexai.jsonnet", - "workbench-ui": import "components/workbench-ui.jsonnet", - "null": {}, - + // ReAct agent "agent-manager-react": import "components/agent-manager-react.jsonnet", - // FIXME: Dupes - "cassandra": import "components/cassandra.jsonnet", - "neo4j": import "components/neo4j.jsonnet", - "memgraph": import "components/memgraph.jsonnet", - "qdrant": import "components/qdrant.jsonnet", - "pinecone": import "components/pinecone.jsonnet", - "milvus": import "components/milvus.jsonnet", - "falkordb": import "components/falkordb.jsonnet", - "trustgraph": import "components/trustgraph.jsonnet", + // Optional UI + "workbench-ui": import "components/workbench-ui.jsonnet", + + // Does nothing. But, can be a hack to overwrite parameters + "null": {}, } diff --git a/templates/components/azure-openai-rag.jsonnet b/templates/components/azure-openai-rag.jsonnet new file mode 100644 index 00000000..fab83266 --- /dev/null +++ b/templates/components/azure-openai-rag.jsonnet @@ -0,0 +1,61 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["ollama-rag-" + key]:: value, + }, + + "azure-openai-rag-model":: "GPT-3.5-Turbo", + "azure-openai-rag-max-output-tokens":: 4192, + "azure-openai-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("azure-openai-credentials") + .with_env_var("AZURE_TOKEN", "azure-token"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-azure", + "-p", + url.pulsar, + "-x", + std.toString($["azure-openai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["azure-openai-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/azure-openai.jsonnet b/templates/components/azure-openai.jsonnet index 8afcaf11..f6bda306 100644 --- a/templates/components/azure-openai.jsonnet +++ b/templates/components/azure-openai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["azure-openai-" + key]:: value, + }, + "azure-openai-model":: "GPT-3.5-Turbo", "azure-openai-max-output-tokens":: 4192, "azure-openai-temperature":: 0.0, @@ -34,48 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-azure", - "-p", - url.pulsar, - "-x", - std.toString($["azure-openai-max-output-tokens"]), - "-t", - "%0.3f" % $["azure-openai-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/azure-rag.jsonnet b/templates/components/azure-rag.jsonnet new file mode 100644 index 00000000..da2f1692 --- /dev/null +++ b/templates/components/azure-rag.jsonnet @@ -0,0 +1,60 @@ +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["azure-rag-" + key]:: value, + }, + + "azure-rag-max-output-tokens":: 4096, + "azure-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("azure-credentials") + .with_env_var("AZURE_TOKEN", "azure-token") + .with_env_var("AZURE_ENDPOINT", "azure-endpoint"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-azure", + "-p", + url.pulsar, + "-x", + std.toString($["azure-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["azure-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + } + +} + prompts + diff --git a/templates/components/azure.jsonnet b/templates/components/azure.jsonnet index cf10dc66..bd364d8d 100644 --- a/templates/components/azure.jsonnet +++ b/templates/components/azure.jsonnet @@ -1,10 +1,14 @@ -local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["azure-" + key]:: value, + }, + "azure-max-output-tokens":: 4096, "azure-temperature":: 0.0, @@ -32,48 +36,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-azure", - "-p", - url.pulsar, - "-x", - std.toString($["azure-max-output-tokens"]), - "-t", - "%0.3f" % $["azure-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) } diff --git a/templates/components/bedrock-rag.jsonnet b/templates/components/bedrock-rag.jsonnet new file mode 100644 index 00000000..68aa9b0a --- /dev/null +++ b/templates/components/bedrock-rag.jsonnet @@ -0,0 +1,66 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; +local chunker = import "chunker-recursive.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["bedrock-rag-" + key]:: value, + }, + + "bedrock-rag-max-output-tokens":: 4096, + "bedrock-rag-temperature":: 0.0, + "bedrock-rag-model":: "mistral.mixtral-8x7b-instruct-v0:1", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("bedrock-credentials") + .with_env_var("AWS_ID_KEY", "aws-id-key") + .with_env_var("AWS_SECRET", "aws-secret") + .with_env_var("AWS_REGION", "aws-region"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-bedrock", + "-p", + url.pulsar, + "-x", + std.toString($["bedrock-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["bedrock-rag-temperature"], + "-m", + $["bedrock-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + chunker + diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet index 6ccaa1c5..922e0409 100644 --- a/templates/components/bedrock.jsonnet +++ b/templates/components/bedrock.jsonnet @@ -6,6 +6,11 @@ local chunker = import "chunker-recursive.jsonnet"; { + with:: function(key, value) + self + { + ["bedrock-" + key]:: value, + }, + "bedrock-max-output-tokens":: 4096, "bedrock-temperature":: 0.0, "bedrock-model":: "mistral.mixtral-8x7b-instruct-v0:1", @@ -37,50 +42,18 @@ local chunker = import "chunker-recursive.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-bedrock", - "-p", - url.pulsar, - "-x", - std.toString($["bedrock-max-output-tokens"]), - "-t", - "%0.3f" % $["bedrock-temperature"], - "-m", - $["bedrock-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/claude-rag.jsonnet b/templates/components/claude-rag.jsonnet new file mode 100644 index 00000000..71a214fc --- /dev/null +++ b/templates/components/claude-rag.jsonnet @@ -0,0 +1,63 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["claude-rag-" + key]:: value, + }, + + "claude-rag-model":: "claude-3-sonnet-20240229", + "claude-rag-max-output-tokens":: 4096, + "claude-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("claude-credentials") + .with_env_var("CLAUDE_KEY", "claude-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-claude", + "-p", + url.pulsar, + "-x", + std.toString($["claude-rag-max-output-tokens"]), + "-m", + $["claude-rag-model"], + "-t", + "%0.3f" % $["claude-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/claude.jsonnet b/templates/components/claude.jsonnet index d1c24392..fd1650d2 100644 --- a/templates/components/claude.jsonnet +++ b/templates/components/claude.jsonnet @@ -5,6 +5,12 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["claude-" + key]:: value, + }, + + "claude-model":: "claude-3-sonnet-20240229", "claude-max-output-tokens":: 4096, "claude-temperature":: 0.0, @@ -24,6 +30,8 @@ local prompts = import "prompts/mixtral.jsonnet"; url.pulsar, "-x", std.toString($["claude-max-output-tokens"]), + "-m", + $["claude-model"], "-t", "%0.3f" % $["claude-temperature"], ]) @@ -31,48 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-claude", - "-p", - url.pulsar, - "-x", - std.toString($["claude-max-output-tokens"]), - "-t", - "%0.3f" % $["claude-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/cohere-rag.jsonnet b/templates/components/cohere-rag.jsonnet new file mode 100644 index 00000000..62829445 --- /dev/null +++ b/templates/components/cohere-rag.jsonnet @@ -0,0 +1,56 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["cohere-rag-" + key]:: value, + }, + + "cohere-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("cohere-credentials") + .with_env_var("COHERE_KEY", "cohere-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-cohere", + "-p", + url.pulsar, + "-t", + "%0.3f" % $["cohere-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/cohere.jsonnet b/templates/components/cohere.jsonnet index 5bc9b39c..2f559ad0 100644 --- a/templates/components/cohere.jsonnet +++ b/templates/components/cohere.jsonnet @@ -5,9 +5,10 @@ local prompts = import "prompts/mixtral.jsonnet"; { - // Override chunking - "chunk-size":: 150, - "chunk-overlap":: 10, + with:: function(key, value) + self + { + ["cohere-" + key]:: value, + }, "cohere-temperature":: 0.0, @@ -31,45 +32,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-cohere", - "-p", - url.pulsar, - "-t", - "%0.3f" % $["cohere-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/googleaistudio-rag.jsonnet b/templates/components/googleaistudio-rag.jsonnet new file mode 100644 index 00000000..3bfa9618 --- /dev/null +++ b/templates/components/googleaistudio-rag.jsonnet @@ -0,0 +1,65 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["googleaistudio-rag-" + key]:: value, + }, + + "googleaistudio-rag-max-output-tokens":: 4096, + "googleaistudio-rag-temperature":: 0.0, + "googleaistudio-rag-model":: "gemini-1.5-flash-002", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("googleaistudio-key") + .with_env_var("GOOGLE_AI_STUDIO_KEY", "googleaistudio-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-googleaistudio", + "-p", + url.pulsar, + "-x", + std.toString( + $["googleaistudio-rag-max-output-tokens"] + ), + "-t", + "%0.3f" % $["googleaistudio-rag-temperature"], + "-m", + $["googleaistudio-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet index 85e91520..00b2bd0e 100644 --- a/templates/components/googleaistudio.jsonnet +++ b/templates/components/googleaistudio.jsonnet @@ -34,50 +34,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-googleaistudio", - "-p", - url.pulsar, - "-x", - std.toString($["googleaistudio-max-output-tokens"]), - "-t", - "%0.3f" % $["googleaistudio-temperature"], - "-m", - $["googleaistudio-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/llamafile-rag.jsonnet b/templates/components/llamafile-rag.jsonnet new file mode 100644 index 00000000..34af69d0 --- /dev/null +++ b/templates/components/llamafile-rag.jsonnet @@ -0,0 +1,57 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/slm.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["llamafile-rag-" + key]:: value, + }, + + "llamafile-rag-model":: "LLaMA_CPP", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("llamafile-credentials") + .with_env_var("LLAMAFILE_URL", "llamafile-url"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-llamafile", + "-p", + url.pulsar, + "-m", + $["llamafile-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/llamafile.jsonnet b/templates/components/llamafile.jsonnet index bc1a011c..bc16dfaa 100644 --- a/templates/components/llamafile.jsonnet +++ b/templates/components/llamafile.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/slm.jsonnet"; { + with:: function(key, value) + self + { + ["llamafile-" + key]:: value, + }, + "llamafile-model":: "LLaMA_CPP", "text-completion" +: { @@ -28,46 +33,18 @@ local prompts = import "prompts/slm.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-llamafile", - "-p", - url.pulsar, - "-m", - $["llamafile-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/ollama-rag.jsonnet b/templates/components/ollama-rag.jsonnet new file mode 100644 index 00000000..e3ca97aa --- /dev/null +++ b/templates/components/ollama-rag.jsonnet @@ -0,0 +1,57 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["ollama-rag-" + key]:: value, + }, + + "ollama-rag-model":: "gemma2:9b", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("ollama-credentials") + .with_env_var("OLLAMA_HOST", "ollama-host"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-ollama", + "-p", + url.pulsar, + "-m", + $["ollama-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/ollama.jsonnet b/templates/components/ollama.jsonnet index 8da00848..28bf36b9 100644 --- a/templates/components/ollama.jsonnet +++ b/templates/components/ollama.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["ollama-" + key]:: value, + }, + "ollama-model":: "gemma2:9b", "text-completion" +: { @@ -28,46 +33,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-ollama", - "-p", - url.pulsar, - "-m", - $["ollama-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/openai-rag.jsonnet b/templates/components/openai-rag.jsonnet new file mode 100644 index 00000000..d9ecd964 --- /dev/null +++ b/templates/components/openai-rag.jsonnet @@ -0,0 +1,63 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["openai-rag-" + key]:: value, + }, + + "openai-rag-max-output-tokens":: 4096, + "openai-rag-temperature":: 0.0, + "openai-rag-model":: "GPT-3.5-Turbo", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("openai-credentials") + .with_env_var("OPENAI_TOKEN", "openai-token"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-openai", + "-p", + url.pulsar, + "-x", + std.toString($["openai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["openai-rag-temperature"], + "-m", + $["openai-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/openai.jsonnet b/templates/components/openai.jsonnet index 27725cb6..50194bfc 100644 --- a/templates/components/openai.jsonnet +++ b/templates/components/openai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["openai-" + key]:: value, + }, + "openai-max-output-tokens":: 4096, "openai-temperature":: 0.0, "openai-model":: "GPT-3.5-Turbo", @@ -34,50 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-openai", - "-p", - url.pulsar, - "-x", - std.toString($["openai-max-output-tokens"]), - "-t", - "%0.3f" % $["openai-temperature"], - "-m", - $["openai-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/prompt-overrides.jsonnet b/templates/components/prompt-overrides.jsonnet index 648e5b66..852ec09d 100644 --- a/templates/components/prompt-overrides.jsonnet +++ b/templates/components/prompt-overrides.jsonnet @@ -1,7 +1,3 @@ -local base = import "base/base.jsonnet"; -local images = import "values/images.jsonnet"; -local url = import "values/url.jsonnet"; -local prompts = import "prompts/mixtral.jsonnet"; local default_prompts = import "prompts/default-prompts.jsonnet"; { diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 541beeed..f63740aa 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -1,7 +1,6 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; -local prompt = import "prompt-template.jsonnet"; { @@ -181,5 +180,5 @@ local prompt = import "prompt-template.jsonnet"; }, -} + prompt +} diff --git a/templates/components/vertexai-rag.jsonnet b/templates/components/vertexai-rag.jsonnet new file mode 100644 index 00000000..9c7eaa00 --- /dev/null +++ b/templates/components/vertexai-rag.jsonnet @@ -0,0 +1,74 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["vertexai-rag-" + key]:: value, + }, + + "vertexai-rag-model":: "gemini-1.0-pro-001", + "vertexai-rag-private-key":: "/vertexai/private.json", + "vertexai-rag-region":: "us-central1", + "vertexai-rag-max-output-tokens":: 4096, + "vertexai-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local cfgVol = engine.secretVolume( + "vertexai-creds", + "./vertexai", + { + "private.json": importstr "vertexai/private.json", + } + ); + + local container = + engine.container("text-completion-rag") + .with_image(images.trustgraph) + .with_command([ + "text-completion-vertexai", + "-p", + url.pulsar, + "-k", + $["vertexai-rag-private-key"], + "-r", + $["vertexai-rag-region"], + "-x", + std.toString($["vertexai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["vertexai-rag-temperature"], + "-m", + $["vertexai-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_limits("0.5", "256M") + .with_reservations("0.1", "256M") + .with_volume_mount(cfgVol, "/vertexai"); + + local containerSet = engine.containers( + "text-completion-rag", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + cfgVol, + containerSet, + service, + ]) + + } + +} + prompts + diff --git a/templates/components/vertexai.jsonnet b/templates/components/vertexai.jsonnet index ef193156..fc3192a4 100644 --- a/templates/components/vertexai.jsonnet +++ b/templates/components/vertexai.jsonnet @@ -61,59 +61,5 @@ local prompts = import "prompts/mixtral.jsonnet"; }, - "text-completion-rag" +: { - - create:: function(engine) - - local cfgVol = engine.secretVolume( - "vertexai-creds", - "./vertexai", - { - "private.json": importstr "vertexai/private.json", - } - ); - - local container = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-vertexai", - "-p", - url.pulsar, - "-k", - $["vertexai-private-key"], - "-r", - $["vertexai-region"], - "-x", - std.toString($["vertexai-max-output-tokens"]), - "-t", - "%0.3f" % $["vertexai-temperature"], - "-m", - $["vertexai-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_limits("0.5", "256M") - .with_reservations("0.1", "256M") - .with_volume_mount(cfgVol, "/vertexai"); - - local containerSet = engine.containers( - "text-completion-rag", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - cfgVol, - containerSet, - service, - ]) - - } - } + prompts diff --git a/templates/generate-all b/templates/generate-all index 1ae71402..22c9a5b0 100755 --- a/templates/generate-all +++ b/templates/generate-all @@ -87,8 +87,18 @@ def full_config_object( ): return config_object([ - graph_store, "pulsar", vector_store, embeddings, - "graph-rag", "grafana", "trustgraph", llm, "workbench-ui", + "triple-store-" + graph_store, + "pulsar", + "vector-store-" + vector_store, + embeddings, + "graph-rag", + "grafana", + "trustgraph-base", + llm, + llm + "-rag", + "workbench-ui", + "prompt-template", + "agent-manager-react", ]) def generate_config( From c603caa3cc81f6f3a61aa1733689d8d64d7122bd Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 11 Jan 2025 12:44:49 +0000 Subject: [PATCH 26/53] Fix neo4j: container doesn't start currently (#264) * Adjust Neo4j memory settings to avoid startup fail * Upgrade Neo4j to community 5.26.0 * Add Neo4j indexes --- templates/stores/neo4j.jsonnet | 6 ++- templates/values/images.jsonnet | 2 +- .../storage/triples/memgraph/write.py | 8 ++++ .../trustgraph/storage/triples/neo4j/write.py | 44 +++++++++++++++++++ 4 files changed, 57 insertions(+), 3 deletions(-) diff --git a/templates/stores/neo4j.jsonnet b/templates/stores/neo4j.jsonnet index 55cccc5f..3a8bb783 100644 --- a/templates/stores/neo4j.jsonnet +++ b/templates/stores/neo4j.jsonnet @@ -14,12 +14,14 @@ local images = import "values/images.jsonnet"; .with_image(images.neo4j) .with_environment({ NEO4J_AUTH: "neo4j/password", + NEO4J_server_memory_pagecache_size: "512m", + NEO4J_server_memory_heap_max__size: "512m", // NEO4J_server_bolt_listen__address: "0.0.0.0:7687", // NEO4J_server_default__listen__address: "0.0.0.0", // NEO4J_server_http_listen__address: "0.0.0.0:7474", }) - .with_limits("1.0", "768M") - .with_reservations("0.5", "768M") + .with_limits("1.0", "1536M") + .with_reservations("0.5", "1536M") .with_port(7474, 7474, "api") .with_port(7687, 7687, "api2") .with_volume_mount(vol, "/data"); diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 9da4e89f..b0416eb3 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -1,7 +1,7 @@ local version = import "version.jsonnet"; { cassandra: "docker.io/cassandra:4.1.6", - neo4j: "docker.io/neo4j:5.22.0-community-bullseye", + neo4j: "docker.io/neo4j:5.26.0-community-bullseye", pulsar: "docker.io/apachepulsar/pulsar:3.3.1", pulsar_manager: "docker.io/apachepulsar/pulsar-manager:v0.4.0", etcd: "quay.io/coreos/etcd:v3.5.15", diff --git a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py index f106170a..7295e691 100755 --- a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py @@ -55,6 +55,14 @@ class Processor(Consumer): def create_indexes(self, session): + # Race condition, index creation failure is ignored. Right thing + # to do if the index already exists. Wrong thing to do if it's + # because the store is not up yet + + # In real-world cases, Memgraph will start up quicker than Pulsar + # and this process will restart several times until Pulsar arrives, + # so should be safe + print("Create indexes...", flush=True) try: diff --git a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py index 1aa25aa8..18b40129 100755 --- a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py @@ -50,6 +50,50 @@ class Processor(Consumer): self.io = GraphDatabase.driver(graph_host, auth=(username, password)) + with self.io.session(database=self.db) as session: + self.create_indexes(session) + + def create_indexes(self, session): + + # Race condition, index creation failure is ignored. Right thing + # to do if the index already exists. Wrong thing to do if it's + # because the store is not up yet + + # In real-world cases, Neo4j will start up quicker than Pulsar + # and this process will restart several times until Pulsar arrives, + # so should be safe + + print("Create indexes...", flush=True) + + try: + session.run( + "CREATE INDEX Node_uri FOR (n:Node) ON (n.uri)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + try: + session.run( + "CREATE INDEX Literal_value FOR (n:Literal) ON (n.value)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + try: + session.run( + "CREATE INDEX Rel_uri FOR ()-[r:Rel]-() ON (r.uri)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + print("Index creation done", flush=True) + def create_node(self, uri): print("Create node", uri) From cd9a2084321ee202345e87488a30eff035e6e272 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 11 Jan 2025 18:10:04 +0000 Subject: [PATCH 27/53] Processors use shared queues, means there can be more than process on a queue to share load (#265) --- trustgraph-base/trustgraph/base/consumer.py | 2 ++ trustgraph-base/trustgraph/base/consumer_producer.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index eeaf83a1..f346f1bc 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -1,5 +1,6 @@ from pulsar.schema import JsonSchema +import pulsar from prometheus_client import Histogram, Info, Counter, Enum import time @@ -51,6 +52,7 @@ class Consumer(BaseProcessor): self.consumer = self.client.subscribe( input_queue, subscriber, + consumer_type=pulsar.ConsumerType.Shared, schema=JsonSchema(input_schema), ) diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py index 31441cda..6d386894 100644 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ b/trustgraph-base/trustgraph/base/consumer_producer.py @@ -1,5 +1,6 @@ from pulsar.schema import JsonSchema +import pulsar from prometheus_client import Histogram, Info, Counter, Enum import time @@ -71,6 +72,7 @@ class ConsumerProducer(BaseProcessor): self.consumer = self.client.subscribe( input_queue, subscriber, + consumer_type=pulsar.ConsumerType.Shared, schema=JsonSchema(input_schema), ) From 1280af3eff3b4de9ef7dfd03caec647789f47d3b Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 13 Jan 2025 14:42:33 +0000 Subject: [PATCH 28/53] Fix/kg core save (#267) * Add a 'localhost' Pulsar endpoint for docker env * - Fix broken socket endpoint streamers - Add unused listener endpoints for publisher/subscriber - Fix graph embedding serialisation * Fix GE load * Remove Gossip settling delay, this is single-node Cassandra. --- templates/components/pulsar.jsonnet | 2 +- templates/stores/cassandra.jsonnet | 2 +- trustgraph-cli/scripts/tg-load-kg-core | 9 +++++++-- trustgraph-cli/scripts/tg-save-kg-core | 9 +++++++-- .../trustgraph/gateway/graph_embeddings_load.py | 11 ++++++++--- .../trustgraph/gateway/graph_embeddings_stream.py | 13 +++++++++++++ trustgraph-flow/trustgraph/gateway/publisher.py | 5 +++-- trustgraph-flow/trustgraph/gateway/requestor.py | 5 ++++- trustgraph-flow/trustgraph/gateway/serialize.py | 9 +++++++-- trustgraph-flow/trustgraph/gateway/subscriber.py | 4 +++- .../trustgraph/gateway/triples_stream.py | 13 +++++++++++++ 11 files changed, 67 insertions(+), 15 deletions(-) diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 96062989..fdec2c33 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -100,7 +100,7 @@ local url = import "values/url.jsonnet"; "managedLedgerDefaultWriteQuorum": "1", "managedLedgerDefaultAckQuorum": "1", "advertisedAddress": "pulsar", - "advertisedListeners": "external:pulsar://pulsar:6650", + "advertisedListeners": "external:pulsar://pulsar:6650,localhost:pulsar://localhost:6650", "PULSAR_MEM": "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=256m", }) .with_port(6650, 6650, "pulsar") diff --git a/templates/stores/cassandra.jsonnet b/templates/stores/cassandra.jsonnet index 0c90421e..2a9d6d7a 100644 --- a/templates/stores/cassandra.jsonnet +++ b/templates/stores/cassandra.jsonnet @@ -13,7 +13,7 @@ local images = import "values/images.jsonnet"; engine.container("cassandra") .with_image(images.cassandra) .with_environment({ - JVM_OPTS: "-Xms300M -Xmx300M", + JVM_OPTS: "-Xms300M -Xmx300M -Dcassandra.skip_wait_for_gossip_to_settle=0", }) .with_limits("1.0", "1000M") .with_reservations("0.5", "1000M") diff --git a/trustgraph-cli/scripts/tg-load-kg-core b/trustgraph-cli/scripts/tg-load-kg-core index 4e76e525..b79ec992 100755 --- a/trustgraph-cli/scripts/tg-load-kg-core +++ b/trustgraph-cli/scripts/tg-load-kg-core @@ -51,8 +51,13 @@ async def load_ge(running, queue, url): "user": msg["m"]["u"], "collection": msg["m"]["c"], }, - "vectors": msg["v"], - "entity": msg["e"], + "entities": [ + { + "entity": ent["e"], + "vectors": ent["v"], + } + for ent in msg["e"] + ], } try: diff --git a/trustgraph-cli/scripts/tg-save-kg-core b/trustgraph-cli/scripts/tg-save-kg-core index e52cd7dc..298f2e84 100755 --- a/trustgraph-cli/scripts/tg-save-kg-core +++ b/trustgraph-cli/scripts/tg-save-kg-core @@ -57,8 +57,13 @@ async def fetch_ge(running, queue, user, collection, url): "u": data["metadata"]["user"], "c": data["metadata"]["collection"], }, - "v": data["vectors"], - "e": data["entity"], + "e": [ + { + "e": ent["entity"], + "v": ent["vectors"], + } + for ent in data["entities"] + ] } ]) if msg.type == aiohttp.WSMsgType.ERROR: diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index 18a2e6fe..86f2016b 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py @@ -5,7 +5,7 @@ import uuid from aiohttp import WSMsgType from .. schema import Metadata -from .. schema import GraphEmbeddings +from .. schema import GraphEmbeddings, EntityEmbeddings from .. schema import graph_embeddings_store_queue from . publisher import Publisher @@ -50,8 +50,13 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): user=data["metadata"]["user"], collection=data["metadata"]["collection"], ), - entity=to_value(data["entity"]), - vectors=data["vectors"], + entities=[ + EntityEmbeddings( + entity=to_value(ent["entity"]), + vectors=ent["vectors"], + ) + for ent in data["entities"] + ] ) self.publisher.send(None, elt) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py index f0b4dd86..b6b5403f 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py @@ -29,6 +29,16 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): schema=JsonSchema(GraphEmbeddings) ) + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(GraphEmbeddingsStreamEndpoint, self).listener(ws, running) + + await worker + async def start(self): self.subscriber.start() @@ -44,6 +54,9 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_graph_embeddings(resp)) + except TimeoutError: + continue + except queue.Empty: continue diff --git a/trustgraph-flow/trustgraph/gateway/publisher.py b/trustgraph-flow/trustgraph/gateway/publisher.py index e3298a47..f79cfa1e 100644 --- a/trustgraph-flow/trustgraph/gateway/publisher.py +++ b/trustgraph-flow/trustgraph/gateway/publisher.py @@ -7,12 +7,13 @@ import threading class Publisher: def __init__(self, pulsar_host, topic, schema=None, max_size=10, - chunking_enabled=True): + chunking_enabled=True, listener=None): self.pulsar_host = pulsar_host self.topic = topic self.schema = schema self.q = queue.Queue(maxsize=max_size) self.chunking_enabled = chunking_enabled + self.listener_name = listener def start(self): self.task = threading.Thread(target=self.run) @@ -25,7 +26,7 @@ class Publisher: try: client = pulsar.Client( - self.pulsar_host, + self.pulsar_host, listener_name=self.listener_name ) producer = client.create_producer( diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index 5bfedf1e..567c9de7 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -60,7 +60,10 @@ class ServiceRequestor: while True: try: - resp = await asyncio.to_thread(q.get, timeout=self.timeout) + resp = await asyncio.to_thread( + q.get, + timeout=self.timeout + ) except Exception as e: raise RuntimeError("Timeout") diff --git a/trustgraph-flow/trustgraph/gateway/serialize.py b/trustgraph-flow/trustgraph/gateway/serialize.py index 35932382..5f9930ad 100644 --- a/trustgraph-flow/trustgraph/gateway/serialize.py +++ b/trustgraph-flow/trustgraph/gateway/serialize.py @@ -51,7 +51,12 @@ def serialize_graph_embeddings(message): "user": message.metadata.user, "collection": message.metadata.collection, }, - "vectors": message.vectors, - "entity": serialize_value(message.entity), + "entities": [ + { + "vectors": entity.vectors, + "entity": serialize_value(entity.entity), + } + for entity in message.entities + ], } diff --git a/trustgraph-flow/trustgraph/gateway/subscriber.py b/trustgraph-flow/trustgraph/gateway/subscriber.py index cccfc5b4..2eca5a36 100644 --- a/trustgraph-flow/trustgraph/gateway/subscriber.py +++ b/trustgraph-flow/trustgraph/gateway/subscriber.py @@ -7,7 +7,7 @@ import time class Subscriber: def __init__(self, pulsar_host, topic, subscription, consumer_name, - schema=None, max_size=100): + schema=None, max_size=100, listener=None): self.pulsar_host = pulsar_host self.topic = topic self.subscription = subscription @@ -17,6 +17,7 @@ class Subscriber: self.full = {} self.max_size = max_size self.lock = threading.Lock() + self.listener_name = listener def start(self): self.task = threading.Thread(target=self.run) @@ -30,6 +31,7 @@ class Subscriber: client = pulsar.Client( self.pulsar_host, + listener_name=self.listener_name, ) consumer = client.subscribe( diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/triples_stream.py index 92ada132..8048d3cc 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/triples_stream.py @@ -27,6 +27,16 @@ class TriplesStreamEndpoint(SocketEndpoint): schema=JsonSchema(Triples) ) + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(TriplesStreamEndpoint, self).listener(ws, running) + + await worker + async def start(self): self.subscriber.start() @@ -42,6 +52,9 @@ class TriplesStreamEndpoint(SocketEndpoint): resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_triples(resp)) + except TimeoutError: + continue + except queue.Empty: continue From acdd3efe51ea6499d1b6ae0b6dfd3432e847aad5 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 15 Jan 2025 23:59:50 +0000 Subject: [PATCH 29/53] Fix argument usage not overriding defaults in Google AI Studio & VertexAI templates (#268) --- templates/components/googleaistudio.jsonnet | 5 +++++ templates/components/vertexai.jsonnet | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet index 00b2bd0e..ac122243 100644 --- a/templates/components/googleaistudio.jsonnet +++ b/templates/components/googleaistudio.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["googleaistudio-" + key]:: value, + }, + "googleaistudio-max-output-tokens":: 4096, "googleaistudio-temperature":: 0.0, "googleaistudio-model":: "gemini-1.5-flash-002", diff --git a/templates/components/vertexai.jsonnet b/templates/components/vertexai.jsonnet index fc3192a4..49243701 100644 --- a/templates/components/vertexai.jsonnet +++ b/templates/components/vertexai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["vertexai-" + key]:: value, + }, + "vertexai-model":: "gemini-1.0-pro-001", "vertexai-private-key":: "/vertexai/private.json", "vertexai-region":: "us-central1", From bed7423c26396dfe7647e45f2b010447305bb6a8 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Thu, 16 Jan 2025 00:00:59 +0000 Subject: [PATCH 30/53] - Added load/save API for document embeddings (#269) - tg-load-doc-embeds and tg-save-doc-embeds command line utils --- trustgraph-cli/scripts/tg-load-doc-embeds | 224 ++++++++++++++++++ trustgraph-cli/scripts/tg-save-doc-embeds | 198 ++++++++++++++++ trustgraph-cli/setup.py | 2 + .../gateway/document_embeddings_load.py | 64 +++++ .../gateway/document_embeddings_stream.py | 72 ++++++ .../trustgraph/gateway/serialize.py | 17 ++ trustgraph-flow/trustgraph/gateway/service.py | 10 + 7 files changed, 587 insertions(+) create mode 100755 trustgraph-cli/scripts/tg-load-doc-embeds create mode 100755 trustgraph-cli/scripts/tg-save-doc-embeds create mode 100644 trustgraph-flow/trustgraph/gateway/document_embeddings_load.py create mode 100644 trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py diff --git a/trustgraph-cli/scripts/tg-load-doc-embeds b/trustgraph-cli/scripts/tg-load-doc-embeds new file mode 100755 index 00000000..d445ec5a --- /dev/null +++ b/trustgraph-cli/scripts/tg-load-doc-embeds @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 + +"""This utility takes a knowledge core and loads it into a running TrustGraph +through the API. The knowledge core should be in msgpack format, which is the +default format produce by tg-save-kg-core. +""" + +import aiohttp +import asyncio +import msgpack +import json +import sys +import argparse +import os +import signal + +class Running: + def __init__(self): self.running = True + def get(self): return self.running + def stop(self): self.running = False + +de_counts = 0 + +async def load_de(running, queue, url): + + global de_counts + + async with aiohttp.ClientSession() as session: + + async with session.ws_connect(f"{url}load/document-embeddings") as ws: + + while running.get(): + + try: + msg = await asyncio.wait_for(queue.get(), 1) + + # End of load + if msg is None: + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + msg = { + "metadata": { + "id": msg["m"]["i"], + "metadata": msg["m"]["m"], + "user": msg["m"]["u"], + "collection": msg["m"]["c"], + }, + "chunks": [ + { + "chunk": chunk["c"], + "vectors": chunk["v"], + } + for chunk in msg["c"] + ], + } + + try: + await ws.send_json(msg) + except Exception as e: + print(e) + + de_counts += 1 + +async def stats(running): + + global de_counts + + while running.get(): + + await asyncio.sleep(2) + + print( + f"Graph embeddings: {de_counts:10d}" + ) + +async def loader(running, de_queue, path, format, user, collection): + + if format == "json": + + raise RuntimeError("Not implemented") + + else: + + with open(path, "rb") as f: + + unpacker = msgpack.Unpacker(f, raw=False) + + while running.get(): + + try: + unpacked = unpacker.unpack() + except: + break + + if user: + unpacked["metadata"]["user"] = user + + if collection: + unpacked["metadata"]["collection"] = collection + + if unpacked[0] == "de": + qtype = de_queue + + while running.get(): + + try: + await asyncio.wait_for(qtype.put(unpacked[1]), 0.5) + + # Successful put message, move on + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + if not running.get(): break + + # Put 'None' on end of queue to finish + while running.get(): + + try: + await asyncio.wait_for(de_queue.put(None), 1) + + # Successful put message, move on + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + +async def run(running, **args): + + # Maxsize on queues reduces back-pressure so tg-load-kg-core doesn't + # grow to eat all memory + de_q = asyncio.Queue(maxsize=10) + + load_task = asyncio.create_task( + loader( + running=running, + de_queue=de_q, + path=args["input_file"], format=args["format"], + user=args["user"], collection=args["collection"], + ) + + ) + + de_task = asyncio.create_task( + load_de( + running=running, + queue=de_q, url=args["url"] + "api/v1/" + ) + ) + + stats_task = asyncio.create_task(stats(running)) + + await de_task + + running.stop() + + await load_task + await stats_task + +async def main(running): + + parser = argparse.ArgumentParser( + prog='tg-load-kg-core', + description=__doc__, + ) + + default_url = os.getenv("TRUSTGRAPH_API", "http://localhost:8088/") + default_user = "trustgraph" + collection = "default" + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'TrustGraph API URL (default: {default_url})', + ) + + parser.add_argument( + '-i', '--input-file', + # Make it mandatory, difficult to over-write an existing file + required=True, + help=f'Output file' + ) + + parser.add_argument( + '--format', + default="msgpack", + choices=["msgpack", "json"], + help=f'Output format (default: msgpack)', + ) + + parser.add_argument( + '--user', + help=f'User ID to load as (default: from input)' + ) + + parser.add_argument( + '--collection', + help=f'Collection ID to load as (default: from input)' + ) + + args = parser.parse_args() + + await run(running, **vars(args)) + +running = Running() + +def interrupt(sig, frame): + running.stop() + print('Interrupt') + +signal.signal(signal.SIGINT, interrupt) + +asyncio.run(main(running)) + diff --git a/trustgraph-cli/scripts/tg-save-doc-embeds b/trustgraph-cli/scripts/tg-save-doc-embeds new file mode 100755 index 00000000..95f8b748 --- /dev/null +++ b/trustgraph-cli/scripts/tg-save-doc-embeds @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 + +""" +This utility connects to a running TrustGraph through the API and creates +a knowledge core from the data streaming through the processing queues. +For completeness of data, tg-save-kg-core should be initiated before data +loading takes place. The default output format, msgpack should be used. +JSON output format is also available - msgpack produces a more compact +representation, which is also more performant to load. +""" + +import aiohttp +import asyncio +import msgpack +import json +import sys +import argparse +import os +import signal + +class Running: + def __init__(self): self.running = True + def get(self): return self.running + def stop(self): self.running = False + +async def fetch_de(running, queue, user, collection, url): + + async with aiohttp.ClientSession() as session: + + de_url = f"{url}stream/document-embeddings" + + async with session.ws_connect(de_url) as ws: + + while running.get(): + + try: + msg = await asyncio.wait_for(ws.receive(), 1) + except: + continue + + if msg.type == aiohttp.WSMsgType.TEXT: + + data = msg.json() + + if user: + if data["metadata"]["user"] != user: + continue + + if collection: + if data["metadata"]["collection"] != collection: + continue + + await queue.put([ + "de", + { + "m": { + "i": data["metadata"]["id"], + "m": data["metadata"]["metadata"], + "u": data["metadata"]["user"], + "c": data["metadata"]["collection"], + }, + "c": [ + { + "c": chunk["chunk"], + "v": chunk["vectors"], + } + for chunk in data["chunks"] + ] + } + ]) + if msg.type == aiohttp.WSMsgType.ERROR: + print("Error") + break + +de_counts = 0 + +async def stats(running): + + global t_counts + global de_counts + + while running.get(): + + await asyncio.sleep(2) + + print( + f"Document embeddings: {de_counts:10d}" + ) + +async def output(running, queue, path, format): + + global t_counts + global de_counts + + with open(path, "wb") as f: + + while running.get(): + + try: + msg = await asyncio.wait_for(queue.get(), 0.5) + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + if format == "msgpack": + f.write(msgpack.packb(msg, use_bin_type=True)) + else: + f.write(json.dumps(msg).encode("utf-8")) + + if msg[0] == "de": + de_counts += 1 + + print("Output file closed") + +async def run(running, **args): + + q = asyncio.Queue() + + de_task = asyncio.create_task( + fetch_de( + running=running, + queue=q, user=args["user"], collection=args["collection"], + url=args["url"] + "api/v1/" + ) + ) + + output_task = asyncio.create_task( + output( + running=running, queue=q, + path=args["output_file"], format=args["format"], + ) + + ) + + stats_task = asyncio.create_task(stats(running)) + + await output_task + await de_task + await stats_task + + print("Exiting") + +async def main(running): + + parser = argparse.ArgumentParser( + prog='tg-save-kg-core', + description=__doc__, + ) + + default_url = os.getenv("TRUSTGRAPH_API", "http://localhost:8088/") + default_user = "trustgraph" + collection = "default" + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'TrustGraph API URL (default: {default_url})', + ) + + parser.add_argument( + '-o', '--output-file', + # Make it mandatory, difficult to over-write an existing file + required=True, + help=f'Output file' + ) + + parser.add_argument( + '--format', + default="msgpack", + choices=["msgpack", "json"], + help=f'Output format (default: msgpack)', + ) + + parser.add_argument( + '--user', + help=f'User ID to filter on (default: no filter)' + ) + + parser.add_argument( + '--collection', + help=f'Collection ID to filter on (default: no filter)' + ) + + args = parser.parse_args() + + await run(running, **vars(args)) + +running = Running() + +def interrupt(sig, frame): + running.stop() + print('Interrupt') + +signal.signal(signal.SIGINT, interrupt) + +asyncio.run(main(running)) + diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 67c70158..7edffd4b 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -54,10 +54,12 @@ setuptools.setup( "scripts/tg-invoke-llm", "scripts/tg-invoke-prompt", "scripts/tg-load-kg-core", + "scripts/tg-load-doc-embeds", "scripts/tg-load-pdf", "scripts/tg-load-text", "scripts/tg-load-turtle", "scripts/tg-processor-state", "scripts/tg-save-kg-core", + "scripts/tg-save-doc-embeds", ] ) diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py new file mode 100644 index 00000000..1a7f635d --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py @@ -0,0 +1,64 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid +from aiohttp import WSMsgType + +from .. schema import Metadata +from .. schema import DocumentEmbeddings, ChunkEmbeddings +from .. schema import document_embeddings_store_queue + +from . publisher import Publisher +from . socket import SocketEndpoint +from . serialize import to_subgraph + +class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): + + def __init__( + self, pulsar_host, auth, path="/api/v1/load/document-embeddings", + ): + + super(DocumentEmbeddingsLoadEndpoint, self).__init__( + endpoint_path=path, auth=auth, + ) + + self.pulsar_host=pulsar_host + + self.publisher = Publisher( + self.pulsar_host, document_embeddings_store_queue, + schema=JsonSchema(DocumentEmbeddings) + ) + + async def start(self): + + self.publisher.start() + + async def listener(self, ws, running): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.ERROR: + break + else: + + data = msg.json() + + elt = DocumentEmbeddings( + metadata=Metadata( + id=data["metadata"]["id"], + metadata=to_subgraph(data["metadata"]["metadata"]), + user=data["metadata"]["user"], + collection=data["metadata"]["collection"], + ), + chunks=[ + ChunkEmbeddings( + chunk=de["chunk"].encode("utf-8"), + vectors=de["vectors"], + ) + for de in data["chunks"] + ], + ) + + self.publisher.send(None, elt) + + running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py new file mode 100644 index 00000000..99cfb0a9 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py @@ -0,0 +1,72 @@ + +import asyncio +import queue +from pulsar.schema import JsonSchema +import uuid + +from .. schema import DocumentEmbeddings +from .. schema import document_embeddings_store_queue + +from . subscriber import Subscriber +from . socket import SocketEndpoint +from . serialize import serialize_document_embeddings + +class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): + + def __init__( + self, pulsar_host, auth, path="/api/v1/stream/document-embeddings" + ): + + super(DocumentEmbeddingsStreamEndpoint, self).__init__( + endpoint_path=path, auth=auth, + ) + + self.pulsar_host=pulsar_host + + self.subscriber = Subscriber( + self.pulsar_host, document_embeddings_store_queue, + "api-gateway", "api-gateway", + schema=JsonSchema(DocumentEmbeddings) + ) + + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(DocumentEmbeddingsStreamEndpoint, self).listener( + ws, running + ) + + await worker + + async def start(self): + + self.subscriber.start() + + async def async_thread(self, ws, running): + + id = str(uuid.uuid4()) + + q = self.subscriber.subscribe_all(id) + + while running.get(): + try: + resp = await asyncio.to_thread(q.get, timeout=0.5) + await ws.send_json(serialize_document_embeddings(resp)) + + except TimeoutError: + continue + + except queue.Empty: + continue + + except Exception as e: + print(f"Exception: {str(e)}", flush=True) + break + + self.subscriber.unsubscribe_all(id) + + running.stop() + diff --git a/trustgraph-flow/trustgraph/gateway/serialize.py b/trustgraph-flow/trustgraph/gateway/serialize.py index 5f9930ad..40b6efc5 100644 --- a/trustgraph-flow/trustgraph/gateway/serialize.py +++ b/trustgraph-flow/trustgraph/gateway/serialize.py @@ -60,3 +60,20 @@ def serialize_graph_embeddings(message): ], } +def serialize_document_embeddings(message): + return { + "metadata": { + "id": message.metadata.id, + "metadata": serialize_subgraph(message.metadata.metadata), + "user": message.metadata.user, + "collection": message.metadata.collection, + }, + "chunks": [ + { + "vectors": chunk.vectors, + "chunk": chunk.chunk.decode("utf-8"), + } + for chunk in message.chunks + ], + } + diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index b329660f..644731e2 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -41,8 +41,10 @@ from . dbpedia import DbpediaRequestor from . internet_search import InternetSearchRequestor from . triples_stream import TriplesStreamEndpoint from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint +from . document_embeddings_stream import DocumentEmbeddingsStreamEndpoint from . triples_load import TriplesLoadEndpoint from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint +from . document_embeddings_load import DocumentEmbeddingsLoadEndpoint from . mux import MuxEndpoint from . document_load import DocumentLoadSender from . text_load import TextLoadSender @@ -203,6 +205,10 @@ class Api: pulsar_host=self.pulsar_host, auth = self.auth, ), + DocumentEmbeddingsStreamEndpoint( + pulsar_host=self.pulsar_host, + auth = self.auth, + ), TriplesLoadEndpoint( pulsar_host=self.pulsar_host, auth = self.auth, @@ -211,6 +217,10 @@ class Api: pulsar_host=self.pulsar_host, auth = self.auth, ), + DocumentEmbeddingsLoadEndpoint( + pulsar_host=self.pulsar_host, + auth = self.auth, + ), MuxEndpoint( pulsar_host=self.pulsar_host, auth = self.auth, From 9bf916f3c695e5ec1fbddad5cc4b27f4ae5acc99 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 16 Jan 2025 00:02:18 +0000 Subject: [PATCH 31/53] Update Python files for 0.20 --- trustgraph-bedrock/setup.py | 2 +- trustgraph-cli/setup.py | 2 +- trustgraph-embeddings-hf/setup.py | 4 ++-- trustgraph-flow/setup.py | 2 +- trustgraph-vertexai/setup.py | 2 +- trustgraph/setup.py | 12 ++++++------ 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/trustgraph-bedrock/setup.py b/trustgraph-bedrock/setup.py index 09785c16..58330aeb 100644 --- a/trustgraph-bedrock/setup.py +++ b/trustgraph-bedrock/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", "pulsar-client", "prometheus-client", "boto3", diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 7edffd4b..4ad9fa51 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", "requests", "pulsar-client", "rdflib", diff --git a/trustgraph-embeddings-hf/setup.py b/trustgraph-embeddings-hf/setup.py index 9af55352..baad468c 100644 --- a/trustgraph-embeddings-hf/setup.py +++ b/trustgraph-embeddings-hf/setup.py @@ -34,8 +34,8 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", - "trustgraph-flow>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", + "trustgraph-flow>=0.20,<0.21", "torch", "urllib3", "transformers", diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 83e4f4f7..83b027dd 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", "urllib3", "rdflib", "pymilvus", diff --git a/trustgraph-vertexai/setup.py b/trustgraph-vertexai/setup.py index 4bc2e716..e02160a6 100644 --- a/trustgraph-vertexai/setup.py +++ b/trustgraph-vertexai/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", "pulsar-client", "google-cloud-aiplatform", "prometheus-client", diff --git a/trustgraph/setup.py b/trustgraph/setup.py index 0b14cfac..3abcf1d7 100644 --- a/trustgraph/setup.py +++ b/trustgraph/setup.py @@ -34,12 +34,12 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.19,<0.20", - "trustgraph-bedrock>=0.19,<0.20", - "trustgraph-cli>=0.19,<0.20", - "trustgraph-embeddings-hf>=0.19,<0.20", - "trustgraph-flow>=0.19,<0.20", - "trustgraph-vertexai>=0.19,<0.20", + "trustgraph-base>=0.20,<0.21", + "trustgraph-bedrock>=0.20,<0.21", + "trustgraph-cli>=0.20,<0.21", + "trustgraph-embeddings-hf>=0.20,<0.21", + "trustgraph-flow>=0.20,<0.21", + "trustgraph-vertexai>=0.20,<0.21", ], scripts=[ ] From 8fc22155dd63ed24652ac01fca30d4f24a9a8dc0 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 20 Jan 2025 14:05:51 +0000 Subject: [PATCH 32/53] Agent tweaks (#272) - Removed unused LLM client configuration from agent-manager-react - Change agent-manager-react template to use prompt-rag instead of prompt - Changed TextCompletion tool to use 'question' instead of 'computation' for its parameter. --- .../components/agent-manager-react.jsonnet | 4 +++ .../trustgraph/agent/react/service.py | 29 ------------------- .../trustgraph/agent/react/tools.py | 2 +- 3 files changed, 5 insertions(+), 30 deletions(-) diff --git a/templates/components/agent-manager-react.jsonnet b/templates/components/agent-manager-react.jsonnet index a995dba5..5f252f8d 100644 --- a/templates/components/agent-manager-react.jsonnet +++ b/templates/components/agent-manager-react.jsonnet @@ -19,6 +19,10 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; "agent-manager-react", "-p", url.pulsar, + "--prompt-request-queue", + "non-persistent://tg/request/prompt-rag", + "--prompt-response-queue", + "non-persistent://tg/response/prompt-rag", "--tool-type", ] + [ tool.id + "=" + tool.type diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 8799816b..2414ea13 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -14,8 +14,6 @@ from ... schema import AgentRequest, AgentResponse, AgentStep from ... schema import agent_request_queue, agent_response_queue from ... schema import prompt_request_queue as pr_request_queue from ... schema import prompt_response_queue as pr_response_queue -from ... schema import text_completion_request_queue as tc_request_queue -from ... schema import text_completion_response_queue as tc_response_queue from ... schema import graph_rag_request_queue as gr_request_queue from ... schema import graph_rag_response_queue as gr_response_queue from ... clients.prompt_client import PromptClient @@ -133,12 +131,6 @@ class Processor(ConsumerProducer): prompt_response_queue = params.get( "prompt_response_queue", pr_response_queue ) - text_completion_request_queue = params.get( - "text_completion_request_queue", tc_request_queue - ) - text_completion_response_queue = params.get( - "text_completion_response_queue", tc_response_queue - ) graph_rag_request_queue = params.get( "graph_rag_request_queue", gr_request_queue ) @@ -155,8 +147,6 @@ class Processor(ConsumerProducer): "output_schema": AgentResponse, "prompt_request_queue": prompt_request_queue, "prompt_response_queue": prompt_response_queue, - "text_completion_request_queue": tc_request_queue, - "text_completion_response_queue": tc_response_queue, "graph_rag_request_queue": gr_request_queue, "graph_rag_response_queue": gr_response_queue, } @@ -169,13 +159,6 @@ class Processor(ConsumerProducer): pulsar_host = self.pulsar_host ) - self.llm = LlmClient( - subscriber=subscriber, - input_queue=text_completion_request_queue, - output_queue=text_completion_response_queue, - pulsar_host = self.pulsar_host - ) - self.graph_rag = GraphRagClient( subscriber=subscriber, input_queue=graph_rag_request_queue, @@ -340,18 +323,6 @@ class Processor(ConsumerProducer): help=f'Prompt response queue (default: {pr_response_queue})', ) - parser.add_argument( - '--text-completion-request-queue', - default=tc_request_queue, - help=f'Text completion request queue (default: {tc_request_queue})', - ) - - parser.add_argument( - '--text-completion-response-queue', - default=tc_response_queue, - help=f'Text completion response queue (default: {tc_response_queue})', - ) - parser.add_argument( '--graph-rag-request-queue', default=gr_request_queue, diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py index d9bc846f..941610be 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -14,6 +14,6 @@ class TextCompletionImpl: self.context = context def invoke(self, **arguments): return self.context.prompt.request( - "question", { "question": arguments.get("computation") } + "question", { "question": arguments.get("question") } ) From ac7bf6bfd5ef1beec7e60b169d9087d1bc0c6426 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 20 Jan 2025 19:17:43 +0000 Subject: [PATCH 33/53] Fix Minikube launch issues (#275) - Fix Pulsar service naming - broker -> pulsar - Fix Google AI Studio secret name for RAG to be consistent with other text-completion. --- templates/components/googleaistudio-rag.jsonnet | 2 +- templates/components/pulsar.jsonnet | 2 +- templates/generate | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/templates/components/googleaistudio-rag.jsonnet b/templates/components/googleaistudio-rag.jsonnet index 3bfa9618..96d80273 100644 --- a/templates/components/googleaistudio-rag.jsonnet +++ b/templates/components/googleaistudio-rag.jsonnet @@ -18,7 +18,7 @@ local prompts = import "prompts/mixtral.jsonnet"; create:: function(engine) - local envSecrets = engine.envSecrets("googleaistudio-key") + local envSecrets = engine.envSecrets("googleaistudio-credentials") .with_env_var("GOOGLE_AI_STUDIO_KEY", "googleaistudio-key"); local containerRag = diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index fdec2c33..59569be0 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -134,7 +134,7 @@ local url = import "values/url.jsonnet"; ); local bookieContainerSet = engine.containers( - "bookie", + "pulsar", [ bookieContainer, ] diff --git a/templates/generate b/templates/generate index e8772483..b2edf847 100755 --- a/templates/generate +++ b/templates/generate @@ -97,7 +97,7 @@ def main(): logger.info(f"Adding {name}...") out.writestr(name, content) - fname = "tg-launch.yaml" + fname = "docker-compose.yaml" platform = "docker-compose" From 8e227e256e9b3ac10ac6f1fb020ece4ef0e189e9 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 20 Jan 2025 23:54:33 +0000 Subject: [PATCH 34/53] Fix bookie & pulsar service naming in K8s (#276) --- templates/components/pulsar.jsonnet | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 59569be0..67e970c5 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -134,14 +134,14 @@ local url = import "values/url.jsonnet"; ); local bookieContainerSet = engine.containers( - "pulsar", + "bookie", [ bookieContainer, ] ); local brokerContainerSet = engine.containers( - "broker", + "pulsar", [ brokerContainer, ] From 605ade84f07250e43417e0ced304a71fc6b62182 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Thu, 23 Jan 2025 15:29:16 -0800 Subject: [PATCH 35/53] Added memgraph volume --- templates/stores/memgraph.jsonnet | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/templates/stores/memgraph.jsonnet b/templates/stores/memgraph.jsonnet index 75faf5f0..70ad127a 100644 --- a/templates/stores/memgraph.jsonnet +++ b/templates/stores/memgraph.jsonnet @@ -7,6 +7,8 @@ local images = import "values/images.jsonnet"; create:: function(engine) + local vol = engine.volume("memgraph").with_size("20G"); + local container = engine.container("memgraph") .with_image(images.memgraph_mage) @@ -16,7 +18,8 @@ local images = import "values/images.jsonnet"; .with_limits("1.0", "1000M") .with_reservations("0.5", "1000M") .with_port(7474, 7474, "api") - .with_port(7687, 7687, "api2"); + .with_port(7687, 7687, "api2") + .with_volume_mount(vol, "/var/lib/memgraph"); local containerSet = engine.containers( "memgraph", [ container ] @@ -28,6 +31,7 @@ local images = import "values/images.jsonnet"; .with_port(7687, 7687, "api2"); engine.resources([ + vol, containerSet, service, ]) @@ -65,4 +69,3 @@ local images = import "values/images.jsonnet"; }, } - From 26a586034cdea4760b7821eb4a2c1ac23c9ce5f4 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 17:04:09 +0000 Subject: [PATCH 36/53] Add cryptography dep (#281) --- trustgraph-flow/setup.py | 41 ++++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 83b027dd..bea1d496 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -35,32 +35,33 @@ setuptools.setup( download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ "trustgraph-base>=0.20,<0.21", - "urllib3", - "rdflib", - "pymilvus", - "langchain", - "langchain-core", - "langchain-text-splitters", - "langchain-community", - "requests", - "cassandra-driver", - "pulsar-client", - "pypdf", - "qdrant-client", - "tabulate", + "aiohttp", "anthropic", - "pyyaml", - "prometheus-client", + "cassandra-driver", "cohere", - "openai", - "neo4j", - "tiktoken", + "cryptography", + "falkordb", "google-generativeai", "ibis", "jsonschema", - "aiohttp", + "langchain", + "langchain-community", + "langchain-core", + "langchain-text-splitters", + "neo4j", + "openai", "pinecone[grpc]", - "falkordb", + "prometheus-client", + "pulsar-client", + "pymilvus", + "pypdf", + "pyyaml", + "qdrant-client", + "rdflib", + "requests", + "tabulate", + "tiktoken", + "urllib3", ], scripts=[ "scripts/agent-manager-react", From 0e03bc05a4fbb6ba82570b9f4af4abe30aa8aa5e Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 17:04:49 +0000 Subject: [PATCH 37/53] Refactor rate limit handling (#280) * - Refactored retry for rate limits into the base class - ConsumerProducer is derived from Consumer to simplify code - Added rate_limit_count metrics for rate limit events * Add rate limit events to VertexAI and Google AI Studio * Added Grafana rate limit dashboard * Add rate limit handling to all LLMs --- grafana/dashboards/dashboard.json | 4 +- trustgraph-base/trustgraph/base/consumer.py | 115 ++++++++++++++---- .../trustgraph/base/consumer_producer.py | 114 +++-------------- trustgraph-base/trustgraph/exceptions.py | 4 - .../model/text_completion/bedrock/llm.py | 1 + .../model/text_completion/azure/llm.py | 20 +-- .../model/text_completion/azure_openai/llm.py | 32 +++-- .../model/text_completion/claude/llm.py | 34 ++---- .../model/text_completion/cohere/llm.py | 22 +--- .../text_completion/googleaistudio/llm.py | 38 +++--- .../model/text_completion/llamafile/llm.py | 21 +--- .../model/text_completion/ollama/llm.py | 21 +--- .../model/text_completion/openai/llm.py | 26 +--- .../model/text_completion/vertexai/llm.py | 20 +-- 14 files changed, 174 insertions(+), 298 deletions(-) diff --git a/grafana/dashboards/dashboard.json b/grafana/dashboards/dashboard.json index 04561863..c484dffa 100644 --- a/grafana/dashboards/dashboard.json +++ b/grafana/dashboards/dashboard.json @@ -577,7 +577,7 @@ "disableTextWrap": false, "editorMode": "builder", "exemplar": false, - "expr": "increase(processing_count_total{status!=\"success\"}[$__rate_interval])", + "expr": "sum by(job) (increase(rate_limit_count_total[$__rate_interval]))", "format": "time_series", "fullMetaSearch": false, "includeNullMetadata": true, @@ -588,7 +588,7 @@ "useBackend": false } ], - "title": "Errors", + "title": "Rate limit events", "type": "timeseries" }, { diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index f346f1bc..521dd3c1 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -7,6 +7,9 @@ import time from . base_processor import BaseProcessor from .. exceptions import TooManyRequests +default_rate_limit_retry = 10 +default_rate_limit_timeout = 7200 + class Consumer(BaseProcessor): def __init__(self, **params): @@ -22,11 +25,18 @@ class Consumer(BaseProcessor): super(Consumer, self).__init__(**params) - input_queue = params.get("input_queue") - subscriber = params.get("subscriber") - input_schema = params.get("input_schema") + self.input_queue = params.get("input_queue") + self.subscriber = params.get("subscriber") + self.input_schema = params.get("input_schema") - if input_schema == None: + self.rate_limit_retry = params.get( + "rate_limit_retry", default_rate_limit_retry + ) + self.rate_limit_timeout = params.get( + "rate_limit_timeout", default_rate_limit_timeout + ) + + if self.input_schema == None: raise RuntimeError("input_schema must be specified") if not hasattr(__class__, "request_metric"): @@ -44,18 +54,27 @@ class Consumer(BaseProcessor): 'processing_count', 'Processing count', ["status"] ) + if not hasattr(__class__, "rate_limit_metric"): + __class__.rate_limit_metric = Counter( + 'rate_limit_count', 'Rate limit event count', + ) + __class__.pubsub_metric.info({ - "input_queue": input_queue, - "subscriber": subscriber, - "input_schema": input_schema.__name__, + "input_queue": self.input_queue, + "subscriber": self.subscriber, + "input_schema": self.input_schema.__name__, + "rate_limit_retry": str(self.rate_limit_retry), + "rate_limit_timeout": str(self.rate_limit_timeout), }) self.consumer = self.client.subscribe( - input_queue, subscriber, + self.input_queue, self.subscriber, consumer_type=pulsar.ConsumerType.Shared, - schema=JsonSchema(input_schema), + schema=JsonSchema(self.input_schema), ) + print("Initialised consumer.", flush=True) + def run(self): __class__.state_metric.state('running') @@ -64,31 +83,61 @@ class Consumer(BaseProcessor): msg = self.consumer.receive() - try: + expiry = time.time() + self.rate_limit_timeout - with __class__.request_metric.time(): - self.handle(msg) + # This loop is for retry on rate-limit / resource limits + while True: - # Acknowledge successful processing of the message - self.consumer.acknowledge(msg) + if time.time() > expiry: - __class__.processing_metric.labels(status="success").inc() + print("Gave up waiting for rate-limit retry", flush=True) - except TooManyRequests: - self.consumer.negative_acknowledge(msg) - print("TooManyRequests: will retry") - __class__.processing_metric.labels(status="rate-limit").inc() - time.sleep(5) - continue + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) + + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break + + try: + + with __class__.request_metric.time(): + self.handle(msg) + + # Acknowledge successful processing of the message + self.consumer.acknowledge(msg) + + __class__.processing_metric.labels(status="success").inc() + + # Break out of retry loop + break + + except TooManyRequests: + + print("TooManyRequests: will retry...", flush=True) + + __class__.rate_limit_metric.inc() + + # Sleep + time.sleep(self.rate_limit_retry) + + # Contine from retry loop, just causes a reprocessing + continue - except Exception as e: + except Exception as e: - print("Exception:", e, flush=True) + print("Exception:", e, flush=True) - # Message failed to be processed - self.consumer.negative_acknowledge(msg) + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) - __class__.processing_metric.labels(status="error").inc() + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break @staticmethod def add_args(parser, default_input_queue, default_subscriber): @@ -107,3 +156,17 @@ class Consumer(BaseProcessor): help=f'Queue subscriber name (default: {default_subscriber})' ) + parser.add_argument( + '--rate-limit-retry', + type=int, + default=default_rate_limit_retry, + help=f'Rate limit retry (default: {default_rate_limit_retry})' + ) + + parser.add_argument( + '--rate-limit-timeout', + type=int, + default=default_rate_limit_timeout, + help=f'Rate limit timeout (default: {default_rate_limit_timeout})' + ) + diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py index 6d386894..be9915ce 100644 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ b/trustgraph-base/trustgraph/base/consumer_producer.py @@ -4,111 +4,43 @@ import pulsar from prometheus_client import Histogram, Info, Counter, Enum import time -from . base_processor import BaseProcessor +from . consumer import Consumer from .. exceptions import TooManyRequests -# FIXME: Derive from consumer? And producer? - -class ConsumerProducer(BaseProcessor): +class ConsumerProducer(Consumer): def __init__(self, **params): - if not hasattr(__class__, "state_metric"): - __class__.state_metric = Enum( - 'processor_state', 'Processor state', - states=['starting', 'running', 'stopped'] - ) - __class__.state_metric.state('starting') + super(ConsumerProducer, self).__init__(**params) - __class__.state_metric.state('starting') - - input_queue = params.get("input_queue") - output_queue = params.get("output_queue") - subscriber = params.get("subscriber") - input_schema = params.get("input_schema") - output_schema = params.get("output_schema") - - if not hasattr(__class__, "request_metric"): - __class__.request_metric = Histogram( - 'request_latency', 'Request latency (seconds)' - ) + self.output_queue = params.get("output_queue") + self.output_schema = params.get("output_schema") if not hasattr(__class__, "output_metric"): __class__.output_metric = Counter( 'output_count', 'Output items created' ) - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration' - ) - - if not hasattr(__class__, "processing_metric"): - __class__.processing_metric = Counter( - 'processing_count', 'Processing count', ["status"] - ) - __class__.pubsub_metric.info({ - "input_queue": input_queue, - "output_queue": output_queue, - "subscriber": subscriber, - "input_schema": input_schema.__name__, - "output_schema": output_schema.__name__, + "input_queue": self.input_queue, + "output_queue": self.output_queue, + "subscriber": self.subscriber, + "input_schema": self.input_schema.__name__, + "output_schema": self.output_schema.__name__, + "rate_limit_retry": str(self.rate_limit_retry), + "rate_limit_timeout": str(self.rate_limit_timeout), }) - super(ConsumerProducer, self).__init__(**params) - - if input_schema == None: - raise RuntimeError("input_schema must be specified") - - if output_schema == None: + if self.output_schema == None: raise RuntimeError("output_schema must be specified") self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(output_schema), + topic=self.output_queue, + schema=JsonSchema(self.output_schema), chunking_enabled=True, ) - self.consumer = self.client.subscribe( - input_queue, subscriber, - consumer_type=pulsar.ConsumerType.Shared, - schema=JsonSchema(input_schema), - ) - - def run(self): - - __class__.state_metric.state('running') - - while True: - - msg = self.consumer.receive() - - try: - - with __class__.request_metric.time(): - resp = self.handle(msg) - - # Acknowledge successful processing of the message - self.consumer.acknowledge(msg) - - __class__.processing_metric.labels(status="success").inc() - - except TooManyRequests: - self.consumer.negative_acknowledge(msg) - print("TooManyRequests: will retry") - __class__.processing_metric.labels(status="rate-limit").inc() - time.sleep(5) - continue - - except Exception as e: - - print("Exception:", e, flush=True) - - # Message failed to be processed - self.consumer.negative_acknowledge(msg) - - __class__.processing_metric.labels(status="error").inc() + print("Initialised consumer/producer.") def send(self, msg, properties={}): self.producer.send(msg, properties) @@ -120,19 +52,7 @@ class ConsumerProducer(BaseProcessor): default_output_queue, ): - BaseProcessor.add_args(parser) - - parser.add_argument( - '-i', '--input-queue', - default=default_input_queue, - help=f'Input queue (default: {default_input_queue})' - ) - - parser.add_argument( - '-s', '--subscriber', - default=default_subscriber, - help=f'Queue subscriber name (default: {default_subscriber})' - ) + Consumer.add_args(parser, default_input_queue, default_subscriber) parser.add_argument( '-o', '--output-queue', diff --git a/trustgraph-base/trustgraph/exceptions.py b/trustgraph-base/trustgraph/exceptions.py index 16f9956c..afe72ccc 100644 --- a/trustgraph-base/trustgraph/exceptions.py +++ b/trustgraph-base/trustgraph/exceptions.py @@ -8,7 +8,3 @@ class LlmError(Exception): class ParseError(Exception): pass - - - - diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index a9c05cc8..3b2aced5 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -267,6 +267,7 @@ class Processor(ConsumerProducer): except Exception as e: + print(type(e)) print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index 4db7dbf1..90be6962 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -158,25 +158,15 @@ class Processor(ConsumerProducer): except TooManyRequests: - print("Send rate limit response...", flush=True) + print("Rate limit...") - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index a3edb859..f5ecb8d6 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -4,10 +4,9 @@ Simple LLM service, performs text prompt completion using the Azure OpenAI endpoit service. Input is prompt, output is response. """ -import requests import json from prometheus_client import Histogram -from openai import AzureOpenAI +from openai import AzureOpenAI, RateLimitError import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error @@ -126,30 +125,27 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp.choices[0].message.content, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) - self.producer.send(r, properties={"id": id}) - - except TooManyRequests: - - print("Send rate limit response...", flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, + response=resp.choices[0].message.content, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model ) self.producer.send(r, properties={"id": id}) - self.consumer.acknowledge(msg) + except RateLimitError: + + print("Send rate limit response...", flush=True) + + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 01ce837d..5cfd8907 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -87,8 +87,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits? - with __class__.text_completion_metric.time(): response = message = self.claude.messages.create( @@ -117,34 +115,26 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) + r = TextCompletionResponse( + response=resp, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ) self.send(r, properties={"id": id}) print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: + except anthropic.RateLimitError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index d03e1554..5b8e3ba9 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -112,27 +112,15 @@ class Processor(ConsumerProducer): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except TooManyRequests: + except cohere.TooManyRequestsError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index a249998d..5d5b23a0 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -88,7 +88,8 @@ class Processor(ConsumerProducer): HarmCategory.HARM_CATEGORY_HARASSMENT: block_level, HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: block_level, HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: block_level, - # There is a documentation conflict on whether or not CIVIC_INTEGRITY is a valid category + # There is a documentation conflict on whether or not + # CIVIC_INTEGRITY is a valid category # HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level, } @@ -122,8 +123,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits? - with __class__.text_completion_metric.time(): chat_session = self.llm.start_chat( @@ -140,35 +139,30 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) + r = TextCompletionResponse( + response=resp, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ) self.send(r, properties={"id": id}) print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit except ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the default handler + raise TooManyRequests() except Exception as e: - print(f"Exception: {e}") + # Apart from rate limits, treat all exceptions as unrecoverable + + print(type(e), flush=True) + print(f"Exception: {e}", flush=True) print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index 274948a8..65a2b171 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -126,26 +126,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: - - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # SLM, presumably there aren't rate limits except Exception as e: diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 00d44f6d..8c5bd3dc 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -100,26 +100,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: - - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # SLM, presumably no rate limits except Exception as e: diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index c874943e..c2b948d5 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -4,7 +4,7 @@ Simple LLM service, performs text prompt completion using OpenAI. Input is prompt, output is response. """ -from openai import OpenAI +from openai import OpenAI, RateLimitError from prometheus_client import Histogram import os @@ -87,8 +87,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits - with __class__.text_completion_metric.time(): resp = self.openai.chat.completions.create( @@ -134,27 +132,15 @@ class Processor(ConsumerProducer): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except TooManyRequests: + except openai.RateLimitError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index cb817836..d6a2efec 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -178,25 +178,15 @@ class Processor(ConsumerProducer): except google.api_core.exceptions.ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) From 0ebd770cb5d2f31b2d5c22c3b27e9fd249748392 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 19:58:38 +0000 Subject: [PATCH 38/53] Back out cryptography (#282) --- trustgraph-flow/setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index bea1d496..66adb7b8 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -39,7 +39,6 @@ setuptools.setup( "anthropic", "cassandra-driver", "cohere", - "cryptography", "falkordb", "google-generativeai", "ibis", From 5a5daa527fa548ce536a2ee2d53e4f332937cd7e Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 19:59:08 +0000 Subject: [PATCH 39/53] Load default embeddings model into container (#283) --- Containerfile | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Containerfile b/Containerfile index 73c9285f..ffb897fa 100644 --- a/Containerfile +++ b/Containerfile @@ -20,6 +20,10 @@ RUN pip3 install anthropic boto3 cohere openai google-cloud-aiplatform ollama go neo4j tiktoken falkordb && \ pip3 cache purge +# Most commonly used embeddings model, just build it into the container +# image +RUN huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 + # ---------------------------------------------------------------------------- # Build a container which contains the built Python packages. The build # creates a bunch of left-over cruft, a separate phase means this is only From 552637c1f710223bc5a655ac410ef037656e8f41 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 21:57:22 +0000 Subject: [PATCH 40/53] Fix deps (#284) * Add cryptography module back in * Lock in known working version of the AI modules, suspecting a problem in sentence-transformers with Python 3.12? --- Containerfile | 17 ++++++++++++----- trustgraph-flow/setup.py | 1 + 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/Containerfile b/Containerfile index ffb897fa..ee4e9fff 100644 --- a/Containerfile +++ b/Containerfile @@ -11,12 +11,19 @@ ENV PIP_BREAK_SYSTEM_PACKAGES=1 RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ python3-rdflib -RUN pip3 install torch --index-url https://download.pytorch.org/whl/cpu +RUN pip3 install torch==2.5.1+cpu \ + --index-url https://download.pytorch.org/whl/cpu -RUN pip3 install anthropic boto3 cohere openai google-cloud-aiplatform ollama google-generativeai \ - langchain langchain-core langchain-huggingface langchain-text-splitters \ - langchain-community pymilvus sentence-transformers transformers \ - huggingface-hub pulsar-client cassandra-driver pyyaml \ +RUN pip3 install \ + anthropic boto3 cohere openai google-cloud-aiplatform \ + ollama google-generativeai \ + langchain==0.3.13 langchain-core==0.3.28 langchain-huggingface==0.1.2 \ + langchain-text-splitters==0.3.4 \ + langchain-community==0.3.13 \ + sentence-transformers==3.4.0 transformers==4.47.1 \ + huggingface-hub==0.27.0 \ + pymilvus \ + pulsar-client==3.5.0 cassandra-driver pyyaml \ neo4j tiktoken falkordb && \ pip3 cache purge diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 66adb7b8..bea1d496 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -39,6 +39,7 @@ setuptools.setup( "anthropic", "cassandra-driver", "cohere", + "cryptography", "falkordb", "google-generativeai", "ibis", From 75a72b0d2d87547f565e1bc4e393b51b705e510f Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Mon, 27 Jan 2025 23:47:15 +0000 Subject: [PATCH 41/53] Fix ollama embeddings client to work (#285) --- trustgraph-flow/setup.py | 1 + .../trustgraph/embeddings/ollama/processor.py | 23 +++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index bea1d496..f5a0bc3c 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -49,6 +49,7 @@ setuptools.setup( "langchain-core", "langchain-text-splitters", "neo4j", + "ollama", "openai", "pinecone[grpc]", "prometheus-client", diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index 6682a79f..fc54cbb8 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -3,12 +3,13 @@ Embeddings service, applies an embeddings model selected from HuggingFace. Input is text, output is embeddings vector. """ -from langchain_community.embeddings import OllamaEmbeddings from ... schema import EmbeddingsRequest, EmbeddingsResponse from ... schema import embeddings_request_queue, embeddings_response_queue from ... log_level import LogLevel from ... base import ConsumerProducer +from ollama import Client +import os module = ".".join(__name__.split(".")[1:-1]) @@ -16,7 +17,7 @@ default_input_queue = embeddings_request_queue default_output_queue = embeddings_response_queue default_subscriber = module default_model="mxbai-embed-large" -default_ollama = 'http://localhost:11434' +default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434') class Processor(ConsumerProducer): @@ -26,6 +27,9 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) + ollama = params.get("ollama", default_ollama) + model = params.get("model", default_model) + super(Processor, self).__init__( **params | { "input_queue": input_queue, @@ -33,10 +37,13 @@ class Processor(ConsumerProducer): "subscriber": subscriber, "input_schema": EmbeddingsRequest, "output_schema": EmbeddingsResponse, + "ollama": ollama, + "model": model, } ) - self.embeddings = OllamaEmbeddings(base_url=ollama, model=model) + self.client = Client(host=ollama) + self.model = model def handle(self, msg): @@ -49,10 +56,16 @@ class Processor(ConsumerProducer): print(f"Handling input {id}...", flush=True) text = v.text - embeds = self.embeddings.embed_query([text]) + embeds = self.client.embed( + model = self.model, + input = text + ) print("Send response...", flush=True) - r = EmbeddingsResponse(vectors=[embeds]) + r = EmbeddingsResponse( + vectors=embeds.embeddings, + error=None, + ) self.producer.send(r, properties={"id": id}) From 6c3d2e7f9756bf56822e21598a84cbfa8e17c2df Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Tue, 28 Jan 2025 18:26:17 +0000 Subject: [PATCH 42/53] Feature/fastembed (#286) * Copied in base files for fastembed * Added fastembed implementation * Added template support for fastembed --- .../components/embeddings-fastembed.jsonnet | 43 +++++++++ trustgraph-flow/scripts/embeddings-fastembed | 6 ++ trustgraph-flow/setup.py | 2 + .../embeddings/fastembed/__init__.py | 3 + .../embeddings/fastembed/__main__.py | 7 ++ .../embeddings/fastembed/processor.py | 89 +++++++++++++++++++ .../trustgraph/embeddings/ollama/processor.py | 2 +- 7 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 templates/components/embeddings-fastembed.jsonnet create mode 100755 trustgraph-flow/scripts/embeddings-fastembed create mode 100644 trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py create mode 100755 trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py create mode 100755 trustgraph-flow/trustgraph/embeddings/fastembed/processor.py diff --git a/templates/components/embeddings-fastembed.jsonnet b/templates/components/embeddings-fastembed.jsonnet new file mode 100644 index 00000000..a515617f --- /dev/null +++ b/templates/components/embeddings-fastembed.jsonnet @@ -0,0 +1,43 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + "embeddings-model":: "sentence-transformers/all-MiniLM-L6-v2", + + embeddings +: { + + create:: function(engine) + + local container = + engine.container("embeddings") + .with_image(images.trustgraph) + .with_command([ + "embeddings-fastembed", + "-p", + url.pulsar, + "-m", + $["embeddings-model"], + ]) + .with_limits("1.0", "400M") + .with_reservations("0.5", "400M"); + + local containerSet = engine.containers( + "embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + +} + diff --git a/trustgraph-flow/scripts/embeddings-fastembed b/trustgraph-flow/scripts/embeddings-fastembed new file mode 100755 index 00000000..e1322269 --- /dev/null +++ b/trustgraph-flow/scripts/embeddings-fastembed @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.fastembed import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index f5a0bc3c..a6191cf5 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -41,6 +41,7 @@ setuptools.setup( "cohere", "cryptography", "falkordb", + "fastembed", "google-generativeai", "ibis", "jsonschema", @@ -78,6 +79,7 @@ setuptools.setup( "scripts/document-embeddings", "scripts/document-rag", "scripts/embeddings-ollama", + "scripts/embeddings-fastembed", "scripts/ge-query-milvus", "scripts/ge-query-pinecone", "scripts/ge-query-qdrant", diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py b/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py new file mode 100644 index 00000000..9d16af90 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py @@ -0,0 +1,3 @@ + +from . processor import * + diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py b/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py new file mode 100755 index 00000000..986c0257 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . processor import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py new file mode 100755 index 00000000..635387b8 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -0,0 +1,89 @@ + +""" +Embeddings service, applies an embeddings model selected from HuggingFace. +Input is text, output is embeddings vector. +""" + +from ... schema import EmbeddingsRequest, EmbeddingsResponse +from ... schema import embeddings_request_queue, embeddings_response_queue +from ... log_level import LogLevel +from ... base import ConsumerProducer +from fastembed import TextEmbedding +import os + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = embeddings_request_queue +default_output_queue = embeddings_response_queue +default_subscriber = module +default_model="sentence-transformers/all-MiniLM-L6-v2" + +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) + + model = params.get("model", default_model) + + super(Processor, self).__init__( + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": EmbeddingsRequest, + "output_schema": EmbeddingsResponse, + "model": model, + } + ) + + self.embeddings = TextEmbedding(model_name = model) + + def handle(self, msg): + + v = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + text = v.text + vecs = self.embeddings.embed([text]) + + vecs = [ + v.tolist() + for v in vecs + ] + + print("Send response...", flush=True) + r = EmbeddingsResponse( + vectors=list(vecs), + error=None, + ) + + self.producer.send(r, properties={"id": id}) + + print("Done.", flush=True) + + @staticmethod + def add_args(parser): + + ConsumerProducer.add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ) + + parser.add_argument( + '-m', '--model', + default=default_model, + help=f'Embeddings model (default: {default_model})' + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index fc54cbb8..5baf64aa 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -1,6 +1,6 @@ """ -Embeddings service, applies an embeddings model selected from HuggingFace. +Embeddings service, applies an embeddings model hosted on a local Ollama. Input is text, output is embeddings vector. """ From edcdc4d59d395b2c3369e48efdda9aab6549ba77 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Tue, 28 Jan 2025 19:36:05 +0000 Subject: [PATCH 43/53] Feature/separate containers (#287) * Separate containerfiles * Add push to Makefile * Update image names in the templates --- Makefile | 20 ++++- containers/Containerfile.base | 48 ++++++++++++ containers/Containerfile.bedrock | 48 ++++++++++++ containers/Containerfile.flow | 60 +++++++++++++++ containers/Containerfile.hf | 75 +++++++++++++++++++ containers/Containerfile.vertexai | 51 +++++++++++++ .../components/agent-manager-react.jsonnet | 2 +- templates/components/azure-openai-rag.jsonnet | 2 +- templates/components/azure-openai.jsonnet | 2 +- templates/components/azure-rag.jsonnet | 2 +- templates/components/azure.jsonnet | 2 +- templates/components/bedrock-rag.jsonnet | 2 +- templates/components/bedrock.jsonnet | 2 +- templates/components/cassandra.jsonnet | 4 +- .../components/chunker-recursive.jsonnet | 2 +- templates/components/claude-rag.jsonnet | 2 +- templates/components/claude.jsonnet | 2 +- templates/components/cohere-rag.jsonnet | 2 +- templates/components/cohere.jsonnet | 2 +- templates/components/document-rag.jsonnet | 4 +- .../components/embeddings-fastembed.jsonnet | 2 +- templates/components/embeddings-hf.jsonnet | 2 +- .../components/embeddings-ollama.jsonnet | 2 +- templates/components/falkordb.jsonnet | 4 +- .../components/googleaistudio-rag.jsonnet | 2 +- templates/components/googleaistudio.jsonnet | 2 +- templates/components/graph-rag.jsonnet | 10 +-- templates/components/llamafile-rag.jsonnet | 2 +- templates/components/llamafile.jsonnet | 2 +- templates/components/memgraph.jsonnet | 4 +- templates/components/milvus.jsonnet | 8 +- templates/components/neo4j.jsonnet | 4 +- templates/components/ollama-rag.jsonnet | 2 +- templates/components/ollama.jsonnet | 2 +- templates/components/openai-rag.jsonnet | 2 +- templates/components/openai.jsonnet | 2 +- templates/components/pinecone.jsonnet | 8 +- templates/components/prompt-template.jsonnet | 4 +- templates/components/pulsar.jsonnet | 2 +- templates/components/qdrant.jsonnet | 8 +- templates/components/trustgraph.jsonnet | 10 +-- templates/components/vertexai-rag.jsonnet | 2 +- templates/components/vertexai.jsonnet | 2 +- templates/values/images.jsonnet | 6 +- 44 files changed, 363 insertions(+), 65 deletions(-) create mode 100644 containers/Containerfile.base create mode 100644 containers/Containerfile.bedrock create mode 100644 containers/Containerfile.flow create mode 100644 containers/Containerfile.hf create mode 100644 containers/Containerfile.vertexai diff --git a/Makefile b/Makefile index f408a31f..80cdefc7 100644 --- a/Makefile +++ b/Makefile @@ -30,7 +30,7 @@ packages: update-package-versions pypi-upload: twine upload dist/*-${VERSION}.* -CONTAINER=docker.io/trustgraph/trustgraph-flow +CONTAINER_BASE=docker.io/trustgraph update-package-versions: mkdir -p trustgraph-cli/trustgraph @@ -44,11 +44,23 @@ update-package-versions: echo __version__ = \"${VERSION}\" > trustgraph/trustgraph/trustgraph_version.py container: update-package-versions - ${DOCKER} build -f Containerfile -t ${CONTAINER}:${VERSION} \ - --format docker + ${DOCKER} build -f containers/Containerfile.base \ + -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} . + ${DOCKER} build -f containers/Containerfile.flow \ + -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} . + ${DOCKER} build -f containers/Containerfile.bedrock \ + -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} . + ${DOCKER} build -f containers/Containerfile.vertexai \ + -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . + ${DOCKER} build -f containers/Containerfile.hf \ + -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} . push: - ${DOCKER} push ${CONTAINER}:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-base:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-flow:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-hf:${VERSION} clean: rm -rf wheels/ diff --git a/containers/Containerfile.base b/containers/Containerfile.base new file mode 100644 index 00000000..b4f5bbbf --- /dev/null +++ b/containers/Containerfile.base @@ -0,0 +1,48 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp && \ + dnf clean all + +RUN pip3 install --no-cache-dir pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_cli-* && \ + rm -rf /root/wheels + +WORKDIR / + + + diff --git a/containers/Containerfile.bedrock b/containers/Containerfile.bedrock new file mode 100644 index 00000000..21819973 --- /dev/null +++ b/containers/Containerfile.bedrock @@ -0,0 +1,48 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir boto3 pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-bedrock/ /root/build/trustgraph-bedrock/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-bedrock/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_bedrock-* && \ + rm -rf /root/wheels + +WORKDIR / + diff --git a/containers/Containerfile.flow b/containers/Containerfile.flow new file mode 100644 index 00000000..8d47effe --- /dev/null +++ b/containers/Containerfile.flow @@ -0,0 +1,60 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir \ + anthropic cohere openai google-generativeai \ + ollama \ + langchain==0.3.13 langchain-core==0.3.28 \ + langchain-text-splitters==0.3.4 \ + langchain-community==0.3.13 \ + pymilvus \ + pulsar-client==3.5.0 cassandra-driver pyyaml \ + neo4j tiktoken falkordb && \ + pip3 cache purge + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-flow/ /root/build/trustgraph-flow/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-flow/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_flow-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_cli-* && \ + rm -rf /root/wheels + +WORKDIR / + diff --git a/containers/Containerfile.hf b/containers/Containerfile.hf new file mode 100644 index 00000000..4076db28 --- /dev/null +++ b/containers/Containerfile.hf @@ -0,0 +1,75 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS ai + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install torch==2.5.1+cpu \ + --index-url https://download.pytorch.org/whl/cpu + +RUN pip3 install --no-cache-dir \ + langchain==0.3.13 langchain-core==0.3.28 langchain-huggingface==0.1.2 \ + langchain-community==0.3.13 \ + sentence-transformers==3.4.0 transformers==4.47.1 \ + huggingface-hub==0.27.0 \ + pulsar-client==3.5.0 + +# Most commonly used embeddings model, just build it into the container +# image +RUN huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM ai AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-flow/ /root/build/trustgraph-flow/ +COPY trustgraph-vertexai/ /root/build/trustgraph-vertexai/ +COPY trustgraph-bedrock/ /root/build/trustgraph-bedrock/ +COPY trustgraph-embeddings-hf/ /root/build/trustgraph-embeddings-hf/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-flow/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-vertexai/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-bedrock/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-embeddings-hf/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM ai + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install /root/wheels/trustgraph_base-* && \ + pip3 install /root/wheels/trustgraph_flow-* && \ + pip3 install /root/wheels/trustgraph_vertexai-* && \ + pip3 install /root/wheels/trustgraph_bedrock-* && \ + pip3 install /root/wheels/trustgraph_embeddings_hf-* && \ + pip3 install /root/wheels/trustgraph_cli-* && \ + pip3 cache purge && \ + rm -rf /root/wheels + +WORKDIR / + +CMD sleep 1000000 + diff --git a/containers/Containerfile.vertexai b/containers/Containerfile.vertexai new file mode 100644 index 00000000..72d21bde --- /dev/null +++ b/containers/Containerfile.vertexai @@ -0,0 +1,51 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir \ + google-cloud-aiplatform pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-vertexai/ /root/build/trustgraph-vertexai/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-vertexai/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_vertexai-* && \ + rm -rf /root/wheels + +WORKDIR / + + + diff --git a/templates/components/agent-manager-react.jsonnet b/templates/components/agent-manager-react.jsonnet index 5f252f8d..672a0439 100644 --- a/templates/components/agent-manager-react.jsonnet +++ b/templates/components/agent-manager-react.jsonnet @@ -14,7 +14,7 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("agent-manager") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "agent-manager-react", "-p", diff --git a/templates/components/azure-openai-rag.jsonnet b/templates/components/azure-openai-rag.jsonnet index fab83266..33355707 100644 --- a/templates/components/azure-openai-rag.jsonnet +++ b/templates/components/azure-openai-rag.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure", "-p", diff --git a/templates/components/azure-openai.jsonnet b/templates/components/azure-openai.jsonnet index f6bda306..3ecbbdac 100644 --- a/templates/components/azure-openai.jsonnet +++ b/templates/components/azure-openai.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure-openai", "-p", diff --git a/templates/components/azure-rag.jsonnet b/templates/components/azure-rag.jsonnet index da2f1692..20b7306e 100644 --- a/templates/components/azure-rag.jsonnet +++ b/templates/components/azure-rag.jsonnet @@ -22,7 +22,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure", "-p", diff --git a/templates/components/azure.jsonnet b/templates/components/azure.jsonnet index bd364d8d..c7746e23 100644 --- a/templates/components/azure.jsonnet +++ b/templates/components/azure.jsonnet @@ -22,7 +22,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure", "-p", diff --git a/templates/components/bedrock-rag.jsonnet b/templates/components/bedrock-rag.jsonnet index 68aa9b0a..bc94c8ee 100644 --- a/templates/components/bedrock-rag.jsonnet +++ b/templates/components/bedrock-rag.jsonnet @@ -26,7 +26,7 @@ local chunker = import "chunker-recursive.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_bedrock) .with_command([ "text-completion-bedrock", "-p", diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet index 922e0409..bb70d55f 100644 --- a/templates/components/bedrock.jsonnet +++ b/templates/components/bedrock.jsonnet @@ -26,7 +26,7 @@ local chunker = import "chunker-recursive.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_bedrock) .with_command([ "text-completion-bedrock", "-p", diff --git a/templates/components/cassandra.jsonnet b/templates/components/cassandra.jsonnet index b52d4b04..92ecf69f 100644 --- a/templates/components/cassandra.jsonnet +++ b/templates/components/cassandra.jsonnet @@ -12,7 +12,7 @@ cassandra + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-cassandra", "-p", @@ -44,7 +44,7 @@ cassandra + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-cassandra", "-p", diff --git a/templates/components/chunker-recursive.jsonnet b/templates/components/chunker-recursive.jsonnet index 0b64b712..4a174366 100644 --- a/templates/components/chunker-recursive.jsonnet +++ b/templates/components/chunker-recursive.jsonnet @@ -14,7 +14,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("chunker") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "chunker-recursive", "-p", diff --git a/templates/components/claude-rag.jsonnet b/templates/components/claude-rag.jsonnet index 71a214fc..06d58db2 100644 --- a/templates/components/claude-rag.jsonnet +++ b/templates/components/claude-rag.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-claude", "-p", diff --git a/templates/components/claude.jsonnet b/templates/components/claude.jsonnet index fd1650d2..e43e7504 100644 --- a/templates/components/claude.jsonnet +++ b/templates/components/claude.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-claude", "-p", diff --git a/templates/components/cohere-rag.jsonnet b/templates/components/cohere-rag.jsonnet index 62829445..6a142519 100644 --- a/templates/components/cohere-rag.jsonnet +++ b/templates/components/cohere-rag.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-cohere", "-p", diff --git a/templates/components/cohere.jsonnet b/templates/components/cohere.jsonnet index 2f559ad0..093436fd 100644 --- a/templates/components/cohere.jsonnet +++ b/templates/components/cohere.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-cohere", "-p", diff --git a/templates/components/document-rag.jsonnet b/templates/components/document-rag.jsonnet index 11dc9c13..ec125ed5 100644 --- a/templates/components/document-rag.jsonnet +++ b/templates/components/document-rag.jsonnet @@ -11,7 +11,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("document-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "document-rag", "-p", @@ -45,7 +45,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("document-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "document-embeddings", "-p", diff --git a/templates/components/embeddings-fastembed.jsonnet b/templates/components/embeddings-fastembed.jsonnet index a515617f..c1fe35ff 100644 --- a/templates/components/embeddings-fastembed.jsonnet +++ b/templates/components/embeddings-fastembed.jsonnet @@ -13,7 +13,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "embeddings-fastembed", "-p", diff --git a/templates/components/embeddings-hf.jsonnet b/templates/components/embeddings-hf.jsonnet index b46feac7..29ebbc48 100644 --- a/templates/components/embeddings-hf.jsonnet +++ b/templates/components/embeddings-hf.jsonnet @@ -13,7 +13,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_hf) .with_command([ "embeddings-hf", "-p", diff --git a/templates/components/embeddings-ollama.jsonnet b/templates/components/embeddings-ollama.jsonnet index 425a1c47..a26ad0ba 100644 --- a/templates/components/embeddings-ollama.jsonnet +++ b/templates/components/embeddings-ollama.jsonnet @@ -13,7 +13,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "embeddings-ollama", "-p", diff --git a/templates/components/falkordb.jsonnet b/templates/components/falkordb.jsonnet index e238cebe..c08896d3 100644 --- a/templates/components/falkordb.jsonnet +++ b/templates/components/falkordb.jsonnet @@ -13,7 +13,7 @@ falkordb + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-falkordb", "-p", @@ -45,7 +45,7 @@ falkordb + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-falkordb", "-p", diff --git a/templates/components/googleaistudio-rag.jsonnet b/templates/components/googleaistudio-rag.jsonnet index 96d80273..332749e8 100644 --- a/templates/components/googleaistudio-rag.jsonnet +++ b/templates/components/googleaistudio-rag.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-googleaistudio", "-p", diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet index ac122243..58c7807d 100644 --- a/templates/components/googleaistudio.jsonnet +++ b/templates/components/googleaistudio.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-googleaistudio", "-p", diff --git a/templates/components/graph-rag.jsonnet b/templates/components/graph-rag.jsonnet index eb72754e..27035b35 100644 --- a/templates/components/graph-rag.jsonnet +++ b/templates/components/graph-rag.jsonnet @@ -14,7 +14,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-definitions") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-definitions", "-p", @@ -44,7 +44,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-relationships") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-relationships", "-p", @@ -74,7 +74,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-topics") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-topics", "-p", @@ -104,7 +104,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("graph-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "graph-rag", "-p", @@ -144,7 +144,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "graph-embeddings", "-p", diff --git a/templates/components/llamafile-rag.jsonnet b/templates/components/llamafile-rag.jsonnet index 34af69d0..262f586e 100644 --- a/templates/components/llamafile-rag.jsonnet +++ b/templates/components/llamafile-rag.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/slm.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-llamafile", "-p", diff --git a/templates/components/llamafile.jsonnet b/templates/components/llamafile.jsonnet index bc16dfaa..f3e1efd3 100644 --- a/templates/components/llamafile.jsonnet +++ b/templates/components/llamafile.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/slm.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-llamafile", "-p", diff --git a/templates/components/memgraph.jsonnet b/templates/components/memgraph.jsonnet index 609da3a2..21684a61 100644 --- a/templates/components/memgraph.jsonnet +++ b/templates/components/memgraph.jsonnet @@ -14,7 +14,7 @@ memgraph + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-memgraph", "-p", @@ -48,7 +48,7 @@ memgraph + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-memgraph", "-p", diff --git a/templates/components/milvus.jsonnet b/templates/components/milvus.jsonnet index b3044f98..27e5e316 100644 --- a/templates/components/milvus.jsonnet +++ b/templates/components/milvus.jsonnet @@ -12,7 +12,7 @@ milvus + { local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-milvus", "-p", @@ -44,7 +44,7 @@ milvus + { local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-milvus", "-p", @@ -76,7 +76,7 @@ milvus + { local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-milvus", "-p", @@ -108,7 +108,7 @@ milvus + { local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-milvus", "-p", diff --git a/templates/components/neo4j.jsonnet b/templates/components/neo4j.jsonnet index b70562fe..7cebdc71 100644 --- a/templates/components/neo4j.jsonnet +++ b/templates/components/neo4j.jsonnet @@ -13,7 +13,7 @@ neo4j + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-neo4j", "-p", @@ -45,7 +45,7 @@ neo4j + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-neo4j", "-p", diff --git a/templates/components/ollama-rag.jsonnet b/templates/components/ollama-rag.jsonnet index e3ca97aa..680adea5 100644 --- a/templates/components/ollama-rag.jsonnet +++ b/templates/components/ollama-rag.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-ollama", "-p", diff --git a/templates/components/ollama.jsonnet b/templates/components/ollama.jsonnet index 28bf36b9..95f1abf0 100644 --- a/templates/components/ollama.jsonnet +++ b/templates/components/ollama.jsonnet @@ -21,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-ollama", "-p", diff --git a/templates/components/openai-rag.jsonnet b/templates/components/openai-rag.jsonnet index d9ecd964..bfb7dd98 100644 --- a/templates/components/openai-rag.jsonnet +++ b/templates/components/openai-rag.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local containerRag = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-openai", "-p", diff --git a/templates/components/openai.jsonnet b/templates/components/openai.jsonnet index 50194bfc..9e0212d2 100644 --- a/templates/components/openai.jsonnet +++ b/templates/components/openai.jsonnet @@ -23,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-openai", "-p", diff --git a/templates/components/pinecone.jsonnet b/templates/components/pinecone.jsonnet index 3422952a..ede383a5 100644 --- a/templates/components/pinecone.jsonnet +++ b/templates/components/pinecone.jsonnet @@ -17,7 +17,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-pinecone", "-p", @@ -52,7 +52,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-pinecone", "-p", @@ -87,7 +87,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-pinecone", "-p", @@ -122,7 +122,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-pinecone", "-p", diff --git a/templates/components/prompt-template.jsonnet b/templates/components/prompt-template.jsonnet index 3dadf337..b3187c9b 100644 --- a/templates/components/prompt-template.jsonnet +++ b/templates/components/prompt-template.jsonnet @@ -44,7 +44,7 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("prompt") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "prompt-template", "-p", @@ -84,7 +84,7 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("prompt-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "prompt-template", "-p", diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 67e970c5..d111f616 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -109,7 +109,7 @@ local url = import "values/url.jsonnet"; // Trustgraph Pulsar initialisation local adminContainer = engine.container("init-trustgraph") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "tg-init-pulsar", "-p", diff --git a/templates/components/qdrant.jsonnet b/templates/components/qdrant.jsonnet index f923e84f..352cb741 100644 --- a/templates/components/qdrant.jsonnet +++ b/templates/components/qdrant.jsonnet @@ -12,7 +12,7 @@ qdrant + { local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-qdrant", "-p", @@ -44,7 +44,7 @@ qdrant + { local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-qdrant", "-p", @@ -76,7 +76,7 @@ qdrant + { local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-qdrant", "-p", @@ -108,7 +108,7 @@ qdrant + { local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-qdrant", "-p", diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index f63740aa..833d932b 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -21,7 +21,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("api-gateway") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "api-gateway", "-p", @@ -60,7 +60,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("chunker") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "chunker-token", "-p", @@ -94,7 +94,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("pdf-decoder") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "pdf-decoder", "-p", @@ -124,7 +124,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("metering") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "metering", "-p", @@ -154,7 +154,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("metering-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "metering", "-p", diff --git a/templates/components/vertexai-rag.jsonnet b/templates/components/vertexai-rag.jsonnet index 9c7eaa00..0b5cf9a3 100644 --- a/templates/components/vertexai-rag.jsonnet +++ b/templates/components/vertexai-rag.jsonnet @@ -30,7 +30,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_vertexai) .with_command([ "text-completion-vertexai", "-p", diff --git a/templates/components/vertexai.jsonnet b/templates/components/vertexai.jsonnet index 49243701..0e3550c5 100644 --- a/templates/components/vertexai.jsonnet +++ b/templates/components/vertexai.jsonnet @@ -30,7 +30,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_vertexai) .with_command([ "text-completion-vertexai", "-p", diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index b0416eb3..d515a450 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -9,7 +9,11 @@ local version = import "version.jsonnet"; milvus: "docker.io/milvusdb/milvus:v2.4.9", prometheus: "docker.io/prom/prometheus:v2.53.2", grafana: "docker.io/grafana/grafana:11.1.4", - trustgraph: "docker.io/trustgraph/trustgraph-flow:" + version, + trustgraph_base: "docker.io/trustgraph/trustgraph-base:" + version, + trustgraph_flow: "docker.io/trustgraph/trustgraph-flow:" + version, + trustgraph_bedrock: "docker.io/trustgraph/trustgraph-bedrock:" + version, + trustgraph_vertexai: "docker.io/trustgraph/trustgraph-vertexai:" + version, + trustgraph_hf: "docker.io/trustgraph/trustgraph-hf:" + version, qdrant: "docker.io/qdrant/qdrant:v1.11.1", memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", From 34bde24086c784e07d829162d6ba3cd2067f38d9 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 29 Jan 2025 12:40:46 +0000 Subject: [PATCH 44/53] Update generate script to handle more cases (#290) --- templates/generate | 198 +++++++++++++++++++++++++++++++++------------ 1 file changed, 145 insertions(+), 53 deletions(-) diff --git a/templates/generate b/templates/generate index b2edf847..6f915cda 100755 --- a/templates/generate +++ b/templates/generate @@ -7,6 +7,8 @@ import logging import os import sys import zipfile +import pathlib +from io import BytesIO logger = logging.getLogger("generate") logging.basicConfig(level=logging.INFO, format='%(message)s') @@ -15,9 +17,13 @@ private_json = "Put your GCP private.json here" class Generator: - def __init__(self, config, base="./templates/", version="0.0.0"): + def __init__( + self, config, templates="./templates/", resources="./resources", + version="0.0.0", + ): - self.jsonnet_base = base + self.templates = pathlib.Path(templates) + self.resources = pathlib.Path(resources) self.config = config self.version = f"\"{version}\"".encode("utf-8") @@ -34,25 +40,27 @@ class Generator: path = os.path.join(".", dir, filename) return str(path), self.config - if filename == "version.jsonnet" and dir == "./templates/values/": + if filename == "version.jsonnet" and dir == "templates/values/": path = os.path.join(".", dir, filename) return str(path), self.version if dir: candidates = [ - os.path.join(".", dir, filename), - os.path.join(".", filename) + self.templates.joinpath(dir, filename), + self.templates.joinpath(filename), + self.resources.joinpath(dir, filename), + self.resources.joinpath(filename), ] else: candidates = [ - os.path.join(".", filename) + self.templates.joinpath(filename), ] try: if filename == "vertexai/private.json": - return candidates[0], private_json.encode("utf-8") + return str(candidates[0]), private_json.encode("utf-8") for c in candidates: logger.debug("Try: %s", c) @@ -68,73 +76,157 @@ class Generator: except: - path = os.path.join(self.jsonnet_base, filename) + path = os.path.join(self.templates, filename) logger.debug("Try: %s", path) with open(path, "rb") as f: logger.debug("Loaded: %s", path) return str(path), f.read() +class Packager: + + def __init__(self): + self.templates = pathlib.Path("./templates") + self.resources = pathlib.Path("./resources") + + def process( + self, config, version="0.0.0", platform="docker-compose", + ): + + config = config.encode("utf-8") + + gen = Generator( + config, templates=self.templates, resources=self.resources, + version=version + ) + + path = self.templates.joinpath( + f"config-to-{platform}.jsonnet" + ) + wrapper = path.read_text() + + processed = gen.process(wrapper) + + return processed + + def generate(self, config, version, platform): + + logger.info(f"Generating for platform={platform} version={version}") + + try: + + if platform in set(["docker-compose", "podman-compose"]): + return self.generate_docker_compose( + "docker-compose", version, config + ) + elif platform in set(["minikube-k8s", "gcp-k8s"]): + return self.generate_k8s( + platform, version, config + ) + else: + raise RuntimeError("Bad configuration") + + except Exception as e: + logging.error(f"Exception: {e}") + raise e + + def generate_docker_compose(self, platform, version, config): + + processed = self.process( + config, platform=platform, version=version + ) + + y = yaml.dump(processed) + + mem = BytesIO() + + with zipfile.ZipFile(mem, mode='w') as out: + + def output(name, content): + logger.info(f"Adding {name}...") + out.writestr(name, content) + + fname = "docker-compose.yaml" + + output(fname, y) + + # Grafana config + path = self.resources.joinpath( + "grafana/dashboards/dashboard.json" + ) + res = path.read_text() + output("grafana/dashboards/dashboard.json", res) + + path = self.resources.joinpath( + "grafana/provisioning/dashboard.yml" + ) + res = path.read_text() + output("grafana/provisioning/dashboard.yml", res) + + path = self.resources.joinpath( + "grafana/provisioning/datasource.yml" + ) + res = path.read_text() + output("grafana/provisioning/datasource.yml", res) + + # Prometheus config + path = self.resources.joinpath( + "prometheus/prometheus.yml" + ) + res = path.read_text() + output("prometheus/prometheus.yml", res) + + logger.info("Generation complete.") + + return mem.getvalue() + + def generate_k8s(self, platform, version, config): + + processed = self.process( + config, platform=platform, version=version + ) + + y = yaml.dump(processed) + + mem = BytesIO() + + with zipfile.ZipFile(mem, mode='w') as out: + + def output(name, content): + logger.info(f"Adding {name}...") + out.writestr(name, content) + + fname = "resources.yaml" + + output(fname, y) + + logger.info("Generation complete.") + + return mem.getvalue() + def main(): - if len(sys.argv) != 3: + if len(sys.argv) != 4: print() print("Usage:") - print(" generate < input.json") + print(" generate < input.json") print() sys.exit(1) outfile = sys.argv[1] version = sys.argv[2] + platform = sys.argv[3] cfg = sys.stdin.read() - cfg = json.loads(cfg) logger.info(f"Outputting to {outfile}...") - with zipfile.ZipFile(outfile, mode='w') as out: + p = Packager() + resp = p.generate(cfg, version, platform) - def output(name, content): - logger.info(f"Adding {name}...") - out.writestr(name, content) + with open(outfile, "wb") as f: + f.write(resp) - fname = "docker-compose.yaml" - - platform = "docker-compose" - - with open(f"./templates/config-to-{platform}.jsonnet", "r") as f: - wrapper = f.read() - - gen = Generator(json.dumps(cfg).encode("utf-8"), version=version) - - processed = gen.process(wrapper) - - y = yaml.dump(processed) - - output(fname, y) - - # Placeholder for the private.json file. Won't put actual credentials - # here. - output("docker-compose/vertexai/private.json", private_json) - - # Grafana config - with open("grafana/dashboards/dashboard.json") as f: - output( - "docker-compose/grafana/dashboards/dashboard.json", f.read() - ) - - with open("grafana/provisioning/dashboard.yml") as f: - output( - "docker-compose/grafana/provisioning/dashboard.yml", f.read() - ) - - with open("grafana/provisioning/datasource.yml") as f: - output( - "docker-compose/grafana/provisioning/datasource.yml", f.read() - ) - - # Prometheus config - with open("prometheus/prometheus.yml") as f: - output("docker-compose/prometheus/prometheus.yml", f.read()) + return main() From 2c5289765ddefc710d031503c9570e2c14c63e31 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 29 Jan 2025 12:41:01 +0000 Subject: [PATCH 45/53] Updated prometheus targets for current flows (#289) --- prometheus/prometheus.yml | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/prometheus/prometheus.yml b/prometheus/prometheus.yml index 785036c9..0fa70314 100644 --- a/prometheus/prometheus.yml +++ b/prometheus/prometheus.yml @@ -44,11 +44,17 @@ scrape_configs: - targets: - 'chunker:8000' - - job_name: 'vectorize' + - job_name: 'document-embeddings' scrape_interval: 5s static_configs: - targets: - - 'vectorize:8000' + - 'document-embeddings:8000' + + - job_name: 'graph-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'graph-embeddings:8000' - job_name: 'embeddings' scrape_interval: 5s @@ -86,6 +92,12 @@ scrape_configs: - targets: - 'metering-rag:8000' + - job_name: 'store-doc-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'store-doc-embeddings:8000' + - job_name: 'store-graph-embeddings' scrape_interval: 5s static_configs: @@ -116,6 +128,12 @@ scrape_configs: - targets: - 'graph-rag:8000' + - job_name: 'document-rag' + scrape_interval: 5s + static_configs: + - targets: + - 'document-rag:8000' + - job_name: 'prompt' scrape_interval: 5s static_configs: @@ -134,6 +152,12 @@ scrape_configs: - targets: - 'query-graph-embeddings:8000' + - job_name: 'query-doc-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'query-doc-embeddings:8000' + - job_name: 'query-triples' scrape_interval: 5s static_configs: @@ -157,3 +181,7 @@ scrape_configs: static_configs: - targets: - 'workbench-ui:8000' + +# Cassandra +# qdrant + From 1543a0650d7e12dc320fd7fa1f313ecf6919288d Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 29 Jan 2025 14:38:16 +0000 Subject: [PATCH 46/53] Better aws integration (#291) * - More AWS Boto3 settings (profile and session key) - Align environment variable and profile setting names with AWS conventions. Hopefully this should be able to run from an EC2 instance just with role setting. * Tweak naming to all make sense, added rate limit detect --- .../model/text_completion/bedrock/llm.py | 79 ++++++++++++------- 1 file changed, 49 insertions(+), 30 deletions(-) diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index 3b2aced5..f87cf917 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -5,6 +5,7 @@ Input is prompt, output is response. Mistral is default. """ import boto3 +from botocore.errorfactory import ThrottlingException import json from prometheus_client import Histogram import os @@ -24,32 +25,48 @@ default_subscriber = module default_model = 'mistral.mistral-large-2407-v1:0' default_temperature = 0.0 default_max_output = 2048 -default_aws_id_key = os.getenv("AWS_ID_KEY", None) -default_aws_secret = os.getenv("AWS_SECRET", None) -default_aws_region = os.getenv("AWS_REGION", 'us-west-2') + +# Actually, these could all just be None, no need to get environment +# variables, as Boto3 would pick all these up if not passed in as args +default_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", None) +default_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY", None) +default_session_token = os.getenv("AWS_SESSION_TOKEN", None) +default_profile = os.getenv("AWS_PROFILE", None) +default_region = os.getenv("AWS_DEFAULT_REGION", None) class Processor(ConsumerProducer): def __init__(self, **params): + + print(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) + model = params.get("model", default_model) - aws_id_key = params.get("aws_id_key", default_aws_id_key) - aws_secret = params.get("aws_secret", default_aws_secret) - aws_region = params.get("aws_region", default_aws_region) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) - if aws_id_key is None: - raise RuntimeError("AWS ID not specified") + aws_access_key_id = params.get( + "aws_access_key_id", default_access_key_id + ) - if aws_secret is None: - raise RuntimeError("AWS secret not specified") + aws_secret_access_key = params.get( + "aws_secret_access_key", default_secret_access_key + ) - if aws_region is None: - raise RuntimeError("AWS region not specified") + aws_session_token = params.get( + "aws_session_token", default_session_token + ) + + aws_region = params.get( + "aws_region", default_region + ) + + aws_profile = params.get( + "aws_profile", default_profile + ) super(Processor, self).__init__( **params | { @@ -82,9 +99,11 @@ class Processor(ConsumerProducer): self.max_output = max_output self.session = boto3.Session( - aws_access_key_id=aws_id_key, - aws_secret_access_key=aws_secret, - region_name=aws_region + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + profile_name=aws_profile, + region_name=aws_region, ) self.bedrock = self.session.client(service_name='bedrock-runtime') @@ -179,9 +198,6 @@ class Processor(ConsumerProducer): accept = 'application/json' contentType = 'application/json' - # FIXME: Consider catching request limits and raise TooManyRequests - # See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html - with __class__.text_completion_metric.time(): response = self.bedrock.invoke_model( body=promptbody, modelId=self.model, accept=accept, @@ -243,10 +259,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - - # FIXME: Wrong exception, don't know what Bedrock throws - # for a rate limit - except TooManyRequests: + except ThrottlingException: print("Send rate limit response...", flush=True) @@ -300,21 +313,27 @@ class Processor(ConsumerProducer): ) parser.add_argument( - '-z', '--aws-id-key', - default=default_aws_id_key, - help=f'AWS ID Key' + '-z', '--aws-access-key-id', + default=default_access_key_id, + help=f'AWS access key ID' ) parser.add_argument( - '-k', '--aws-secret', - default=default_aws_secret, - help=f'AWS Secret Key' + '-k', '--aws-secret-access-key', + default=default_secret_access_key, + help=f'AWS secret access key' ) parser.add_argument( '-r', '--aws-region', - default=default_aws_region, - help=f'AWS Region' + default=default_region, + help=f'AWS region' + ) + + parser.add_argument( + '--aws-profile', '--profile', + default=default_profile, + help=f'AWS profile name' ) parser.add_argument( From 701ec1e27e2214cdf438fddf16b956539b37f44a Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 29 Jan 2025 19:11:08 +0000 Subject: [PATCH 47/53] Fix startup error on import (#292) --- .../trustgraph/model/text_completion/bedrock/llm.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index f87cf917..21c81fec 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -5,7 +5,6 @@ Input is prompt, output is response. Mistral is default. """ import boto3 -from botocore.errorfactory import ThrottlingException import json from prometheus_client import Histogram import os @@ -259,7 +258,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - except ThrottlingException: + except self.bedrock.exceptions.ThrottlingException: print("Send rate limit response...", flush=True) From d1e9577e7feb60a4c47f6970523734d4da3342d6 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Wed, 29 Jan 2025 21:13:17 +0000 Subject: [PATCH 48/53] Fix rate limit handler, incomplete (#293) --- .../model/text_completion/bedrock/llm.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index 21c81fec..f46e35fb 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -258,27 +258,17 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - except self.bedrock.exceptions.ThrottlingException: + except self.bedrock.exceptions.ThrottlingException as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(type(e)) print(f"Exception: {e}") From 15c8afa97b14d95a8920d7d489c8b88f37d89d0b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 30 Jan 2025 21:40:44 +0000 Subject: [PATCH 49/53] Changed setup.py for 0.21 release branch --- trustgraph-bedrock/setup.py | 2 +- trustgraph-cli/setup.py | 2 +- trustgraph-embeddings-hf/setup.py | 4 ++-- trustgraph-flow/setup.py | 2 +- trustgraph-vertexai/setup.py | 2 +- trustgraph/setup.py | 12 ++++++------ 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/trustgraph-bedrock/setup.py b/trustgraph-bedrock/setup.py index 58330aeb..8db4520b 100644 --- a/trustgraph-bedrock/setup.py +++ b/trustgraph-bedrock/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", "pulsar-client", "prometheus-client", "boto3", diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 4ad9fa51..7d3f513d 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", "requests", "pulsar-client", "rdflib", diff --git a/trustgraph-embeddings-hf/setup.py b/trustgraph-embeddings-hf/setup.py index baad468c..8cf5beb4 100644 --- a/trustgraph-embeddings-hf/setup.py +++ b/trustgraph-embeddings-hf/setup.py @@ -34,8 +34,8 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", - "trustgraph-flow>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", + "trustgraph-flow>=0.21,<0.22", "torch", "urllib3", "transformers", diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index a6191cf5..e56aa694 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", "aiohttp", "anthropic", "cassandra-driver", diff --git a/trustgraph-vertexai/setup.py b/trustgraph-vertexai/setup.py index e02160a6..1258fea9 100644 --- a/trustgraph-vertexai/setup.py +++ b/trustgraph-vertexai/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", "pulsar-client", "google-cloud-aiplatform", "prometheus-client", diff --git a/trustgraph/setup.py b/trustgraph/setup.py index 3abcf1d7..d7185e66 100644 --- a/trustgraph/setup.py +++ b/trustgraph/setup.py @@ -34,12 +34,12 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.20,<0.21", - "trustgraph-bedrock>=0.20,<0.21", - "trustgraph-cli>=0.20,<0.21", - "trustgraph-embeddings-hf>=0.20,<0.21", - "trustgraph-flow>=0.20,<0.21", - "trustgraph-vertexai>=0.20,<0.21", + "trustgraph-base>=0.21,<0.22", + "trustgraph-bedrock>=0.21,<0.22", + "trustgraph-cli>=0.21,<0.22", + "trustgraph-embeddings-hf>=0.21,<0.22", + "trustgraph-flow>=0.21,<0.22", + "trustgraph-vertexai>=0.21,<0.22", ], scripts=[ ] From 5029ab7ad90ae5f34877edb54f705c13c668a637 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Fri, 31 Jan 2025 12:40:23 +0000 Subject: [PATCH 50/53] Fix AWS env vars in templates (#296) --- docs/README.quickstart-docker-compose.md | 5 +++-- templates/components/bedrock-rag.jsonnet | 6 +++--- templates/components/bedrock.jsonnet | 6 +++--- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/docs/README.quickstart-docker-compose.md b/docs/README.quickstart-docker-compose.md index 12cb8cf0..76f7e1f5 100644 --- a/docs/README.quickstart-docker-compose.md +++ b/docs/README.quickstart-docker-compose.md @@ -113,8 +113,9 @@ Choose one of the `Docker Compose` files that meets your preferred model and gra ### AWS Bedrock API ``` -export AWS_ID_KEY= -export AWS_SECRET_KEY= +export AWS_ACCESS_KEY_ID= +export AWS_SECRET_ACCESS_KEY= +export AWS_DEFAULT_REGION= docker compose -f tg-launch-bedrock-cassandra.yaml up -d # Using Cassandra as the graph store docker compose -f tg-launch-bedrock-neo4j.yaml up -d # Using Neo4j as the graph store ``` diff --git a/templates/components/bedrock-rag.jsonnet b/templates/components/bedrock-rag.jsonnet index bc94c8ee..b265a9f2 100644 --- a/templates/components/bedrock-rag.jsonnet +++ b/templates/components/bedrock-rag.jsonnet @@ -20,9 +20,9 @@ local chunker = import "chunker-recursive.jsonnet"; create:: function(engine) local envSecrets = engine.envSecrets("bedrock-credentials") - .with_env_var("AWS_ID_KEY", "aws-id-key") - .with_env_var("AWS_SECRET", "aws-secret") - .with_env_var("AWS_REGION", "aws-region"); + .with_env_var("AWS_ACCESS_KEY_ID", "aws-id-key") + .with_env_var("AWS_SECRET_ACCESS_KEY", "aws-secret") + .with_env_var("AWS_DEFAULT_REGION", "aws-region"); local containerRag = engine.container("text-completion-rag") diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet index bb70d55f..6b599057 100644 --- a/templates/components/bedrock.jsonnet +++ b/templates/components/bedrock.jsonnet @@ -20,9 +20,9 @@ local chunker = import "chunker-recursive.jsonnet"; create:: function(engine) local envSecrets = engine.envSecrets("bedrock-credentials") - .with_env_var("AWS_ID_KEY", "aws-id-key") - .with_env_var("AWS_SECRET", "aws-secret") - .with_env_var("AWS_REGION", "aws-region"); + .with_env_var("AWS_ACCESS_KEY_ID", "aws-id-key") + .with_env_var("AWS_SECRET_ACCESS_KEY", "aws-secret") + .with_env_var("AWS_DEFAULT_REGION", "aws-region"); local container = engine.container("text-completion") From 1000a9de8f2fc8b07d0635d29f5d6f0128ea4f47 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Tue, 4 Feb 2025 12:24:46 +0000 Subject: [PATCH 51/53] Add missing aiohttp dep to trustgraph-cli, needed for knowledge cores (#297) * Add missing aiohttp dep to trustgraph-cli, needed for knowledge cores --- trustgraph-cli/setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 7d3f513d..822ab765 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -37,6 +37,7 @@ setuptools.setup( "trustgraph-base>=0.21,<0.22", "requests", "pulsar-client", + "aiohttp", "rdflib", "tabulate", "msgpack", From 41ccb6c976d571cdc4ea865fe1a908f84be6ede0 Mon Sep 17 00:00:00 2001 From: Tyler Oliver <4535788+toliver38@users.noreply.github.com> Date: Sat, 8 Feb 2025 11:39:23 +0000 Subject: [PATCH 52/53] Add user and password auth for Cassandra (#301) --- .../trustgraph/direct/cassandra.py | 9 +++-- .../query/triples/cassandra/service.py | 34 +++++++++++++++--- .../storage/rows/cassandra/write.py | 24 +++++++++++-- .../storage/triples/cassandra/write.py | 35 ++++++++++++++++--- 4 files changed, 89 insertions(+), 13 deletions(-) diff --git a/trustgraph-flow/trustgraph/direct/cassandra.py b/trustgraph-flow/trustgraph/direct/cassandra.py index 568411a9..d06b270f 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra.py +++ b/trustgraph-flow/trustgraph/direct/cassandra.py @@ -6,7 +6,7 @@ class TrustGraph: def __init__( self, hosts=None, - keyspace="trustgraph", table="default", + keyspace="trustgraph", table="default", username=None, password=None ): if hosts is None: @@ -14,8 +14,13 @@ class TrustGraph: self.keyspace = keyspace self.table = table + self.username = username - self.cluster = Cluster(hosts) + if username and password: + auth_provider = PlainTextAuthProvider(username=username, password=password) + self.cluster = Cluster(hosts, auth_provider=auth_provider) + else: + self.cluster = Cluster(hosts) self.session = self.cluster.connect() self.init() diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 4245784d..22fbf84d 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -26,6 +26,8 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -35,10 +37,14 @@ class Processor(ConsumerProducer): "input_schema": TriplesQueryRequest, "output_schema": TriplesQueryResponse, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) self.graph_host = [graph_host] + self.username = graph_username + self.password = graph_password self.table = None def create_value(self, ent): @@ -56,10 +62,17 @@ class Processor(ConsumerProducer): table = (v.user, v.collection) if table != self.table: - self.tg = TrustGraph( - hosts=self.graph_host, - keyspace=v.user, table=v.collection, - ) + if self.username and self.password: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.user, table=v.collection, + username=self.username, password=self.password + ) + else: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.user, table=v.collection, + ) self.table = table # Sender-produced ID @@ -176,6 +189,19 @@ class Processor(ConsumerProducer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) + def run(): diff --git a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py index d44864fe..fc8f6686 100755 --- a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py @@ -29,6 +29,8 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -36,10 +38,16 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": Rows, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) - - self.cluster = Cluster(graph_host.split(",")) + + if graph_username and graph_password: + auth_provider = PlainTextAuthProvider(username=graph_username, password=graph_password) + self.cluster = Cluster(graph_host.split(","), auth_provider=auth_provider) + else: + self.cluster = Cluster(graph_host.split(",")) self.session = self.cluster.connect() self.tables = set() @@ -120,6 +128,18 @@ class Processor(Consumer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index e7078e08..d940d0ec 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -28,6 +28,8 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -35,10 +37,14 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": Triples, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) - + self.graph_host = [graph_host] + self.username = graph_username + self.password = graph_password self.table = None def handle(self, msg): @@ -52,10 +58,17 @@ class Processor(Consumer): self.tg = None try: - self.tg = TrustGraph( - hosts=self.graph_host, - keyspace=v.metadata.user, table=v.metadata.collection, - ) + if self.username and self.password: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.metadata.user, table=v.metadata.collection, + username=self.username, password=self.password + ) + else: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.metadata.user, table=v.metadata.collection, + ) except Exception as e: print("Exception", e, flush=True) time.sleep(1) @@ -82,6 +95,18 @@ class Processor(Consumer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) def run(): From e99c0ac238f3a92e0c6d1298557aed3b9dff8766 Mon Sep 17 00:00:00 2001 From: Tyler Oliver <4535788+toliver38@users.noreply.github.com> Date: Sat, 8 Feb 2025 11:45:52 +0000 Subject: [PATCH 53/53] Add support for Qdrant API Auth (#300) Added the necessary changes to support API Key in Qdrant Client Query+Storage - Doc Embeddings - Graph Embeddings --- .../query/doc_embeddings/qdrant/service.py | 13 +++++++++++-- .../query/graph_embeddings/qdrant/service.py | 12 ++++++++++-- .../storage/doc_embeddings/qdrant/write.py | 10 +++++++++- .../storage/graph_embeddings/qdrant/write.py | 10 +++++++++- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index 7bb5133a..dd53862f 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -30,6 +30,8 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + #optional api key + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -39,10 +41,11 @@ class Processor(ConsumerProducer): "input_schema": DocumentEmbeddingsRequest, "output_schema": DocumentEmbeddingsResponse, "store_uri": store_uri, + "api_key": api_key, } ) - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) def handle(self, msg): @@ -111,7 +114,13 @@ class Processor(ConsumerProducer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Milvus store URI (default: {default_store_uri})' + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index c2dcaa4c..4e9492c6 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -30,6 +30,7 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -39,10 +40,11 @@ class Processor(ConsumerProducer): "input_schema": GraphEmbeddingsRequest, "output_schema": GraphEmbeddingsResponse, "store_uri": store_uri, + "api_key": api_key, } ) - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -137,7 +139,13 @@ class Processor(ConsumerProducer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Milvus store URI (default: {default_store_uri})' + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index f852e03b..a01fd9e4 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -26,6 +26,7 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -33,6 +34,7 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": DocumentEmbeddings, "store_uri": store_uri, + "api_key": api_key, } ) @@ -97,7 +99,13 @@ class Processor(Consumer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Qdrant store URI (default: {default_store_uri})' + help=f'Qdrant URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 7bc5778c..ffe33565 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -26,6 +26,7 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -33,12 +34,13 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": GraphEmbeddings, "store_uri": store_uri, + "api_key": api_key, } ) self.last_collection = None - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) def get_collection(self, dim, user, collection): @@ -106,6 +108,12 @@ class Processor(Consumer): default=default_store_uri, help=f'Qdrant store URI (default: {default_store_uri})' ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key' + ) def run():