Agent working through client

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

@@ -0,0 +1,192 @@
+
+import ibis
+import logging
+import json
+
+from . types import Action, Final
+
+logger = logging.getLogger(__name__)
+
+class AgentManager:
+
+    template="""Answer the following questions as best you can. You have
+access to the following functions:
+
+{% for tool in tools %}{
+    "function": "{{ tool.tool.name }}",
+    "description": "{{ tool.tool.description }}",
+    "arguments": [
+{% for arg in tool.tool.arguments %}        {
+            "name": "{{ arg.name }}",
+            "type": "{{ arg.type }}",
+            "description": "{{ arg.description }}",
+        }
+{% endfor %}
+    ]
+}
+{% endfor %}
+
+You can either choose to call a function to get more information, or
+return a final answer.
+    
+To call a function, respond with a JSON object of the following format:
+
+{
+    "thought": "your thought about what to do",
+    "action": "the action to take, should be one of [{{tool_names}}]",
+    "arguments": {
+        "argument1": "argument_value",
+        "argument2": "argument_value"
+    }
+}
+
+To provide a final answer, response a JSON object of the following format:
+
+{
+  "thought": "I now know the final answer",
+  "final-answer": "the final answer to the original input question"
+}
+
+Previous steps are included in the input.  Each step has the following
+format in your output:
+
+{
+  "thought": "your thought about what to do",
+  "action": "the action taken",
+  "arguments": {
+      "argument1": action argument,
+      "argument2": action argument2
+  },
+  "observation": "the result of the action",
+}
+
+Respond by describing either one single thought/action/arguments or
+the final-answer.  Pause after providing one action or final-answer.
+Begin!
+
+Question: {{question}}
+
+Input:
+    
+{% for h in history %}
+{
+    "action": "{{h.action}}",
+    "arguments": [
+{% for k, v in h.arguments.items() %}        {
+            "{{k}}": "{{v}}",
+{%endfor%}        }
+    ],
+    "observation": "{{h.observation}}"
+}
+{% endfor %}"""
+
+    def __init__(self, context, tools):
+        self.context = context
+        self.tools = tools
+
+    def reason(self, question, history):
+
+        tpl = ibis.Template(self.template)
+
+        tools = self.tools
+
+        tool_names = ",".join([
+            t.tool.name for t in self.tools
+        ])
+
+        print("HERE!")
+
+        prompt = tpl.render({
+            "tools": tools,
+            "question": question,
+            "tool_names": tool_names,
+            "history": [
+                {
+                    "thought": h.thought,
+                    "action": h.name,
+                    "arguments": h.arguments,
+                    "observation": h.observation,
+                }
+                for h in history
+            ],
+        })
+
+        logger.info(f"prompt: {prompt}")
+
+        print(prompt)
+
+        resp = self.context.prompt.request(
+            "question",
+            {
+                "question": prompt
+            }
+        )
+
+        resp = resp.replace("```json", "")
+        resp = resp.replace("```", "")
+
+        logger.info(f"response: {resp}")
+
+        obj = json.loads(resp)
+
+        if obj.get("final-answer"):
+
+            a = Final(
+                thought = obj.get("thought"),
+                final = obj.get("final-answer"),
+            )
+
+            return a
+
+        else:
+
+            a = Action(
+                thought = obj.get("thought"),
+                name = obj.get("action"),
+                arguments = obj.get("arguments"),
+                observation = ""
+            )
+
+            return a
+
+    def react(self, question, history, think, observe):
+
+        act = self.reason(question, history)
+        logger.info(f"act: {act}")
+
+        print(act)
+
+        if isinstance(act, Final):
+
+            think(act.thought)
+            return act
+
+        else:
+
+            think(act.thought)
+
+            action = None
+
+            for tool in self.tools:
+                if tool.tool.name == act.name:
+                    action = tool
+
+            if action is None:
+                raise RuntimeError(f"No action for {act.name}!")
+
+            print(act.arguments)
+            print("Invoke...")
+            resp = action.invoke(**act.arguments)
+
+            resp = resp.strip()
+
+            logger.info(f"resp: {resp}")
+
+            observe(resp)
+
+            act.observation = resp
+
+            logger.info(f"iter: {act}")
+
+            return act
+

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

@@ -13,7 +13,17 @@ 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 text_completion_request_queue as tc_request_queue
+from ... schema import text_completion_response_queue as tc_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 CatsKb, ShuttleKb, Compute
+from . agent_manager import AgentManager
+from . types import Final, Action
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -34,6 +44,18 @@ class Processor(ConsumerProducer):
         prompt_response_queue = params.get(
             "prompt_response_queue", pr_response_queue
         )
+        text_completion_request_queue = params.get(
+            "text_completion_request_queue", tc_request_queue
+        )
+        text_completion_response_queue = params.get(
+            "text_completion_response_queue", tc_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
+        )
 
         super(Processor, self).__init__(
             **params | {
@@ -44,6 +66,10 @@ class Processor(ConsumerProducer):
                 "output_schema": AgentResponse,
                 "prompt_request_queue": prompt_request_queue,
                 "prompt_response_queue": prompt_response_queue,
+                "text_completion_request_queue": tc_request_queue,
+                "text_completion_response_queue": tc_response_queue,
+                "graph_rag_request_queue": gr_request_queue,
+                "graph_rag_response_queue": gr_response_queue,
             }
         )
 
@@ -54,12 +80,32 @@ class Processor(ConsumerProducer):
             pulsar_host = self.pulsar_host
         )
 
+        self.llm = LlmClient(
+            subscriber=subscriber,
+            input_queue=text_completion_request_queue,
+            output_queue=text_completion_response_queue,
+            pulsar_host = self.pulsar_host
+        )
+
+        self.graph_rag = GraphRagClient(
+            subscriber=subscriber,
+            input_queue=graph_rag_request_queue,
+            output_queue=graph_rag_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),
         )
 
+        tools = [
+            CatsKb(self), ShuttleKb(self), Compute(self),
+        ]
+
+        self.agent = AgentManager(self, tools)
+
     def parse_json(self, text):
         json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
     
@@ -81,63 +127,63 @@ class Processor(ConsumerProducer):
             id = msg.properties()["id"]
 
             if v.history:
-                history = v.history
+                history = [
+                    Action(
+                        thought=h.thought,
+                        name=h.action,
+                        arguments=h.arguments,
+                        observation=h.observation
+                    )
+                    for h in v.history
+                ]
             else:
                 history = []
 
-            if len(history) == 0:
+            print(v.question)
 
-                print("Send response...", flush=True)
+            if len(history) > 10:
+                raise RuntimeError("Too many agent iterations")
 
-                thought = "Maybe, just maybe..."
+            print("History:", history)
+
+            def think(x):
+
+                print("THINK:", x)
 
                 r = AgentResponse(
                     answer=None,
                     error=None,
-                    thought=thought,
+                    thought=x,
                     observation=None,
                 )
 
                 self.producer.send(r, properties={"id": id})
 
+            def observe(x):
+
+                print("OBSERVE:", x)
+
                 r = AgentResponse(
                     answer=None,
                     error=None,
                     thought=None,
-                    observation="I have imagined myself",
+                    observation=x,
                 )
 
                 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",
-                        )
-                    ]
-                )
+            act = self.agent.react(v.question, history, think, observe)
 
-                self.recursive_input.send(r, properties={"id": id})
+            print(act)
 
-                print("Done.", flush=True)
+            print("Send response...", flush=True)
 
-                return
-
-            else:
+            if type(act) == Final:
 
                 print("Send response...", flush=True)
 
                 r = AgentResponse(
-                    answer="Here is the answer",
+                    answer=act.final,
                     error=None,
                     thought=None,
                 )
@@ -148,6 +194,31 @@ class Processor(ConsumerProducer):
 
                 return
 
+            history.append(act)
+
+            print(history)
+
+            r = AgentRequest(
+                question=v.question,
+                plan=v.plan,
+                state=v.state,
+                history=[
+                    AgentStep(
+                        thought=h.thought,
+                        action=h.name,
+                        arguments=h.arguments,
+                        observation=h.observation
+                    )
+                    for h in history
+                ]
+            )
+
+            self.recursive_input.send(r, properties={"id": id})
+
+            print("Done.", flush=True)
+
+            return
+
         except Exception as e:
 
             print(f"Exception: {e}")
@@ -184,6 +255,30 @@ class Processor(ConsumerProducer):
             help=f'Prompt response queue (default: {pr_response_queue})',
         )
 
+        parser.add_argument(
+            '--text-completion-request-queue',
+            default=tc_request_queue,
+            help=f'Text completion request queue (default: {tc_request_queue})',
+        )
+
+        parser.add_argument(
+            '--text-completion-response-queue',
+            default=tc_response_queue,
+            help=f'Text completion response queue (default: {tc_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})',
+        )
+
 def run():
 
     Processor.start(module, __doc__)

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

@@ -0,0 +1,56 @@
+
+from . types import Tool, Argument
+
+class CatsKb:
+    tool = Tool(
+        name = "cats-kb",
+        description = "Query a knowledge base with information about Mark's cats.  The query should be a simple natural language question",
+        arguments = [
+            Argument(
+                name = "query",
+                type = "string",
+                description = "The search query string"
+            )
+        ]
+    )
+    def __init__(self, context):
+        self.context = context
+    def invoke(self, **arguments):
+        return self.context.graph_rag.request(arguments.get("query"))
+
+class ShuttleKb:
+    tool = Tool(
+        name = "shuttle-kb",
+        description = "Query a knowledge base with information about the space shuttle.  The query should be a simple natural language question",
+        arguments = [
+            Argument(
+                name = "query",
+                type = "string",
+                description = "The search query string"
+            )
+        ]
+    )
+    def __init__(self, context):
+        self.context = context
+    def invoke(self, **arguments):
+        return self.context.graph_rag.request(arguments.get("query"))
+
+class Compute:
+    tool = Tool(
+        name = "compute",
+        description = "A computation engine which can answer questions about maths and computation",
+        arguments = [
+            Argument(
+                name = "computation",
+                type = "string",
+                description = "The computation to solve"
+            )
+        ]
+    )
+    def __init__(self, context):
+        self.context = context
+    def invoke(self, **arguments):
+        return self.context.prompt.request(
+            "question", { "question": arguments.get("computation") }
+        )
+

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

@@ -0,0 +1,27 @@
+
+import dataclasses
+
+@dataclasses.dataclass
+class Argument:
+    name : str
+    type : str
+    description : str
+
+@dataclasses.dataclass
+class Tool:
+    name : str
+    description : str
+    arguments : list[Argument]
+    
+@dataclasses.dataclass
+class Action:
+    thought : str
+    name : str
+    arguments : dict
+    observation : str
+    
+@dataclasses.dataclass
+class Final:
+    thought : str
+    final : str
+