trustgraph-flow 1.3.18__tar.gz → 1.3.20__tar.gz

{trustgraph_flow-1.3.18 → trustgraph_flow-1.3.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.3.18
+Version: 1.3.20
 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.20/trustgraph/flow_version.py

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

{trustgraph_flow-1.3.18 → trustgraph_flow-1.3.20}/trustgraph/gateway/dispatch/mux.py

@@ -147,7 +147,7 @@ class Mux:
                 self.running.stop()
 
                 if self.ws:
-                    self.ws.close()
+                    await self.ws.close()
                     self.ws = None
 
                 break
@@ -165,6 +165,6 @@ class Mux:
         self.running.stop()
 
         if self.ws:
-            self.ws.close()
+            await self.ws.close()
             self.ws = None
 

{trustgraph_flow-1.3.18 → trustgraph_flow-1.3.20}/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.18 → trustgraph_flow-1.3.20}/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.18 → trustgraph_flow-1.3.20}/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.18 → trustgraph_flow-1.3.20}/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.18 → trustgraph_flow-1.3.20}/trustgraph_flow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.3.18
+Version: 1.3.20
 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.18/trustgraph/flow_version.py

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