trustgraph-flow 1.3.17__tar.gz → 1.3.19__tar.gz

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.3.17
+Version: 1.3.19
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/pyproject.toml

@@ -90,7 +90,6 @@ metering = "trustgraph.metering:run"
 nlp-query = "trustgraph.retrieval.nlp_query:run"
 objects-write-cassandra = "trustgraph.storage.objects.cassandra:run"
 objects-query-cassandra = "trustgraph.query.objects.cassandra:run"
-oe-write-milvus = "trustgraph.storage.object_embeddings.milvus:run"
 pdf-decoder = "trustgraph.decoding.pdf:run"
 pdf-ocr-mistral = "trustgraph.decoding.mistral_ocr:run"
 prompt-template = "trustgraph.prompt.template:run"

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/direct/milvus_doc_embeddings.py

@@ -2,9 +2,32 @@
 from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType
 import time
 import logging
+import re
 
 logger = logging.getLogger(__name__)
 
+def make_safe_collection_name(user, collection, dimension, prefix):
+    """
+    Create a safe Milvus collection name from user/collection parameters.
+    Milvus only allows letters, numbers, and underscores.
+    """
+    def sanitize(s):
+        # Replace non-alphanumeric characters (except underscore) with underscore
+        # Then collapse multiple underscores into single underscore
+        safe = re.sub(r'[^a-zA-Z0-9_]', '_', s)
+        safe = re.sub(r'_+', '_', safe)
+        # Remove leading/trailing underscores
+        safe = safe.strip('_')
+        # Ensure it's not empty
+        if not safe:
+            safe = 'default'
+        return safe
+    
+    safe_user = sanitize(user)
+    safe_collection = sanitize(collection)
+    
+    return f"{prefix}_{safe_user}_{safe_collection}_{dimension}"
+
 class DocVectors:
 
     def __init__(self, uri="http://localhost:19530", prefix='doc'):
@@ -26,9 +49,9 @@ class DocVectors:
         self.next_reload = time.time() + self.reload_time
         logger.debug(f"Reload at {self.next_reload}")
 
-    def init_collection(self, dimension):
+    def init_collection(self, dimension, user, collection):
 
-        collection_name = self.prefix + "_" + str(dimension)
+        collection_name = make_safe_collection_name(user, collection, dimension, self.prefix)
 
         pkey_field = FieldSchema(
             name="id",
@@ -75,14 +98,14 @@ class DocVectors:
             index_params=index_params
         )
 
-        self.collections[dimension] = collection_name
+        self.collections[(dimension, user, collection)] = collection_name
 
-    def insert(self, embeds, doc):
+    def insert(self, embeds, doc, user, collection):
 
         dim = len(embeds)
 
-        if dim not in self.collections:
-            self.init_collection(dim)
+        if (dim, user, collection) not in self.collections:
+            self.init_collection(dim, user, collection)
     
         data = [
             {
@@ -92,18 +115,18 @@ class DocVectors:
         ]
 
         self.client.insert(
-            collection_name=self.collections[dim],
+            collection_name=self.collections[(dim, user, collection)],
             data=data
         )
 
-    def search(self, embeds, fields=["doc"], limit=10):
+    def search(self, embeds, user, collection, fields=["doc"], limit=10):
 
         dim = len(embeds)
 
-        if dim not in self.collections:
-            self.init_collection(dim)
+        if (dim, user, collection) not in self.collections:
+            self.init_collection(dim, user, collection)
 
-        coll = self.collections[dim]
+        coll = self.collections[(dim, user, collection)]
 
         search_params = {
             "metric_type": "COSINE",

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/direct/milvus_graph_embeddings.py

@@ -2,9 +2,32 @@
 from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType
 import time
 import logging
+import re
 
 logger = logging.getLogger(__name__)
 
+def make_safe_collection_name(user, collection, dimension, prefix):
+    """
+    Create a safe Milvus collection name from user/collection parameters.
+    Milvus only allows letters, numbers, and underscores.
+    """
+    def sanitize(s):
+        # Replace non-alphanumeric characters (except underscore) with underscore
+        # Then collapse multiple underscores into single underscore
+        safe = re.sub(r'[^a-zA-Z0-9_]', '_', s)
+        safe = re.sub(r'_+', '_', safe)
+        # Remove leading/trailing underscores
+        safe = safe.strip('_')
+        # Ensure it's not empty
+        if not safe:
+            safe = 'default'
+        return safe
+    
+    safe_user = sanitize(user)
+    safe_collection = sanitize(collection)
+    
+    return f"{prefix}_{safe_user}_{safe_collection}_{dimension}"
+
 class EntityVectors:
 
     def __init__(self, uri="http://localhost:19530", prefix='entity'):
@@ -26,9 +49,9 @@ class EntityVectors:
         self.next_reload = time.time() + self.reload_time
         logger.debug(f"Reload at {self.next_reload}")
 
-    def init_collection(self, dimension):
+    def init_collection(self, dimension, user, collection):
 
-        collection_name = self.prefix + "_" + str(dimension)
+        collection_name = make_safe_collection_name(user, collection, dimension, self.prefix)
 
         pkey_field = FieldSchema(
             name="id",
@@ -75,14 +98,14 @@ class EntityVectors:
             index_params=index_params
         )
 
-        self.collections[dimension] = collection_name
+        self.collections[(dimension, user, collection)] = collection_name
 
-    def insert(self, embeds, entity):
+    def insert(self, embeds, entity, user, collection):
 
         dim = len(embeds)
 
-        if dim not in self.collections:
-            self.init_collection(dim)
+        if (dim, user, collection) not in self.collections:
+            self.init_collection(dim, user, collection)
     
         data = [
             {
@@ -92,18 +115,18 @@ class EntityVectors:
         ]
 
         self.client.insert(
-            collection_name=self.collections[dim],
+            collection_name=self.collections[(dim, user, collection)],
             data=data
         )
 
-    def search(self, embeds, fields=["entity"], limit=10):
+    def search(self, embeds, user, collection, fields=["entity"], limit=10):
 
         dim = len(embeds)
 
-        if dim not in self.collections:
-            self.init_collection(dim)
+        if (dim, user, collection) not in self.collections:
+            self.init_collection(dim, user, collection)
 
-        coll = self.collections[dim]
+        coll = self.collections[(dim, user, collection)]
 
         search_params = {
             "metric_type": "COSINE",

trustgraph_flow-1.3.19/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "1.3.19"

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/doc_embeddings/milvus/service.py

@@ -43,7 +43,12 @@ class Processor(DocumentEmbeddingsQueryService):
 
             for vec in msg.vectors:
 
-                resp = self.vecstore.search(vec, limit=msg.limit)
+                resp = self.vecstore.search(
+                    vec, 
+                    msg.user, 
+                    msg.collection, 
+                    limit=msg.limit
+                )
 
                 for r in resp:
                     chunk = r["entity"]["doc"]

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/doc_embeddings/pinecone/service.py

@@ -47,6 +47,39 @@ class Processor(DocumentEmbeddingsQueryService):
             }
         )
 
+        self.last_index_name = None
+
+    def ensure_index_exists(self, index_name, dim):
+        """Ensure index exists, create if it doesn't"""
+        if index_name != self.last_index_name:
+            if not self.pinecone.has_index(index_name):
+                try:
+                    self.pinecone.create_index(
+                        name=index_name,
+                        dimension=dim,
+                        metric="cosine",
+                        spec=ServerlessSpec(
+                            cloud="aws",
+                            region="us-east-1",
+                        )
+                    )
+                    logger.info(f"Created index: {index_name}")
+                    
+                    # Wait for index to be ready
+                    import time
+                    for i in range(0, 1000):
+                        if self.pinecone.describe_index(index_name).status["ready"]:
+                            break
+                        time.sleep(1)
+                        
+                    if not self.pinecone.describe_index(index_name).status["ready"]:
+                        raise RuntimeError("Gave up waiting for index creation")
+                        
+                except Exception as e:
+                    logger.error(f"Pinecone index creation failed: {e}")
+                    raise e
+            self.last_index_name = index_name
+
     async def query_document_embeddings(self, msg):
 
         try:
@@ -65,6 +98,8 @@ class Processor(DocumentEmbeddingsQueryService):
                     "d-" + msg.user + "-" + msg.collection + "-" + str(dim)
                 )
 
+                self.ensure_index_exists(index_name, dim)
+
                 index = self.pinecone.Index(index_name)
 
                 results = index.query(

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/doc_embeddings/qdrant/service.py

@@ -38,6 +38,24 @@ class Processor(DocumentEmbeddingsQueryService):
         )
 
         self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
+        self.last_collection = None
+
+    def ensure_collection_exists(self, collection, dim):
+        """Ensure collection exists, create if it doesn't"""
+        if collection != self.last_collection:
+            if not self.qdrant.collection_exists(collection):
+                try:
+                    self.qdrant.create_collection(
+                        collection_name=collection,
+                        vectors_config=VectorParams(
+                            size=dim, distance=Distance.COSINE
+                        ),
+                    )
+                    logger.info(f"Created collection: {collection}")
+                except Exception as e:
+                    logger.error(f"Qdrant collection creation failed: {e}")
+                    raise e
+            self.last_collection = collection
 
     async def query_document_embeddings(self, msg):
 
@@ -53,6 +71,8 @@ class Processor(DocumentEmbeddingsQueryService):
                     str(dim)
                 )
 
+                self.ensure_collection_exists(collection, dim)
+
                 search_result = self.qdrant.query_points(
                     collection_name=collection,
                     query=vec,

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/graph_embeddings/milvus/service.py

@@ -50,7 +50,12 @@ class Processor(GraphEmbeddingsQueryService):
 
             for vec in msg.vectors:
 
-                resp = self.vecstore.search(vec, limit=msg.limit * 2)
+                resp = self.vecstore.search(
+                    vec, 
+                    msg.user, 
+                    msg.collection, 
+                    limit=msg.limit * 2
+                )
 
                 for r in resp:
                     ent = r["entity"]["entity"]

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/graph_embeddings/pinecone/service.py

@@ -49,6 +49,39 @@ class Processor(GraphEmbeddingsQueryService):
             }
         )
 
+        self.last_index_name = None
+
+    def ensure_index_exists(self, index_name, dim):
+        """Ensure index exists, create if it doesn't"""
+        if index_name != self.last_index_name:
+            if not self.pinecone.has_index(index_name):
+                try:
+                    self.pinecone.create_index(
+                        name=index_name,
+                        dimension=dim,
+                        metric="cosine",
+                        spec=ServerlessSpec(
+                            cloud="aws",
+                            region="us-east-1",
+                        )
+                    )
+                    logger.info(f"Created index: {index_name}")
+                    
+                    # Wait for index to be ready
+                    import time
+                    for i in range(0, 1000):
+                        if self.pinecone.describe_index(index_name).status["ready"]:
+                            break
+                        time.sleep(1)
+                        
+                    if not self.pinecone.describe_index(index_name).status["ready"]:
+                        raise RuntimeError("Gave up waiting for index creation")
+                        
+                except Exception as e:
+                    logger.error(f"Pinecone index creation failed: {e}")
+                    raise e
+            self.last_index_name = index_name
+
     def create_value(self, ent):
         if ent.startswith("http://") or ent.startswith("https://"):
             return Value(value=ent, is_uri=True)
@@ -74,6 +107,8 @@ class Processor(GraphEmbeddingsQueryService):
                     "t-" + msg.user + "-" + msg.collection + "-" + str(dim)
                 )
 
+                self.ensure_index_exists(index_name, dim)
+
                 index = self.pinecone.Index(index_name)
 
                 # Heuristic hack, get (2*limit), so that we have more chance

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/query/graph_embeddings/qdrant/service.py

@@ -38,6 +38,24 @@ class Processor(GraphEmbeddingsQueryService):
         )
 
         self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
+        self.last_collection = None
+
+    def ensure_collection_exists(self, collection, dim):
+        """Ensure collection exists, create if it doesn't"""
+        if collection != self.last_collection:
+            if not self.qdrant.collection_exists(collection):
+                try:
+                    self.qdrant.create_collection(
+                        collection_name=collection,
+                        vectors_config=VectorParams(
+                            size=dim, distance=Distance.COSINE
+                        ),
+                    )
+                    logger.info(f"Created collection: {collection}")
+                except Exception as e:
+                    logger.error(f"Qdrant collection creation failed: {e}")
+                    raise e
+            self.last_collection = collection
 
     def create_value(self, ent):
         if ent.startswith("http://") or ent.startswith("https://"):
@@ -60,6 +78,8 @@ class Processor(GraphEmbeddingsQueryService):
                     str(dim)
                 )
 
+                self.ensure_collection_exists(collection, dim)
+
                 # Heuristic hack, get (2*limit), so that we have more chance
                 # of getting (limit) entities
                 search_result = self.qdrant.query_points(

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/storage/doc_embeddings/milvus/write.py

@@ -33,7 +33,11 @@ class Processor(DocumentEmbeddingsStoreService):
             if chunk == "": continue
 
             for vec in emb.vectors:
-                self.vecstore.insert(vec, chunk)
+                self.vecstore.insert(
+                    vec, chunk, 
+                    message.metadata.user, 
+                    message.metadata.collection
+                )
 
     @staticmethod
     def add_args(parser):

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph/storage/graph_embeddings/milvus/write.py

@@ -29,7 +29,11 @@ class Processor(GraphEmbeddingsStoreService):
 
             if entity.entity.value != "" and entity.entity.value is not None:
                 for vec in entity.vectors:
-                    self.vecstore.insert(vec, entity.entity.value)
+                    self.vecstore.insert(
+                        vec, entity.entity.value,
+                        message.metadata.user,
+                        message.metadata.collection
+                    )
 
     @staticmethod
     def add_args(parser):

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph_flow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.3.17
+Version: 1.3.19
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph_flow.egg-info/SOURCES.txt

@@ -40,7 +40,6 @@ trustgraph/direct/__init__.py
 trustgraph/direct/cassandra.py
 trustgraph/direct/milvus_doc_embeddings.py
 trustgraph/direct/milvus_graph_embeddings.py
-trustgraph/direct/milvus_object_embeddings.py
 trustgraph/embeddings/__init__.py
 trustgraph/embeddings/document_embeddings/__init__.py
 trustgraph/embeddings/document_embeddings/__main__.py
@@ -260,9 +259,6 @@ trustgraph/storage/knowledge/__init__.py
 trustgraph/storage/knowledge/__main__.py
 trustgraph/storage/knowledge/store.py
 trustgraph/storage/object_embeddings/__init__.py
-trustgraph/storage/object_embeddings/milvus/__init__.py
-trustgraph/storage/object_embeddings/milvus/__main__.py
-trustgraph/storage/object_embeddings/milvus/write.py
 trustgraph/storage/objects/__init__.py
 trustgraph/storage/objects/cassandra/__init__.py
 trustgraph/storage/objects/cassandra/__main__.py

{trustgraph_flow-1.3.17 → trustgraph_flow-1.3.19}/trustgraph_flow.egg-info/entry_points.txt

@@ -35,7 +35,6 @@ metering = trustgraph.metering:run
 nlp-query = trustgraph.retrieval.nlp_query:run
 objects-query-cassandra = trustgraph.query.objects.cassandra:run
 objects-write-cassandra = trustgraph.storage.objects.cassandra:run
-oe-write-milvus = trustgraph.storage.object_embeddings.milvus:run
 pdf-decoder = trustgraph.decoding.pdf:run
 pdf-ocr-mistral = trustgraph.decoding.mistral_ocr:run
 prompt-template = trustgraph.prompt.template:run

trustgraph_flow-1.3.17/trustgraph/direct/milvus_object_embeddings.py

@@ -1,157 +0,0 @@
-
-from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType
-import time
-import logging
-
-logger = logging.getLogger(__name__)
-
-class ObjectVectors:
-
-    def __init__(self, uri="http://localhost:19530", prefix='obj'):
-
-        self.client = MilvusClient(uri=uri)
-
-        # Strategy is to create collections per dimension.  Probably only
-        # going to be using 1 anyway, but that means we don't need to
-        # hard-code the dimension anywhere, and no big deal if more than
-        # one are created.
-        self.collections = {}
-
-        self.prefix = prefix
-
-        # Time between reloads
-        self.reload_time = 90
-
-        # Next time to reload - this forces a reload at next window
-        self.next_reload = time.time() + self.reload_time
-        logger.debug(f"Reload at {self.next_reload}")
-
-    def init_collection(self, dimension, name):
-
-        collection_name = self.prefix + "_" + name + "_" + str(dimension)
-
-        pkey_field = FieldSchema(
-            name="id",
-            dtype=DataType.INT64,
-            is_primary=True,
-            auto_id=True,
-        )
-
-        vec_field = FieldSchema(
-            name="vector",
-            dtype=DataType.FLOAT_VECTOR,
-            dim=dimension,
-        )
-
-        name_field = FieldSchema(
-            name="name",
-            dtype=DataType.VARCHAR,
-            max_length=65535,
-        )
-
-        key_name_field = FieldSchema(
-            name="key_name",
-            dtype=DataType.VARCHAR,
-            max_length=65535,
-        )
-
-        key_field = FieldSchema(
-            name="key",
-            dtype=DataType.VARCHAR,
-            max_length=65535,
-        )
-
-        schema = CollectionSchema(
-            fields = [
-                pkey_field, vec_field, name_field, key_name_field, key_field
-            ],
-            description = "Object embedding schema",
-        )
-
-        self.client.create_collection(
-            collection_name=collection_name,
-            schema=schema,
-            metric_type="COSINE",
-        )
-
-        index_params = MilvusClient.prepare_index_params()
-
-        index_params.add_index(
-            field_name="vector",
-            metric_type="COSINE",
-            index_type="IVF_SQ8",
-            index_name="vector_index",
-            params={ "nlist": 128 }
-        )
-
-        self.client.create_index(
-            collection_name=collection_name,
-            index_params=index_params
-        )
-
-        self.collections[(dimension, name)] = collection_name
-
-    def insert(self, embeds, name, key_name, key):
-
-        dim = len(embeds)
-
-        if (dim, name) not in self.collections:
-            self.init_collection(dim, name)
-    
-        data = [
-            {
-                "vector": embeds,
-                "name": name,
-                "key_name": key_name,
-                "key": key,
-            }
-        ]
-
-        self.client.insert(
-            collection_name=self.collections[(dim, name)],
-            data=data
-        )
-
-    def search(self, embeds, name, fields=["key_name", "name"], limit=10):
-
-        dim = len(embeds)
-
-        if dim not in self.collections:
-            self.init_collection(dim, name)
-
-        coll = self.collections[(dim, name)]
-
-        search_params = {
-            "metric_type": "COSINE",
-            "params": {
-                "radius": 0.1,
-                "range_filter": 0.8
-            }
-        }
-
-        logger.debug("Loading...")
-        self.client.load_collection(
-            collection_name=coll,
-        )
-
-        logger.debug("Searching...")
-
-        res = self.client.search(
-            collection_name=coll,
-            data=[embeds],
-            limit=limit,
-            output_fields=fields,
-            search_params=search_params,
-        )[0]
-
-
-        # If reload time has passed, unload collection
-        if time.time() > self.next_reload:
-            logger.debug(f"Unloading, reload at {self.next_reload}")
-            self.client.release_collection(
-                collection_name=coll,
-            )
-            self.next_reload = time.time() + self.reload_time
-
-        return res
-

trustgraph_flow-1.3.17/trustgraph/flow_version.py

@@ -1 +0,0 @@
-__version__ = "1.3.17"

trustgraph_flow-1.3.17/trustgraph/storage/object_embeddings/milvus/write.py

@@ -1,61 +0,0 @@
-
-"""
-Accepts entity/vector pairs and writes them to a Milvus store.
-"""
-
-from .... schema import ObjectEmbeddings
-from .... schema import object_embeddings_store_queue
-from .... log_level import LogLevel
-from .... direct.milvus_object_embeddings import ObjectVectors
-from .... base import Consumer
-
-module = "oe-write"
-
-default_input_queue = object_embeddings_store_queue
-default_subscriber = module
-default_store_uri = 'http://localhost:19530'
-
-class Processor(Consumer):
-
-    def __init__(self, **params):
-
-        input_queue = params.get("input_queue", default_input_queue)
-        subscriber = params.get("subscriber", default_subscriber)
-        store_uri = params.get("store_uri", default_store_uri)
-
-        super(Processor, self).__init__(
-            **params | {
-                "input_queue": input_queue,
-                "subscriber": subscriber,
-                "input_schema": ObjectEmbeddings,
-                "store_uri": store_uri,
-            }
-        )
-
-        self.vecstore = ObjectVectors(store_uri)
-
-    async def handle(self, msg):
-
-        v = msg.value()
-
-        if v.id != "" and v.id is not None:
-            for vec in v.vectors:
-                self.vecstore.insert(vec, v.name, v.key_name, v.id)
-
-    @staticmethod
-    def add_args(parser):
-
-        Consumer.add_args(
-            parser, default_input_queue, default_subscriber,
-        )
-
-        parser.add_argument(
-            '-t', '--store-uri',
-            default=default_store_uri,
-            help=f'Milvus store URI (default: {default_store_uri})'
-        )
-
-def run():
-
-    Processor.launch(module, __doc__)
-