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="""
-This knowledge graph describes the Space Shuttle disaster.
-Present 20 facts which are present in the knowledge graph."""
+#query="""
+#This knowledge graph describes the Space Shuttle disaster.
+#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,26 +67,27 @@ class FlowProcessor(AsyncProcessor):
             # Get my flow config
             flow_config = json.loads(config["flows-active"][self.id])
 
-            # Get list of flows which should be running and are currently
-            # running
-            wanted_flows = flow_config.keys()
-            current_flows = self.flows.keys()
-
-            # Start all the flows which arent currently running
-            for flow in wanted_flows:
-                if flow not in current_flows:
-                    await self.start_flow(flow, flow_config[flow])
-
-            # Stop all the unwanted flows which are due to be stopped
-            for flow in current_flows:
-                if flow not in wanted_flows:
-                    await self.stop_flow(flow)
-
-            print("Handled config update")
-
         else:
 
-            print("No configuration settings for me!", flush=True)
+            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()
+        current_flows = self.flows.keys()
+
+        # Start all the flows which arent currently running
+        for flow in wanted_flows:
+            if flow not in current_flows:
+                await self.start_flow(flow, flow_config[flow])
+
+        # Stop all the unwanted flows which are due to be stopped
+        for flow in current_flows:
+            if flow not in wanted_flows:
+                await self.stop_flow(flow)
+
+        print("Handled config update")
 
     # Start threads, just call parent
     async def start(self):

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,14 +55,40 @@ class RequestResponse(Subscriber):
                 properties={"id": id}
             )
 
-            resp = await asyncio.wait_for(
-                q.get(),
-                timeout=timeout
-            )
+        except Exception as e:
 
-            print("Got response.", flush=True)
+            print("Exception:", e)
+            raise e
 
-            return resp
+
+        try:
+
+            while True:
+
+                resp = await asyncio.wait_for(
+                    q.get(),
+                    timeout=timeout
+                )
+
+                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):
-        self.context = context
+    def __init__(self, tools, additional_context=None):
         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(
-            "agent-react",
-            variables
-        )
+        obj = await context("prompt-request").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 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 ... schema import AgentRequest, AgentResponse, AgentStep, Error
 
 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
-default_output_queue = agent_response_queue
-default_subscriber = module
-default_max_iterations = 15
-
-class Processor(ConsumerProducer):
+class Processor(AgentService):
 
     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,
-                "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,
+                "id": id,
+                "max_iterations": self.max_iterations,
+                "config_type": self.config_key,
             }
         )
 
-        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()
-
-            # Sender-produced ID
-            id = msg.properties()["id"]
-
-            if v.history:
+            if request.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,
-                plan=v.plan,
-                state=v.state,
+                question=request.question,
+                plan=request.plan,
+                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(
-            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})',
-        )
+        AgentService.add_args(parser)
 
         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):
-        return self.context.graph_rag.request(arguments.get("question"))
+    async def invoke(self, **arguments):
+        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):
-        return self.context.prompt.request(
-            "question", { "question": arguments.get("question") }
+    async def invoke(self, **arguments):
+        client = self.context("prompt-request")
+        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)
 
-