trustgraph-flow 0.22.9__tar.gz → 0.22.11__tar.gz

{trustgraph-flow-0.22.9 → trustgraph-flow-0.22.11}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 0.22.9
+Version: 0.22.11
 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/v0.22.9.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.22.11.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

{trustgraph-flow-0.22.9 → trustgraph-flow-0.22.11}/trustgraph/agent/react/agent_manager.py

@@ -8,11 +8,12 @@ logger = logging.getLogger(__name__)
 
 class AgentManager:
 
-    def __init__(self, tools, additional_context=None):
+    def __init__(self, context, tools, additional_context=None):
+        self.context = context
         self.tools = tools
         self.additional_context = additional_context
 
-    async def reason(self, question, history, context):
+    def reason(self, question, history):
 
         tools = self.tools
 
@@ -55,7 +56,10 @@ class AgentManager:
 
         logger.info(f"prompt: {variables}")
 
-        obj = await context("prompt-request").agent_react(variables)
+        obj = self.context.prompt.request(
+            "agent-react",
+            variables
+        )
 
         print(json.dumps(obj, indent=4), flush=True)
 
@@ -81,13 +85,9 @@ class AgentManager:
 
             return a
 
-    async def react(self, question, history, think, observe, context):
+    async def react(self, question, history, think, observe):
 
-        act = await self.reason(
-            question = question,
-            history = history,
-            context = context,
-        )
+        act = self.reason(question, history)
         logger.info(f"act: {act}")
 
         if isinstance(act, Final):
@@ -104,12 +104,7 @@ class AgentManager:
             else:
                 raise RuntimeError(f"No action for {act.name}!")
 
-            print("TOOL>>>", act)
-            resp = await action.implementation(context).invoke(
-                **act.arguments
-            )
-
-            print("RSETUL", resp)
+            resp = action.implementation.invoke(**act.arguments)
 
             resp = resp.strip()
 

{trustgraph-flow-0.22.9 → trustgraph-flow-0.22.11}/trustgraph/agent/react/service.py

@@ -6,68 +6,103 @@ import json
 import re
 import sys
 
-from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
-from ... base import GraphRagClientSpec
-
-from ... schema import AgentRequest, AgentResponse, AgentStep, Error
+from pulsar.schema import JsonSchema
+
+from ... base import ConsumerProducer
+from ... schema import Error
+from ... schema import AgentRequest, AgentResponse, AgentStep
+from ... schema import agent_request_queue, agent_response_queue
+from ... schema import prompt_request_queue as pr_request_queue
+from ... schema import prompt_response_queue as pr_response_queue
+from ... schema import graph_rag_request_queue as gr_request_queue
+from ... schema import graph_rag_response_queue as gr_response_queue
+from ... clients.prompt_client import PromptClient
+from ... clients.llm_client import LlmClient
+from ... clients.graph_rag_client import GraphRagClient
 
 from . tools import KnowledgeQueryImpl, TextCompletionImpl
 from . agent_manager import AgentManager
 
 from . types import Final, Action, Tool, Argument
 
-default_ident = "agent-manager"
-default_max_iterations = 10
+module = ".".join(__name__.split(".")[1:-1])
 
-class Processor(AgentService):
+default_input_queue = agent_request_queue
+default_output_queue = agent_response_queue
+default_subscriber = module
+default_max_iterations = 15
 
-    def __init__(self, **params):
+class Processor(ConsumerProducer):
 
-        id = params.get("id")
+    def __init__(self, **params):
 
         self.max_iterations = int(
             params.get("max_iterations", default_max_iterations)
         )
 
+        tools = {}
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        prompt_request_queue = params.get(
+            "prompt_request_queue", pr_request_queue
+        )
+        prompt_response_queue = params.get(
+            "prompt_response_queue", pr_response_queue
+        )
+        graph_rag_request_queue = params.get(
+            "graph_rag_request_queue", gr_request_queue
+        )
+        graph_rag_response_queue = params.get(
+            "graph_rag_response_queue", gr_response_queue
+        )
+
         self.config_key = params.get("config_type", "agent")
 
         super(Processor, self).__init__(
             **params | {
-                "id": id,
-                "max_iterations": self.max_iterations,
-                "config_type": self.config_key,
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": AgentRequest,
+                "output_schema": AgentResponse,
+                "prompt_request_queue": prompt_request_queue,
+                "prompt_response_queue": prompt_response_queue,
+                "graph_rag_request_queue": gr_request_queue,
+                "graph_rag_response_queue": gr_response_queue,
             }
         )
 
-        self.agent = AgentManager(
-            tools=[],
-            additional_context="",
+        self.prompt = PromptClient(
+            subscriber=subscriber,
+            input_queue=prompt_request_queue,
+            output_queue=prompt_response_queue,
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
-        self.config_handlers.append(self.on_tools_config)
-
-        self.register_specification(
-            TextCompletionClientSpec(
-                request_name = "text-completion-request",
-                response_name = "text-completion-response",
-            )
+        self.graph_rag = GraphRagClient(
+            subscriber=subscriber,
+            input_queue=graph_rag_request_queue,
+            output_queue=graph_rag_response_queue,
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
-        self.register_specification(
-            GraphRagClientSpec(
-                request_name = "graph-rag-request",
-                response_name = "graph-rag-response",
-            )
+        # Need to be able to feed requests to myself
+        self.recursive_input = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(AgentRequest),
         )
 
-        self.register_specification(
-            PromptClientSpec(
-                request_name = "prompt-request",
-                response_name = "prompt-response",
-            )
+        self.agent = AgentManager(
+            context=self,
+            tools=[],
+            additional_context="",
         )
 
-    async def on_tools_config(self, config, version):
+    async def on_config(self, version, config):
 
         print("Loading configuration version", version)
 
@@ -103,9 +138,9 @@ class Processor(AgentService):
                 impl_id = data.get("type")
 
                 if impl_id == "knowledge-query":
-                    impl = KnowledgeQueryImpl
+                    impl = KnowledgeQueryImpl(self)
                 elif impl_id == "text-completion":
-                    impl = TextCompletionImpl
+                    impl = TextCompletionImpl(self)
                 else:
                     raise RuntimeError(
                         f"Tool-kind {impl_id} not known"
@@ -120,6 +155,7 @@ class Processor(AgentService):
                 )
 
             self.agent = AgentManager(
+                context=self,
                 tools=tools,
                 additional_context=additional
             )
@@ -128,14 +164,19 @@ class Processor(AgentService):
 
         except Exception as e:
 
-            print("on_tools_config Exception:", e, flush=True)
+            print("Exception:", e, flush=True)
             print("Configuration reload failed", flush=True)
 
-    async def agent_request(self, request, respond, next, flow):
+    async def handle(self, msg):
 
         try:
 
-            if request.history:
+            v = msg.value()
+
+            # Sender-produced ID
+            id = msg.properties()["id"]
+
+            if v.history:
                 history = [
                     Action(
                         thought=h.thought,
@@ -143,12 +184,12 @@ class Processor(AgentService):
                         arguments=h.arguments,
                         observation=h.observation
                     )
-                    for h in request.history
+                    for h in v.history
                 ]
             else:
                 history = []
 
-            print(f"Question: {request.question}", flush=True)
+            print(f"Question: {v.question}", flush=True)
 
             if len(history) >= self.max_iterations:
                 raise RuntimeError("Too many agent iterations")
@@ -166,7 +207,7 @@ class Processor(AgentService):
                     observation=None,
                 )
 
-                await respond(r)
+                await self.send(r, properties={"id": id})
 
             async def observe(x):
 
@@ -179,21 +220,15 @@ class Processor(AgentService):
                     observation=x,
                 )
 
-                await respond(r)
+                await self.send(r, properties={"id": id})
 
-            act = await self.agent.react(
-                question = request.question,
-                history = history,
-                think = think,
-                observe = observe,
-                context = flow,
-            )
+            act = await self.agent.react(v.question, history, think, observe)
 
             print(f"Action: {act}", flush=True)
 
-            if isinstance(act, Final):
+            print("Send response...", flush=True)
 
-                print("Send final response...", flush=True)
+            if type(act) == Final:
 
                 r = AgentResponse(
                     answer=act.final,
@@ -201,20 +236,18 @@ class Processor(AgentService):
                     thought=None,
                 )
 
-                await respond(r)
+                await self.send(r, properties={"id": id})
 
                 print("Done.", flush=True)
 
                 return
 
-            print("Send next...", flush=True)
-
             history.append(act)
 
             r = AgentRequest(
-                question=request.question,
-                plan=request.plan,
-                state=request.state,
+                question=v.question,
+                plan=v.plan,
+                state=v.state,
                 history=[
                     AgentStep(
                         thought=h.thought,
@@ -226,7 +259,7 @@ class Processor(AgentService):
                 ]
             )
 
-            await next(r)
+            self.recursive_input.send(r, properties={"id": id})
 
             print("Done.", flush=True)
 
@@ -234,7 +267,7 @@ class Processor(AgentService):
 
         except Exception as e:
 
-            print(f"agent_request Exception: {e}")
+            print(f"Exception: {e}")
 
             print("Send error response...", flush=True)
 
@@ -246,12 +279,39 @@ class Processor(AgentService):
                 response=None,
             )
 
-            await respond(r)
+            await self.send(r, properties={"id": id})
 
     @staticmethod
     def add_args(parser):
 
-        AgentService.add_args(parser)
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '--prompt-request-queue',
+            default=pr_request_queue,
+            help=f'Prompt request queue (default: {pr_request_queue})',
+        )
+
+        parser.add_argument(
+            '--prompt-response-queue',
+            default=pr_response_queue,
+            help=f'Prompt response queue (default: {pr_response_queue})',
+        )
+
+        parser.add_argument(
+            '--graph-rag-request-queue',
+            default=gr_request_queue,
+            help=f'Graph RAG request queue (default: {gr_request_queue})',
+        )
+
+        parser.add_argument(
+            '--graph-rag-response-queue',
+            default=gr_response_queue,
+            help=f'Graph RAG response queue (default: {gr_response_queue})',
+        )
 
         parser.add_argument(
             '--max-iterations',
@@ -267,5 +327,5 @@ class Processor(AgentService):
 
 def run():
 
-    Processor.launch(default_ident, __doc__)
+    Processor.launch(module, __doc__)
 

trustgraph-flow-0.22.11/trustgraph/agent/react/tools.py

@@ -0,0 +1,19 @@
+
+# This tool implementation knows how to put a question to the graph RAG
+# service
+class KnowledgeQueryImpl:
+    def __init__(self, context):
+        self.context = context
+    def invoke(self, **arguments):
+        return self.context.graph_rag.request(arguments.get("question"))
+
+# This tool implementation knows how to do text completion.  This uses
+# the prompt service, rather than talking  to TextCompletion directly.
+class TextCompletionImpl:
+    def __init__(self, context):
+        self.context = context
+    def invoke(self, **arguments):
+        return self.context.prompt.request(
+            "question", { "question": arguments.get("question") }
+        )
+

{trustgraph-flow-0.22.9 → trustgraph-flow-0.22.11}/trustgraph/chunking/recursive/chunker.py

@@ -7,27 +7,40 @@ as text as separate output objects.
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from prometheus_client import Histogram
 
-from ... schema import TextDocument, Chunk
-from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
+from ... schema import TextDocument, Chunk, Metadata
+from ... schema import text_ingest_queue, chunk_ingest_queue
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
 
-default_ident = "chunker"
+module = ".".join(__name__.split(".")[1:-1])
 
-class Processor(FlowProcessor):
+default_input_queue = text_ingest_queue
+default_output_queue = chunk_ingest_queue
+default_subscriber = module
+
+class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
-        id = params.get("id", default_ident)
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
         chunk_size = params.get("chunk_size", 2000)
         chunk_overlap = params.get("chunk_overlap", 100)
         
         super(Processor, self).__init__(
-            **params | { "id": id }
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": TextDocument,
+                "output_schema": Chunk,
+            }
         )
 
         if not hasattr(__class__, "chunk_metric"):
             __class__.chunk_metric = Histogram(
                 'chunk_size', 'Chunk size',
-                ["id", "flow"],
                 buckets=[100, 160, 250, 400, 650, 1000, 1600,
                          2500, 4000, 6400, 10000, 16000]
             )
@@ -39,24 +52,7 @@ class Processor(FlowProcessor):
             is_separator_regex=False,
         )
 
-        self.register_specification(
-            ConsumerSpec(
-                name = "input",
-                schema = TextDocument,
-                handler = self.on_message,
-            )
-        )
-
-        self.register_specification(
-            ProducerSpec(
-                name = "output",
-                schema = Chunk,
-            )
-        )
-
-        print("Chunker initialised", flush=True)
-
-    async def on_message(self, msg, consumer, flow):
+    async def handle(self, msg):
 
         v = msg.value()
         print(f"Chunking {v.metadata.id}...", flush=True)
@@ -67,25 +63,24 @@ class Processor(FlowProcessor):
 
         for ix, chunk in enumerate(texts):
 
-            print("Chunk", len(chunk.page_content), flush=True)
-
             r = Chunk(
                 metadata=v.metadata,
                 chunk=chunk.page_content.encode("utf-8"),
             )
 
-            __class__.chunk_metric.labels(
-                id=consumer.id, flow=consumer.flow
-            ).observe(len(chunk.page_content))
+            __class__.chunk_metric.observe(len(chunk.page_content))
 
-            await flow("output").send(r)
+            await self.send(r)
 
         print("Done.", flush=True)
 
     @staticmethod
     def add_args(parser):
 
-        FlowProcessor.add_args(parser)
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
 
         parser.add_argument(
             '-z', '--chunk-size',
@@ -103,5 +98,5 @@ class Processor(FlowProcessor):
 
 def run():
 
-    Processor.launch(default_ident, __doc__)
+    Processor.launch(module, __doc__)
 

{trustgraph-flow-0.22.9 → trustgraph-flow-0.22.11}/trustgraph/chunking/token/chunker.py

@@ -7,27 +7,40 @@ as text as separate output objects.
 from langchain_text_splitters import TokenTextSplitter
 from prometheus_client import Histogram
 
-from ... schema import TextDocument, Chunk
-from ... base import FlowProcessor
+from ... schema import TextDocument, Chunk, Metadata
+from ... schema import text_ingest_queue, chunk_ingest_queue
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
 
-default_ident = "chunker"
+module = ".".join(__name__.split(".")[1:-1])
 
-class Processor(FlowProcessor):
+default_input_queue = text_ingest_queue
+default_output_queue = chunk_ingest_queue
+default_subscriber = module
+
+class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
-        id = params.get("id")
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
         chunk_size = params.get("chunk_size", 250)
         chunk_overlap = params.get("chunk_overlap", 15)
         
         super(Processor, self).__init__(
-            **params | { "id": id }
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "subscriber": subscriber,
+                "input_schema": TextDocument,
+                "output_schema": Chunk,
+            }
         )
 
         if not hasattr(__class__, "chunk_metric"):
             __class__.chunk_metric = Histogram(
                 'chunk_size', 'Chunk size',
-                ["id", "flow"],
                 buckets=[100, 160, 250, 400, 650, 1000, 1600,
                          2500, 4000, 6400, 10000, 16000]
             )
@@ -38,24 +51,7 @@ class Processor(FlowProcessor):
             chunk_overlap=chunk_overlap,
         )
 
-        self.register_specification(
-            ConsumerSpec(
-                name = "input",
-                schema = TextDocument,
-                handler = self.on_message,
-            )
-        )
-
-        self.register_specification(
-            ProducerSpec(
-                name = "output",
-                schema = Chunk,
-            )
-        )
-
-        print("Chunker initialised", flush=True)
-
-    async def on_message(self, msg, consumer, flow):
+    async def handle(self, msg):
 
         v = msg.value()
         print(f"Chunking {v.metadata.id}...", flush=True)
@@ -66,25 +62,24 @@ class Processor(FlowProcessor):
 
         for ix, chunk in enumerate(texts):
 
-            print("Chunk", len(chunk.page_content), flush=True)
-
             r = Chunk(
                 metadata=v.metadata,
                 chunk=chunk.page_content.encode("utf-8"),
             )
 
-            __class__.chunk_metric.labels(
-                id=consumer.id, flow=consumer.flow
-            ).observe(len(chunk.page_content))
+            __class__.chunk_metric.observe(len(chunk.page_content))
 
-            await flow("output").send(r)
+            await self.send(r)
 
         print("Done.", flush=True)
 
     @staticmethod
     def add_args(parser):
 
-        FlowProcessor.add_args(parser)
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
 
         parser.add_argument(
             '-z', '--chunk-size',
@@ -102,5 +97,5 @@ class Processor(FlowProcessor):
 
 def run():
 
-    Processor.launch(default_ident, __doc__)
+    Processor.launch(module, __doc__)