From c3471df6cee956b379d356548d06167dca9902ac Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Thu, 5 Mar 2026 19:15:33 +0000
Subject: [PATCH] Fix embeddings to chunk linkage

---
 trustgraph-base/trustgraph/provenance/triples.py    |  2 ++
 .../trustgraph/direct/milvus_graph_embeddings.py    | 13 ++++++++++---
 .../storage/graph_embeddings/milvus/write.py        |  3 ++-
 .../storage/graph_embeddings/pinecone/write.py      |  6 +++++-
 .../storage/graph_embeddings/qdrant/write.py        | 10 +++++++---
 5 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/trustgraph-base/trustgraph/provenance/triples.py b/trustgraph-base/trustgraph/provenance/triples.py
index 5e4dc774..732f2fb7 100644
--- a/trustgraph-base/trustgraph/provenance/triples.py
+++ b/trustgraph-base/trustgraph/provenance/triples.py
@@ -145,6 +145,7 @@ def derived_entity_triples(
 
         # Activity declaration
         _triple(act_uri, RDF_TYPE, _iri(PROV_ACTIVITY)),
+        _triple(act_uri, RDFS_LABEL, _literal(f"{component_name} extraction")),
         _triple(act_uri, PROV_USED, _iri(parent_uri)),
         _triple(act_uri, PROV_WAS_ASSOCIATED_WITH, _iri(agt_uri)),
         _triple(act_uri, PROV_STARTED_AT_TIME, _literal(timestamp)),
@@ -232,6 +233,7 @@ def triple_provenance_triples(
 
         # Activity
         _triple(act_uri, RDF_TYPE, _iri(PROV_ACTIVITY)),
+        _triple(act_uri, RDFS_LABEL, _literal(f"{component_name} extraction")),
         _triple(act_uri, PROV_USED, _iri(chunk_uri)),
         _triple(act_uri, PROV_WAS_ASSOCIATED_WITH, _iri(agt_uri)),
         _triple(act_uri, PROV_STARTED_AT_TIME, _literal(timestamp)),
diff --git a/trustgraph-flow/trustgraph/direct/milvus_graph_embeddings.py b/trustgraph-flow/trustgraph/direct/milvus_graph_embeddings.py
index 4a106f27..dcbf6734 100644
--- a/trustgraph-flow/trustgraph/direct/milvus_graph_embeddings.py
+++ b/trustgraph-flow/trustgraph/direct/milvus_graph_embeddings.py
@@ -90,8 +90,14 @@ class EntityVectors:
             max_length=65535,
         )
 
+        chunk_id_field = FieldSchema(
+            name="chunk_id",
+            dtype=DataType.VARCHAR,
+            max_length=65535,
+        )
+
         schema = CollectionSchema(
-            fields = [pkey_field, vec_field, entity_field],
+            fields = [pkey_field, vec_field, entity_field, chunk_id_field],
             description = "Graph embedding schema",
         )
 
@@ -119,17 +125,18 @@ class EntityVectors:
         self.collections[(dimension, user, collection)] = collection_name
         logger.info(f"Created Milvus collection {collection_name} with dimension {dimension}")
 
-    def insert(self, embeds, entity, user, collection):
+    def insert(self, embeds, entity, user, collection, chunk_id=""):
 
         dim = len(embeds)
 
         if (dim, user, collection) not in self.collections:
             self.init_collection(dim, user, collection)
-    
+
         data = [
             {
                 "vector": embeds,
                 "entity": entity,
+                "chunk_id": chunk_id,
             }
         ]
 
diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
index 21aa21e6..8e1c4485 100755
--- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
+++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
@@ -57,7 +57,8 @@ class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService):
                     self.vecstore.insert(
                         vec, entity_value,
                         message.metadata.user,
-                        message.metadata.collection
+                        message.metadata.collection,
+                        chunk_id=entity.chunk_id or "",
                     )
 
     @staticmethod
diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py
index c4b0065b..f4de7f82 100755
--- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py
+++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py
@@ -137,11 +137,15 @@ class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService):
                 # Generate unique ID for each vector
                 vector_id = str(uuid.uuid4())
 
+                metadata = {"entity": entity_value}
+                if entity.chunk_id:
+                    metadata["chunk_id"] = entity.chunk_id
+
                 records = [
                     {
                         "id": vector_id,
                         "values": vec,
-                        "metadata": { "entity": entity_value },
+                        "metadata": metadata,
                     }
                 ]
 
diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
index 0da59bb9..4877ae96 100755
--- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
+++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
@@ -90,15 +90,19 @@ class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService):
                         )
                     )
 
+                payload = {
+                    "entity": entity_value,
+                }
+                if entity.chunk_id:
+                    payload["chunk_id"] = entity.chunk_id
+
                 self.qdrant.upsert(
                     collection_name=collection,
                     points=[
                         PointStruct(
                             id=str(uuid.uuid4()),
                             vector=vec,
-                            payload={
-                                "entity": entity_value,
-                            }
+                            payload=payload,
                         )
                     ]
                 )