Fix tg-load-knowledge

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"

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)

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)

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,