Fix embeddings to chunk linkage

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)),

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,
             }
         ]
 

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

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,
                     }
                 ]
 

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={
+                            payload=payload,
-                                "entity": entity_value,
-                            }
                         )
                     ]
                 )