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

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

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 1.1.6
+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.6.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.6 → trustgraph-flow-1.1.8}/trustgraph/agent/react/agent_manager.py

@@ -14,12 +14,19 @@ class AgentManager:
 
     async def reason(self, question, history, context):
 
+        print(f"calling reason: {question}", flush=True)
+
         tools = self.tools
 
+        print(f"in reason", flush=True)
+        print(tools, flush=True)
+
         tool_names = ",".join([
             t for t in self.tools.keys()
         ])
 
+        print("Tool names:", tool_names, flush=True)
+
         variables = {
             "question": question,
             "tools": [
@@ -83,6 +90,9 @@ class AgentManager:
 
     async def react(self, question, history, think, observe, context):
 
+        logger.info(f"question: {question}")
+        print(f"question: {question}", flush=True)
+
         act = await self.reason(
             question = question,
             history = history,
@@ -104,13 +114,12 @@ class AgentManager:
             else:
                 raise RuntimeError(f"No action for {act.name}!")
 
-            print("TOOL>>>", act)
+            print("TOOL>>>", act, flush=True)
+
             resp = await action.implementation(context).invoke(
                 **act.arguments
             )
 
-            print("RSETUL", resp)
-
             resp = resp.strip()
 
             logger.info(f"resp: {resp}")

{trustgraph-flow-1.1.6 → trustgraph-flow-1.1.8}/trustgraph/agent/react/service.py

@@ -5,13 +5,14 @@ Simple agent infrastructure broadly implements the ReAct flow.
 import json
 import re
 import sys
+import functools
 
 from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
-from ... base import GraphRagClientSpec
+from ... base import GraphRagClientSpec, ToolClientSpec
 
 from ... schema import AgentRequest, AgentResponse, AgentStep, Error
 
-from . tools import KnowledgeQueryImpl, TextCompletionImpl
+from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl
 from . agent_manager import AgentManager
 
 from . types import Final, Action, Tool, Argument
@@ -67,6 +68,13 @@ class Processor(AgentService):
             )
         )
 
+        self.register_specification(
+            ToolClientSpec(
+                request_name = "mcp-tool-request",
+                response_name = "mcp-tool-response",
+            )
+        )
+
     async def on_tools_config(self, config, version):
 
         print("Loading configuration version", version)
@@ -102,17 +110,21 @@ class Processor(AgentService):
 
                 impl_id = data.get("type")
 
+                name = data.get("name")
+
                 if impl_id == "knowledge-query":
                     impl = KnowledgeQueryImpl
                 elif impl_id == "text-completion":
                     impl = TextCompletionImpl
+                elif impl_id == "mcp-tool":
+                    impl = functools.partial(McpToolImpl, name=k)
                 else:
                     raise RuntimeError(
                         f"Tool-kind {impl_id} not known"
                     )
 
                 tools[data.get("name")] = Tool(
-                    name = data.get("name"),
+                    name = name,
                     description = data.get("description"),
                     implementation = impl,
                     config=data.get("config", {}),
@@ -181,6 +193,8 @@ class Processor(AgentService):
 
                 await respond(r)
 
+            print("Call React", flush=True)
+
             act = await self.agent.react(
                 question = request.question,
                 history = history,

{trustgraph-flow-1.1.6 → trustgraph-flow-1.1.8}/trustgraph/agent/react/tools.py

@@ -1,4 +1,6 @@
 
+import json
+
 # This tool implementation knows how to put a question to the graph RAG
 # service
 class KnowledgeQueryImpl:
@@ -23,3 +25,29 @@ class TextCompletionImpl:
             arguments.get("question")
         )
 
+# This tool implementation knows how to do MCP tool invocation.  This uses
+# the mcp-tool service.
+class McpToolImpl:
+
+    def __init__(self, context, name):
+        self.context = context
+        self.name = name
+
+    async def invoke(self, **arguments):
+
+        client = self.context("mcp-tool-request")
+
+        print(f"MCP tool invocation: {self.name}...", flush=True)
+        output = await client.invoke(
+            name = self.name,
+            parameters = {},
+        )
+
+        print(output)
+
+        if isinstance(output, str):
+            return output
+        else:
+            return json.dumps(output)
+
+        

{trustgraph-flow-1.1.6 → 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.6 → 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__)