trustgraph/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
cybermaggedon 2b9232917c
Fix/extraction prov (#662)
Quoted triple fixes, including...

1. Updated triple_provenance_triples() in triples.py:
   - Now accepts a Triple object directly
   - Creates the reification triple using TRIPLE term type: stmt_uri tg:reifies
         <<extracted_triple>>
   - Includes it in the returned provenance triples
    
2. Updated definitions extractor:
   - Added imports for provenance functions and component version
   - Added ParameterSpec for optional llm-model and ontology flow parameters
   - For each definition triple, generates provenance with reification
    
3. Updated relationships extractor:
   - Same changes as definitions extractor
2026-03-06 12:23:58 +00:00

101 lines
3.1 KiB
Python
Executable file

"""
Accepts entity/vector pairs and writes them to a Milvus store.
"""
import logging
from .... direct.milvus_graph_embeddings import EntityVectors
from .... base import GraphEmbeddingsStoreService, CollectionConfigHandler
from .... base import AsyncProcessor, Consumer, Producer
from .... base import ConsumerMetrics, ProducerMetrics
from .... schema import IRI, LITERAL
# Module logger
logger = logging.getLogger(__name__)
def get_term_value(term):
"""Extract the string value from a Term"""
if term is None:
return None
if term.type == IRI:
return term.iri
elif term.type == LITERAL:
return term.value
else:
# For blank nodes or other types, use id or value
return term.id or term.value
default_ident = "graph-embeddings-write"
default_store_uri = 'http://localhost:19530'
class Processor(CollectionConfigHandler, GraphEmbeddingsStoreService):
def __init__(self, **params):
store_uri = params.get("store_uri", default_store_uri)
super(Processor, self).__init__(
**params | {
"store_uri": store_uri,
}
)
self.vecstore = EntityVectors(store_uri)
# Register for config push notifications
self.register_config_handler(self.on_collection_config)
async def store_graph_embeddings(self, message):
for entity in message.entities:
entity_value = get_term_value(entity.entity)
if entity_value != "" and entity_value is not None:
for vec in entity.vectors:
self.vecstore.insert(
vec, entity_value,
message.metadata.user,
message.metadata.collection,
chunk_id=entity.chunk_id or "",
)
@staticmethod
def add_args(parser):
GraphEmbeddingsStoreService.add_args(parser)
parser.add_argument(
'-t', '--store-uri',
default=default_store_uri,
help=f'Milvus store URI (default: {default_store_uri})'
)
async def create_collection(self, user: str, collection: str, metadata: dict):
"""
Create collection via config push - collections are created lazily on first write
with the correct dimension determined from the actual embeddings.
"""
try:
logger.info(f"Collection create request for {user}/{collection} - will be created lazily on first write")
self.vecstore.create_collection(user, collection)
except Exception as e:
logger.error(f"Failed to create collection {user}/{collection}: {e}", exc_info=True)
raise
async def delete_collection(self, user: str, collection: str):
"""Delete the collection for graph embeddings via config push"""
try:
self.vecstore.delete_collection(user, collection)
logger.info(f"Successfully deleted collection {user}/{collection}")
except Exception as e:
logger.error(f"Failed to delete collection {user}/{collection}: {e}", exc_info=True)
raise
def run():
Processor.launch(default_ident, __doc__)