Agent flows working

tests/test-agent

@@ -20,7 +20,11 @@ def output(text, prefix="> ", width=78):
     )
     print(out)
 
-p = AgentClient(pulsar_host="pulsar://localhost:6650")
+p = AgentClient(
+    pulsar_host="pulsar://pulsar:6650",
+    input_queue = "non-persistent://tg/request/agent:0000",
+    output_queue = "non-persistent://tg/response/agent:0000",
+)
 
 q = "How many cats does Mark have? Calculate that number raised to 0.4 power.  Is that number lower than the numeric part of the mission identifier of the Space Shuttle Challenger on its last mission?  If so, give me an apple pie recipe, otherwise return a poem about cheese."
 

tests/test-graph-rag

@@ -10,9 +10,11 @@ rag = GraphRagClient(
     output_queue = "non-persistent://tg/response/graph-rag:default",
 )
 
-query="""
+#query="""
-This knowledge graph describes the Space Shuttle disaster.
+#This knowledge graph describes the Space Shuttle disaster.
-Present 20 facts which are present in the knowledge graph."""
+#Present 20 facts which are present in the knowledge graph."""
+
+query = "How many cats does Mark have?"
 
 resp = rag.request(query)
 

trustgraph-base/trustgraph/base/__init__.py

@@ -26,4 +26,6 @@ from . document_embeddings_query_service import DocumentEmbeddingsQueryService
 from . graph_embeddings_client import GraphEmbeddingsClientSpec
 from . triples_client import TriplesClientSpec
 from . document_embeddings_client import DocumentEmbeddingsClientSpec
+from . agent_service import AgentService
+from . graph_rag_client import GraphRagClientSpec
 

trustgraph-base/trustgraph/base/agent_client.py

@@ -0,0 +1,39 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import AgentRequest, AgentResponse
+from .. knowledge import Uri, Literal
+
+class AgentClient(RequestResponse):
+    async def request(self, recipient, question, plan=None, state=None,
+                    history=[], timeout=300):
+
+        resp = await self.request(
+            AgentRequest(
+                question = question,
+                plan = plan,
+                state = state,
+                history = history,
+            ),
+            recipient=recipient,
+            timeout=timeout,
+        )
+
+        print(resp, flush=True)
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp
+
+class GraphEmbeddingsClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(GraphEmbeddingsClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = GraphEmbeddingsRequest,
+            response_name = response_name,
+            response_schema = GraphEmbeddingsResponse,
+            impl = GraphEmbeddingsClient,
+        )
+

trustgraph-base/trustgraph/base/agent_service.py

@@ -0,0 +1,100 @@
+
+"""
+Agent manager service completion base class
+"""
+
+import time
+from prometheus_client import Histogram
+
+from .. schema import AgentRequest, AgentResponse, Error
+from .. exceptions import TooManyRequests
+from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
+
+default_ident = "agent-manager"
+
+class AgentService(FlowProcessor):
+
+    def __init__(self, **params):
+
+        id = params.get("id")
+
+        super(AgentService, self).__init__(**params | { "id": id })
+
+        self.register_specification(
+            ConsumerSpec(
+                name = "request",
+                schema = AgentRequest,
+                handler = self.on_request
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "next",
+                schema = AgentRequest
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "response",
+                schema = AgentResponse
+            )
+        )
+
+    async def on_request(self, msg, consumer, flow):
+
+        try:
+
+            request = msg.value()
+
+            # Sender-produced ID
+            id = msg.properties()["id"]
+
+            async def respond(resp):
+
+                await flow("response").send(
+                    resp,
+                    properties={"id": id}
+                )
+
+            async def next(resp):
+
+                await flow("next").send(
+                    resp,
+                    properties={"id": id}
+                )
+
+            await self.agent_request(
+                request = request, respond = respond, next = next,
+                flow = flow
+            )
+
+        except TooManyRequests as e:
+            raise e
+
+        except Exception as e:
+
+            # Apart from rate limits, treat all exceptions as unrecoverable
+            print(f"on_request Exception: {e}")
+
+            print("Send error response...", flush=True)
+
+            await flow.producer["response"].send(
+                AgentResponse(
+                    error=Error(
+                        type = "agent-error",
+                        message = str(e),
+                    ),
+                    thought = None,
+                    observation = None,
+                    answer = None,
+                ),
+                properties={"id": id}
+            )
+
+    @staticmethod
+    def add_args(parser):
+
+        FlowProcessor.add_args(parser)
+

trustgraph-base/trustgraph/base/consumer.py

@@ -89,7 +89,7 @@ class Consumer:
 
             except Exception as e:
 
-                print("Exception:", e, flush=True)
+                print("consumer subs Exception:", e, flush=True)
                 await asyncio.sleep(self.reconnect_time)
                 continue
 
@@ -107,7 +107,7 @@ class Consumer:
 
             except Exception as e:
 
-                print("Exception:", e, flush=True)
+                print("consumer loop exception:", e, flush=True)
                 self.consumer.close()
                 self.consumer = None
                 await asyncio.sleep(self.reconnect_time)
@@ -184,7 +184,7 @@ class Consumer:
 
                 except Exception as e:
 
-                    print("Exception:", e, flush=True)
+                    print("consume exception:", e, flush=True)
 
                     # Message failed to be processed, this causes it to
                     # be retried

trustgraph-base/trustgraph/base/flow_processor.py

@@ -67,6 +67,11 @@ class FlowProcessor(AsyncProcessor):
             # Get my flow config
             flow_config = json.loads(config["flows-active"][self.id])
 
+        else:
+
+            print("No configuration settings for me.", flush=True)
+            flow_config = {}
+
         # Get list of flows which should be running and are currently
         # running
         wanted_flows = flow_config.keys()
@@ -84,10 +89,6 @@ class FlowProcessor(AsyncProcessor):
 
         print("Handled config update")
 
-        else:
-
-            print("No configuration settings for me!", flush=True)
-
     # Start threads, just call parent
     async def start(self):
         await super(FlowProcessor, self).start()

trustgraph-base/trustgraph/base/graph_rag_client.py

@@ -0,0 +1,33 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import GraphRagQuery, GraphRagResponse
+
+class GraphRagClient(RequestResponse):
+    async def rag(self, query, user="trustgraph", collection="default",
+                  timeout=600):
+        resp = await self.request(
+            GraphRagQuery(
+                query = query,
+                user = user,
+                collection = collection,
+            ),
+            timeout=timeout
+        )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp.response
+
+class GraphRagClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(GraphRagClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = GraphRagQuery,
+            response_name = response_name,
+            response_schema = GraphRagResponse,
+            impl = GraphRagClient,
+        )
+

trustgraph-base/trustgraph/base/llm_service.py

@@ -62,8 +62,6 @@ class LlmService(FlowProcessor):
 
             id = msg.properties()["id"]
 
-            prompt = request.system + "\n\n" + request.prompt
-
             with __class__.text_completion_metric.labels(
                     id=self.id,
                     flow=f"{flow.name}-{consumer.name}",

trustgraph-base/trustgraph/base/prompt_client.py

@@ -63,6 +63,22 @@ class PromptClient(RequestResponse):
             timeout = timeout,
         )
 
+    async def agent_react(self, variables, timeout=600):
+        return await self.prompt(
+            id = "agent-react",
+            variables = variables,
+            timeout = timeout,
+        )
+
+    async def question(self, question, timeout=600):
+        return await self.prompt(
+            id = "question",
+            variables = {
+                "question": question,
+            },
+            timeout = timeout,
+        )
+
 class PromptClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,

trustgraph-base/trustgraph/base/request_response_spec.py

@@ -40,7 +40,7 @@ class RequestResponse(Subscriber):
         await self.producer.stop()
         await super(RequestResponse, self).stop()
 
-    async def request(self, req, timeout=300):
+    async def request(self, req, timeout=300, recipient=None):
 
         id = str(uuid.uuid4())
 
@@ -55,6 +55,16 @@ class RequestResponse(Subscriber):
                 properties={"id": id}
             )
 
+        except Exception as e:
+
+            print("Exception:", e)
+            raise e
+
+
+        try:
+
+            while True:
+
                 resp = await asyncio.wait_for(
                     q.get(),
                     timeout=timeout
@@ -62,7 +72,23 @@ class RequestResponse(Subscriber):
 
                 print("Got response.", flush=True)
 
+                if recipient is None:
+
+                    # If no recipient handler, just return the first
+                    # response we get
                     return resp
+                else:
+
+                    # Recipient handler gets to decide when we're done b
+                    # returning a boolean
+                    fin = await recipient(resp)
+
+                    # If done, return the last result otherwise loop round for
+                    # next response
+                    if fin:
+                        return resp
+                    else:
+                        continue
 
         except Exception as e:
 
@@ -108,4 +134,3 @@ class RequestResponseSpec(Spec):
 
         flow.consumer[self.request_name] = rr
 
-

trustgraph-base/trustgraph/schema/agent.py

@@ -26,12 +26,5 @@ class AgentResponse(Record):
     thought = String()
     observation = String()
 
-agent_request_queue = topic(
-    'agent', kind='non-persistent', namespace='request'
-)
-agent_response_queue = topic(
-    'agent', kind='non-persistent', namespace='response'
-)
-
 ############################################################################
 

trustgraph-flow/trustgraph/agent/react/agent_manager.py

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

trustgraph-flow/trustgraph/agent/react/service.py

@@ -6,105 +6,68 @@ import json
 import re
 import sys
 
-from pulsar.schema import JsonSchema
+from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
+from ... base import GraphRagClientSpec
 
-from ... base import ConsumerProducer
+from ... schema import AgentRequest, AgentResponse, AgentStep, Error
-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
 
-module = "agent"
+default_ident = "agent-manager"
+default_max_iterations = 10
 
-default_input_queue = agent_request_queue
+class Processor(AgentService):
-default_output_queue = agent_response_queue
-default_subscriber = module
-default_max_iterations = 15
-
-class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
+        id = params.get("id")
+
         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 | {
-                "input_queue": input_queue,
+                "id": id,
-                "output_queue": output_queue,
+                "max_iterations": self.max_iterations,
-                "subscriber": subscriber,
+                "config_type": self.config_key,
-                "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.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.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,
-        )
-
-        # Need to be able to feed requests to myself
-        self.recursive_input = self.client.create_producer(
-            topic=input_queue,
-            schema=JsonSchema(AgentRequest),
-        )
-
         self.agent = AgentManager(
-            context=self,
             tools=[],
             additional_context="",
         )
 
-        self.config_handlers.append(self.on_config)
+        self.config_handlers.append(self.on_tools_config)
 
-    async def on_config(self, version, config):
+        self.register_specification(
+            TextCompletionClientSpec(
+                request_name = "text-completion-request",
+                response_name = "text-completion-response",
+            )
+        )
+
+        self.register_specification(
+            GraphRagClientSpec(
+                request_name = "graph-rag-request",
+                response_name = "graph-rag-response",
+            )
+        )
+
+        self.register_specification(
+            PromptClientSpec(
+                request_name = "prompt-request",
+                response_name = "prompt-response",
+            )
+        )
+
+    async def on_tools_config(self, config, version):
 
         print("Loading configuration version", version)
 
@@ -140,9 +103,9 @@ class Processor(ConsumerProducer):
                 impl_id = data.get("type")
 
                 if impl_id == "knowledge-query":
-                    impl = KnowledgeQueryImpl(self)
+                    impl = KnowledgeQueryImpl
                 elif impl_id == "text-completion":
-                    impl = TextCompletionImpl(self)
+                    impl = TextCompletionImpl
                 else:
                     raise RuntimeError(
                         f"Tool-kind {impl_id} not known"
@@ -157,7 +120,6 @@ class Processor(ConsumerProducer):
                 )
 
             self.agent = AgentManager(
-                context=self,
                 tools=tools,
                 additional_context=additional
             )
@@ -166,19 +128,14 @@ class Processor(ConsumerProducer):
 
         except Exception as e:
 
-            print("Exception:", e, flush=True)
+            print("on_tools_config Exception:", e, flush=True)
             print("Configuration reload failed", flush=True)
 
-    async def handle(self, msg):
+    async def agent_request(self, request, respond, next, flow):
 
         try:
 
-            v = msg.value()
+            if request.history:
-
-            # Sender-produced ID
-            id = msg.properties()["id"]
-
-            if v.history:
                 history = [
                     Action(
                         thought=h.thought,
@@ -186,12 +143,12 @@ class Processor(ConsumerProducer):
                         arguments=h.arguments,
                         observation=h.observation
                     )
-                    for h in v.history
+                    for h in request.history
                 ]
             else:
                 history = []
 
-            print(f"Question: {v.question}", flush=True)
+            print(f"Question: {request.question}", flush=True)
 
             if len(history) >= self.max_iterations:
                 raise RuntimeError("Too many agent iterations")
@@ -209,7 +166,7 @@ class Processor(ConsumerProducer):
                     observation=None,
                 )
 
-                await self.send(r, properties={"id": id})
+                await respond(r)
 
             async def observe(x):
 
@@ -222,15 +179,21 @@ class Processor(ConsumerProducer):
                     observation=x,
                 )
 
-                await self.send(r, properties={"id": id})
+                await respond(r)
 
-            act = await self.agent.react(v.question, history, think, observe)
+            act = await self.agent.react(
+                question = request.question,
+                history = history,
+                think = think,
+                observe = observe,
+                context = flow,
+            )
 
             print(f"Action: {act}", flush=True)
 
-            print("Send response...", flush=True)
+            if isinstance(act, Final):
 
-            if type(act) == Final:
+                print("Send final response...", flush=True)
 
                 r = AgentResponse(
                     answer=act.final,
@@ -238,18 +201,20 @@ class Processor(ConsumerProducer):
                     thought=None,
                 )
 
-                await self.send(r, properties={"id": id})
+                await respond(r)
 
                 print("Done.", flush=True)
 
                 return
 
+            print("Send next...", flush=True)
+
             history.append(act)
 
             r = AgentRequest(
-                question=v.question,
+                question=request.question,
-                plan=v.plan,
+                plan=request.plan,
-                state=v.state,
+                state=request.state,
                 history=[
                     AgentStep(
                         thought=h.thought,
@@ -261,7 +226,7 @@ class Processor(ConsumerProducer):
                 ]
             )
 
-            self.recursive_input.send(r, properties={"id": id})
+            await next(r)
 
             print("Done.", flush=True)
 
@@ -269,7 +234,7 @@ class Processor(ConsumerProducer):
 
         except Exception as e:
 
-            print(f"Exception: {e}")
+            print(f"agent_request Exception: {e}")
 
             print("Send error response...", flush=True)
 
@@ -281,39 +246,12 @@ class Processor(ConsumerProducer):
                 response=None,
             )
 
-            await self.send(r, properties={"id": id})
+            await respond(r)
 
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
+        AgentService.add_args(parser)
-            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',
@@ -329,5 +267,5 @@ class Processor(ConsumerProducer):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)
 

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

@@ -4,16 +4,22 @@
 class KnowledgeQueryImpl:
     def __init__(self, context):
         self.context = context
-    def invoke(self, **arguments):
+    async def invoke(self, **arguments):
-        return self.context.graph_rag.request(arguments.get("question"))
+        client = self.context("graph-rag-request")
+        print("Graph RAG question...", flush=True)
+        return await client.rag(
+            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):
+    async def invoke(self, **arguments):
-        return self.context.prompt.request(
+        client = self.context("prompt-request")
-            "question", { "question": arguments.get("question") }
+        print("Prompt question...", flush=True)
+        return await client.question(
+            arguments.get("question")
         )
 

trustgraph-flow/trustgraph/gateway/agent.py

@@ -39,4 +39,3 @@ class AgentRequestor(ServiceRequestor):
         # The 2nd boolean expression indicates whether we're done responding
         return resp, (message.answer is not None)
 
-