trustgraph-flow 1.1.7__tar.gz → 1.1.8__tar.gz

{trustgraph-flow-1.1.7 → trustgraph-flow-1.1.8}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 1.1.7
+Version: 1.1.8
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.1.7.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.1.8.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

{trustgraph-flow-1.1.7 → trustgraph-flow-1.1.8}/trustgraph/direct/cassandra.py

@@ -3,6 +3,9 @@ from cassandra.cluster import Cluster
 from cassandra.auth import PlainTextAuthProvider
 from ssl import SSLContext, PROTOCOL_TLSv1_2
 
+# Global list to track clusters for cleanup
+_active_clusters = []
+
 class TrustGraph:
 
     def __init__(
@@ -24,6 +27,9 @@ class TrustGraph:
         else:
             self.cluster = Cluster(hosts)
         self.session = self.cluster.connect()
+        
+        # Track this cluster globally
+        _active_clusters.append(self.cluster)
 
         self.init()
 
@@ -119,3 +125,13 @@ class TrustGraph:
             f"""select s as x from {self.table} where s = %s and p = %s and o = %s limit {limit}""",
             (s, p, o)
         )
+
+    def close(self):
+        """Close the Cassandra session and cluster connections properly"""
+        if hasattr(self, 'session') and self.session:
+            self.session.shutdown()
+        if hasattr(self, 'cluster') and self.cluster:
+            self.cluster.shutdown()
+            # Remove from global tracking
+            if self.cluster in _active_clusters:
+                _active_clusters.remove(self.cluster)

trustgraph-flow-1.1.8/trustgraph/embeddings/ollama/processor.py

@@ -0,0 +1,63 @@
+
+"""
+Embeddings service, applies an embeddings model hosted on a local Ollama.
+Input is text, output is embeddings vector.
+"""
+from ... base import EmbeddingsService
+
+from ollama import Client
+import os
+
+default_ident = "embeddings"
+
+default_model="mxbai-embed-large"
+default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434')
+
+class Processor(EmbeddingsService):
+
+    def __init__(self, **params):
+
+        model = params.get("model", default_model)
+        ollama = params.get("ollama", default_ollama)
+
+        super(Processor, self).__init__(
+            **params | {
+                "ollama": ollama,
+                "model": model
+            }
+        )
+
+        self.client = Client(host=ollama)
+        self.model = model
+
+    async def on_embeddings(self, text):
+
+        embeds = self.client.embed(
+            model = self.model,
+            input = text
+        )
+
+        return embeds.embeddings
+
+    @staticmethod
+    def add_args(parser):
+
+        EmbeddingsService.add_args(parser)
+
+        parser.add_argument(
+            '-m', '--model',
+            default=default_model,
+            help=f'Embeddings model (default: {default_model})'
+        )
+
+        parser.add_argument(
+            '-r', '--ollama',
+            default=default_ollama,
+            help=f'ollama (default: {default_ollama})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)
+
+

trustgraph-flow-1.1.8/trustgraph/flow_version.py

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

trustgraph-flow-1.1.8/trustgraph/query/doc_embeddings/milvus/service.py

@@ -0,0 +1,68 @@
+
+"""
+Document embeddings query service.  Input is vector, output is an array
+of chunks
+"""
+
+from .... direct.milvus_doc_embeddings import DocVectors
+from .... schema import DocumentEmbeddingsResponse
+from .... schema import Error, Value
+from .... base import DocumentEmbeddingsQueryService
+
+default_ident = "de-query"
+default_store_uri = 'http://localhost:19530'
+
+class Processor(DocumentEmbeddingsQueryService):
+
+    def __init__(self, **params):
+
+        store_uri = params.get("store_uri", default_store_uri)
+
+        super(Processor, self).__init__(
+            **params | {
+                "store_uri": store_uri,
+            }
+        )
+
+        self.vecstore = DocVectors(store_uri)
+
+    async def query_document_embeddings(self, msg):
+
+        try:
+
+            # Handle zero limit case
+            if msg.limit <= 0:
+                return []
+
+            chunks = []
+
+            for vec in msg.vectors:
+
+                resp = self.vecstore.search(vec, limit=msg.limit)
+
+                for r in resp:
+                    chunk = r["entity"]["doc"]
+                    chunks.append(chunk)
+
+            return chunks
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+            raise e
+
+    @staticmethod
+    def add_args(parser):
+
+        DocumentEmbeddingsQueryService.add_args(parser)
+
+        parser.add_argument(
+            '-t', '--store-uri',
+            default=default_store_uri,
+            help=f'Milvus store URI (default: {default_store_uri})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)
+

trustgraph-flow-1.1.8/trustgraph/query/doc_embeddings/pinecone/service.py

@@ -0,0 +1,103 @@
+
+"""
+Document embeddings query service.  Input is vector, output is an array
+of chunks.  Pinecone implementation.
+"""
+
+from pinecone import Pinecone, ServerlessSpec
+from pinecone.grpc import PineconeGRPC, GRPCClientConfig
+
+import uuid
+import os
+
+from .... base import DocumentEmbeddingsQueryService
+
+default_ident = "de-query"
+default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
+
+class Processor(DocumentEmbeddingsQueryService):
+
+    def __init__(self, **params):
+
+        self.url = params.get("url", None)
+        self.api_key = params.get("api_key", default_api_key)
+
+        if self.api_key is None or self.api_key == "not-specified":
+            raise RuntimeError("Pinecone API key must be specified")
+
+        if self.url:
+
+            self.pinecone = PineconeGRPC(
+                api_key = self.api_key,
+                host = self.url
+            )
+
+        else:
+
+            self.pinecone = Pinecone(api_key = self.api_key)
+
+        super(Processor, self).__init__(
+            **params | {
+                "url": self.url,
+                "api_key": self.api_key,
+            }
+        )
+
+    async def query_document_embeddings(self, msg):
+
+        try:
+
+            # Handle zero limit case
+            if msg.limit <= 0:
+                return []
+
+            chunks = []
+
+            for vec in msg.vectors:
+
+                dim = len(vec)
+
+                index_name = (
+                    "d-" + msg.user + "-" + msg.collection + "-" + str(dim)
+                )
+
+                index = self.pinecone.Index(index_name)
+
+                results = index.query(
+                    vector=vec,
+                    top_k=msg.limit,
+                    include_values=False,
+                    include_metadata=True
+                )
+
+                for r in results.matches:
+                    doc = r.metadata["doc"]
+                    chunks.append(doc)
+
+            return chunks
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+            raise e
+
+    @staticmethod
+    def add_args(parser):
+
+        DocumentEmbeddingsQueryService.add_args(parser)
+
+        parser.add_argument(
+            '-a', '--api-key',
+            default=default_api_key,
+            help='Pinecone API key. (default from PINECONE_API_KEY)'
+        )
+
+        parser.add_argument(
+            '-u', '--url',
+            help='Pinecone URL.  If unspecified, serverless is used'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)
+

trustgraph-flow-1.1.8/trustgraph/query/graph_embeddings/milvus/service.py

@@ -0,0 +1,95 @@
+
+"""
+Graph embeddings query service.  Input is vector, output is list of
+entities
+"""
+
+from .... direct.milvus_graph_embeddings import EntityVectors
+from .... schema import GraphEmbeddingsResponse
+from .... schema import Error, Value
+from .... base import GraphEmbeddingsQueryService
+
+default_ident = "ge-query"
+default_store_uri = 'http://localhost:19530'
+
+class Processor(GraphEmbeddingsQueryService):
+
+    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)
+
+    def create_value(self, ent):
+        if ent.startswith("http://") or ent.startswith("https://"):
+            return Value(value=ent, is_uri=True)
+        else:
+            return Value(value=ent, is_uri=False)
+        
+    async def query_graph_embeddings(self, msg):
+
+        try:
+
+            entity_set = set()
+            entities = []
+
+            # Handle zero limit case
+            if msg.limit <= 0:
+                return []
+
+            for vec in msg.vectors:
+
+                resp = self.vecstore.search(vec, limit=msg.limit * 2)
+
+                for r in resp:
+                    ent = r["entity"]["entity"]
+                    
+                    # De-dupe entities
+                    if ent not in entity_set:
+                        entity_set.add(ent)
+                        entities.append(ent)
+
+                    # Keep adding entities until limit
+                    if len(entity_set) >= msg.limit: break
+
+                # Keep adding entities until limit
+                if len(entity_set) >= msg.limit: break
+
+            ents2 = []
+
+            for ent in entities:
+                ents2.append(self.create_value(ent))
+
+            entities = ents2
+
+            print("Send response...", flush=True)
+            return entities
+
+            print("Done.", flush=True)
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+            raise e
+
+    @staticmethod
+    def add_args(parser):
+
+        GraphEmbeddingsQueryService.add_args(parser)
+
+        parser.add_argument(
+            '-t', '--store-uri',
+            default=default_store_uri,
+            help=f'Milvus store URI (default: {default_store_uri})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)
+

{trustgraph-flow-1.1.7 → trustgraph-flow-1.1.8}/trustgraph/query/graph_embeddings/pinecone/service.py

@@ -10,30 +10,23 @@ from pinecone.grpc import PineconeGRPC, GRPCClientConfig
 import uuid
 import os
 
-from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
+from .... schema import GraphEmbeddingsResponse
 from .... schema import Error, Value
-from .... schema import graph_embeddings_request_queue
-from .... schema import graph_embeddings_response_queue
-from .... base import ConsumerProducer
+from .... base import GraphEmbeddingsQueryService
 
-module = "ge-query"
-
-default_input_queue = graph_embeddings_request_queue
-default_output_queue = graph_embeddings_response_queue
-default_subscriber = module
+default_ident = "ge-query"
 default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
 
-class Processor(ConsumerProducer):
+class Processor(GraphEmbeddingsQueryService):
 
     def __init__(self, **params):
 
-        input_queue = params.get("input_queue", default_input_queue)
-        output_queue = params.get("output_queue", default_output_queue)
-        subscriber = params.get("subscriber", default_subscriber)
-
         self.url = params.get("url", None)
         self.api_key = params.get("api_key", default_api_key)
 
+        if self.api_key is None or self.api_key == "not-specified":
+            raise RuntimeError("Pinecone API key must be specified")
+
         if self.url:
 
             self.pinecone = PineconeGRPC(
@@ -47,12 +40,8 @@ class Processor(ConsumerProducer):
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": GraphEmbeddingsRequest,
-                "output_schema": GraphEmbeddingsResponse,
                 "url": self.url,
+                "api_key": self.api_key,
             }
         )
 
@@ -62,26 +51,23 @@ class Processor(ConsumerProducer):
         else:
             return Value(value=ent, is_uri=False)
         
-    async def handle(self, msg):
+    async def query_graph_embeddings(self, msg):
 
         try:
 
-            v = msg.value()
-
-            # Sender-produced ID
-            id = msg.properties()["id"]
-
-            print(f"Handling input {id}...", flush=True)
+            # Handle zero limit case
+            if msg.limit <= 0:
+                return []
 
             entity_set = set()
             entities = []
 
-            for vec in v.vectors:
+            for vec in msg.vectors:
 
                 dim = len(vec)
 
                 index_name = (
-                    "t-" + v.user + "-" + str(dim)
+                    "t-" + msg.user + "-" + msg.collection + "-" + str(dim)
                 )
 
                 index = self.pinecone.Index(index_name)
@@ -89,9 +75,8 @@ class Processor(ConsumerProducer):
                 # Heuristic hack, get (2*limit), so that we have more chance
                 # of getting (limit) entities
                 results = index.query(
-                    namespace=v.collection,
                     vector=vec,
-                    top_k=v.limit * 2,
+                    top_k=msg.limit * 2,
                     include_values=False,
                     include_metadata=True
                 )
@@ -106,10 +91,10 @@ class Processor(ConsumerProducer):
                         entities.append(ent)
 
                     # Keep adding entities until limit
-                    if len(entity_set) >= v.limit: break
+                    if len(entity_set) >= msg.limit: break
 
                 # Keep adding entities until limit
-                if len(entity_set) >= v.limit: break
+                if len(entity_set) >= msg.limit: break
 
             ents2 = []
 
@@ -118,37 +103,17 @@ class Processor(ConsumerProducer):
 
             entities = ents2
 
-            print("Send response...", flush=True)
-            r = GraphEmbeddingsResponse(entities=entities, error=None)
-            await self.send(r, properties={"id": id})
-
-            print("Done.", flush=True)
+            return entities
 
         except Exception as e:
 
             print(f"Exception: {e}")
-
-            print("Send error response...", flush=True)
-
-            r = GraphEmbeddingsResponse(
-                error=Error(
-                    type = "llm-error",
-                    message = str(e),
-                ),
-                entities=None,
-            )
-
-            await self.send(r, properties={"id": id})
-
-            self.consumer.acknowledge(msg)
+            raise e
 
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
+        GraphEmbeddingsQueryService.add_args(parser)
 
         parser.add_argument(
             '-a', '--api-key',
@@ -163,5 +128,5 @@ class Processor(ConsumerProducer):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)