From e428b2ea6178c0bfe551ef3222de50b8302e1067 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 4 Feb 2026 14:09:08 +0000 Subject: [PATCH] Fix tg-load-knowledge --- trustgraph-cli/pyproject.toml | 3 ++ .../trustgraph/cli/load_knowledge.py | 29 +++++++++++++---- .../dispatch/document_embeddings_query.py | 31 +++++++++++++++++++ .../trustgraph/gateway/dispatch/manager.py | 2 ++ 4 files changed, 59 insertions(+), 6 deletions(-) create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_query.py diff --git a/trustgraph-cli/pyproject.toml b/trustgraph-cli/pyproject.toml index ab7980ac..09a22bdc 100644 --- a/trustgraph-cli/pyproject.toml +++ b/trustgraph-cli/pyproject.toml @@ -43,6 +43,9 @@ tg-invoke-agent = "trustgraph.cli.invoke_agent:main" tg-invoke-document-rag = "trustgraph.cli.invoke_document_rag:main" tg-invoke-graph-rag = "trustgraph.cli.invoke_graph_rag:main" tg-invoke-llm = "trustgraph.cli.invoke_llm:main" +tg-invoke-embeddings = "trustgraph.cli.invoke_embeddings:main" +tg-invoke-graph-embeddings = "trustgraph.cli.invoke_graph_embeddings:main" +tg-invoke-document-embeddings = "trustgraph.cli.invoke_document_embeddings:main" tg-invoke-mcp-tool = "trustgraph.cli.invoke_mcp_tool:main" tg-invoke-nlp-query = "trustgraph.cli.invoke_nlp_query:main" tg-invoke-objects-query = "trustgraph.cli.invoke_objects_query:main" diff --git a/trustgraph-cli/trustgraph/cli/load_knowledge.py b/trustgraph-cli/trustgraph/cli/load_knowledge.py index ff6ca980..5e96850f 100644 --- a/trustgraph-cli/trustgraph/cli/load_knowledge.py +++ b/trustgraph-cli/trustgraph/cli/load_knowledge.py @@ -87,13 +87,20 @@ class KnowledgeLoader: # Load triples from all files print("Loading triples...") + total_triples = 0 for file in self.files: print(f" Processing {file}...") - triples = self.load_triples_from_file(file) + count = 0 + + def counting_triples(): + nonlocal count + for triple in self.load_triples_from_file(file): + count += 1 + yield triple bulk.import_triples( flow=self.flow, - triples=triples, + triples=counting_triples(), metadata={ "id": self.document_id, "metadata": [], @@ -101,25 +108,33 @@ class KnowledgeLoader: "collection": self.collection } ) + print(f" Loaded {count} triples") + total_triples += count - print("Triples loaded.") + print(f"Triples loaded. Total: {total_triples}") # Load entity contexts from all files print("Loading entity contexts...") + total_contexts = 0 for file in self.files: print(f" Processing {file}...") + count = 0 # Convert tuples to the format expected by import_entity_contexts + # Entity must be in Term format: {"t": "i", "i": uri} for IRI def entity_context_generator(): + nonlocal count for entity, context in self.load_entity_contexts_from_file(file): + count += 1 + # Entities from RDF are URIs, use IRI term format yield { - "entity": {"v": entity, "e": True}, + "entity": {"t": "i", "i": entity}, "context": context } bulk.import_entity_contexts( flow=self.flow, - entities=entity_context_generator(), + contexts=entity_context_generator(), metadata={ "id": self.document_id, "metadata": [], @@ -127,8 +142,10 @@ class KnowledgeLoader: "collection": self.collection } ) + print(f" Loaded {count} entity contexts") + total_contexts += count - print("Entity contexts loaded.") + print(f"Entity contexts loaded. Total: {total_contexts}") except Exception as e: print(f"Error: {e}", flush=True) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_query.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_query.py new file mode 100644 index 00000000..650d4f40 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_query.py @@ -0,0 +1,31 @@ + +from ... schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse +from ... messaging import TranslatorRegistry + +from . requestor import ServiceRequestor + +class DocumentEmbeddingsQueryRequestor(ServiceRequestor): + def __init__( + self, backend, request_queue, response_queue, timeout, + consumer, subscriber, + ): + + super(DocumentEmbeddingsQueryRequestor, self).__init__( + backend=backend, + request_queue=request_queue, + response_queue=response_queue, + request_schema=DocumentEmbeddingsRequest, + response_schema=DocumentEmbeddingsResponse, + subscription = subscriber, + consumer_name = consumer, + timeout=timeout, + ) + + self.request_translator = TranslatorRegistry.get_request_translator("document-embeddings-query") + self.response_translator = TranslatorRegistry.get_response_translator("document-embeddings-query") + + def to_request(self, body): + return self.request_translator.to_pulsar(body) + + def from_response(self, message): + return self.response_translator.from_response_with_completion(message) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 0766e232..2d401cf3 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -26,6 +26,7 @@ from . structured_query import StructuredQueryRequestor from . structured_diag import StructuredDiagRequestor from . embeddings import EmbeddingsRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor +from . document_embeddings_query import DocumentEmbeddingsQueryRequestor from . mcp_tool import McpToolRequestor from . text_load import TextLoad from . document_load import DocumentLoad @@ -55,6 +56,7 @@ request_response_dispatchers = { "document-rag": DocumentRagRequestor, "embeddings": EmbeddingsRequestor, "graph-embeddings": GraphEmbeddingsQueryRequestor, + "document-embeddings": DocumentEmbeddingsQueryRequestor, "triples": TriplesQueryRequestor, "objects": ObjectsQueryRequestor, "nlp-query": NLPQueryRequestor,