Agent infrastructure with client and stub server

tests/test-agent

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+import json
+from trustgraph.clients.agent_client import AgentClient
+
+p = AgentClient(pulsar_host="pulsar://localhost:6650")
+
+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."
+
+def think(x):
+    print(f"think: {x}")
+
+def observe(x):
+    print(f"observe: {x}")
+
+resp = p.request(
+    question=q, think=think, observe=observe,
+)
+
+print(resp)
+

tests/test-llm

@@ -5,9 +5,10 @@ from trustgraph.clients.llm_client import LlmClient
 
 llm = LlmClient(pulsar_host="pulsar://localhost:6650")
 
+system = "You are a lovely assistant."
 prompt="Write a funny limerick about a llama"
 
-resp = llm.request(prompt)
+resp = llm.request(system, prompt)
 
 print(resp)
 

trustgraph-base/trustgraph/clients/agent_client.py

@@ -0,0 +1,66 @@
+
+import _pulsar
+
+from .. schema import AgentRequest, AgentResponse
+from .. schema import agent_request_queue
+from .. schema import agent_response_queue
+from . base import BaseClient
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class AgentClient(BaseClient):
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue is None: input_queue = agent_request_queue
+        if output_queue is None: output_queue = agent_response_queue
+
+        super(AgentClient, self).__init__(
+            log_level=log_level,
+            subscriber=subscriber,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            pulsar_host=pulsar_host,
+            input_schema=AgentRequest,
+            output_schema=AgentResponse,
+        )
+
+    def request(
+            self,
+            question,
+            think=None,
+            observe=None,
+            timeout=300
+    ):
+
+        def inspect(x):
+
+            print("inspect", x)
+
+            if x.thought and think:
+                think(x.thought)
+                return
+
+            if x.observation and observation:
+                think(x.observation)
+                return
+
+            if x.answer:
+                return True
+
+            return False
+
+        return self.call(
+            question=question, inspect=inspect, timeout=timeout
+        ).answer
+

trustgraph-base/trustgraph/clients/base.py

@@ -59,10 +59,14 @@ class BaseClient:
     def call(self, **args):
 
         timeout = args.get("timeout", DEFAULT_TIMEOUT)
+        inspect = args.get("inspect", lambda x: True)
 
         if "timeout" in args:
             del args["timeout"]
 
+        if "inspect" in args:
+            del args["inspect"]
+
         id = str(uuid.uuid4())
 
         r = self.input_schema(**args)
@@ -103,6 +107,10 @@ class BaseClient:
                             f"{value.error.type}: {value.error.message}"
                         )
 
+                complete = inspect(value)
+
+                if not complete: continue
+
                 resp = msg.value()
                 self.consumer.acknowledge(msg)
                 return resp

trustgraph-base/trustgraph/schema/__init__.py

@@ -8,4 +8,5 @@ from . topic import *
 from . graph import *
 from . retrieval import *
 from . metadata import *
+from . agent import *
 

trustgraph-base/trustgraph/schema/agent.py

@@ -1,8 +1,8 @@
 
-from pulsar.schema import Record, Bytes, String, Boolean, Array, Map, Integer
+from pulsar.schema import Record, String, Array, Map
 
 from . topic import topic
-from . types import Error, RowSchema
+from . types import Error
 
 ############################################################################
 
@@ -11,18 +11,20 @@ from . types import Error, RowSchema
 class AgentStep(Record):
     thought = String()
     action = String()
-    action_input = String()
+    arguments = Map(String())
     observation = String()
 
 class AgentRequest(Record):
     question = String()
     plan = String()
+    state = String()
     history = Array(AgentStep())
 
 class AgentResponse(Record):
     answer = String()
-    error = String()
+    error = Error()
     thought = String()
+    observation = String()
 
 agent_request_queue = topic(
     'agent', kind='non-persistent', namespace='request'

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

@@ -0,0 +1,3 @@
+
+from . service import *
+

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

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+if __name__ == '__main__':
+    run()
+

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

@@ -0,0 +1,174 @@
+"""
+Simple agent infrastructure broadly implements the ReAct flow.
+"""
+
+import json
+import re
+
+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 ... clients.prompt_client import PromptClient
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = agent_request_queue
+default_output_queue = agent_response_queue
+default_subscriber = module
+
+class Processor(ConsumerProducer):
+
+    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)
+        prompt_request_queue = params.get(
+            "prompt_request_queue", pr_request_queue
+        )
+        prompt_response_queue = params.get(
+            "prompt_response_queue", pr_response_queue
+        )
+
+        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,
+            }
+        )
+
+        self.prompt = PromptClient(
+            subscriber=subscriber,
+            input_queue=prompt_request_queue,
+            output_queue=prompt_response_queue,
+            pulsar_host = self.pulsar_host
+        )
+
+        # Need to be able to feed requests to myself
+        self.recursive_input = self.client.create_producer(
+            topic=input_queue,
+            schema=JsonSchema(AgentRequest),
+        )
+
+    def parse_json(self, text):
+        json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
+    
+        if json_match:
+            json_str = json_match.group(1).strip()
+        else:
+            # If no delimiters, assume the entire output is JSON
+            json_str = text.strip()
+
+        return json.loads(json_str)
+
+    def handle(self, msg):
+
+        try:
+
+            v = msg.value()
+
+            # Sender-produced ID
+            id = msg.properties()["id"]
+
+            if v.history:
+                history = v.history
+            else:
+                history = []
+
+            if len(history) == 0:
+
+                print("Send response...", flush=True)
+
+                thought = "Maybe, just maybe..."
+
+                r = AgentResponse(
+                    answer=None,
+                    error=None,
+                    thought=thought,
+                )
+
+                self.producer.send(r, properties={"id": id})
+
+                r = AgentRequest(
+                    question=v.question,
+                    plan=v.plan,
+                    state=v.state,
+                    history=[
+                        AgentStep(
+                            thought=thought,
+                            action="asd",
+                            arguments={
+                                "asd": "def",
+                                "ghi": "alskd",
+                            },
+                            observation="hello world",
+                        )
+                    ]
+                )
+
+                self.recursive_input.send(r, properties={"id": id})
+
+                return
+
+            else:
+
+                r = AgentResponse(
+                    answer="Here is the answer",
+                    error=None,
+                    thought=None,
+                )
+
+                self.producer.send(r, properties={"id": id})
+
+                return
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+
+            print("Send error response...", flush=True)
+
+            r = AgentResponse(
+                error=Error(
+                    type = "agent-error",
+                    message = str(e),
+                ),
+                response=None,
+            )
+
+            self.producer.send(r, properties={"id": id})
+
+    @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})',
+        )
+
+def run():
+
+    Processor.start(module, __doc__)
+