Plan-b working

agent/plan-b-4

@@ -0,0 +1,327 @@
+#!/usr/bin/env python3
+
+from trustgraph.clients.prompt_client import PromptClient
+from trustgraph.clients.graph_rag_client import GraphRagClient
+import json
+import textwrap
+import ibis
+import textwrap
+import time
+import dataclasses
+import sys
+import base64
+
+def wrap(text, width=75):
+
+    if text is None: text = "n/a"
+
+    out = textwrap.wrap(
+        text, width=width
+    )
+    return "\n".join(out)
+
+def output(text, prefix="> ", width=78):
+
+    out = textwrap.indent(
+        text, prefix=prefix
+    )
+    print(out)
+
+pulsar_host = "pulsar://localhost:6650"
+
+prompt_client = PromptClient(pulsar_host=pulsar_host)
+rag_client = GraphRagClient(pulsar_host=pulsar_host)
+
+@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
+    
+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, client):
+        self.client = client
+    def invoke(self, **arguments):
+        return self.client.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, client):
+        self.client = client
+    def invoke(self, **arguments):
+        return self.client.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, client):
+        self.client = client
+    def invoke(self, **arguments):
+        return self.client.request(
+            "question", { "question": arguments.get("computation") }
+        )
+
+tools = [
+    CatsKb(rag_client),
+    ShuttleKb(rag_client),
+    Compute(prompt_client),
+]
+
+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, client, tools, question):
+        self.client = client
+        self.tools = tools
+        self.history = []
+        self.question = question
+
+    def iterate(self):
+
+        tpl = ibis.Template(self.template)
+
+        tools = self.tools
+
+        tool_names = ",".join([
+            t.tool.name for t in self.tools
+        ])
+
+        prompt = tpl.render({
+            "tools": tools,
+            "question": self.question,
+            "tool_names": tool_names,
+            "history": [
+                {
+                    "thought": h.thought,
+                    "action": h.name,
+                    "arguments": h.arguments,
+                    "observation": h.observation,
+                }
+                for h in self.history
+            ],
+        })
+
+#         print(prompt)
+
+        resp = prompt_client.request(
+            "question",
+            {
+                "question": prompt
+            }
+        )
+
+        resp = resp.replace("```json", "")
+        resp = resp.replace("```", "")
+
+#         print("---")
+#         print(resp)
+#         print("---")
+
+#         print(base64.b64encode(resp.encode("utf-8")).decode("utf-8"))
+
+        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 = "FIXME"
+            )
+
+            return a
+
+    def think(self, thought):
+
+#         print(thought)
+
+        output(wrap(thought), "? ")
+        print()
+
+    def invoke(self):
+
+        for i in range(0, 10):
+
+#             print("-----------------------")
+#             print("Iter", i)
+
+            act = self.iterate()
+#             print(act)
+
+            if isinstance(act, Final):
+
+                self.think(act.thought)
+                return act.final
+
+            else:
+
+                self.think(act.thought)
+
+#                 print(act.name)
+
+                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}!")
+
+                resp = action.invoke(**act.arguments)
+
+                resp = resp.strip()
+
+#                 print(resp)
+
+                act.observation = resp
+
+#                 print("act>", act)
+
+                self.history.append(act)
+
+#                 print(self.history)
+
+        raise RuntimeError("Too many iterations")
+        
+
+q = "How many cats does Mark have? Calculate that number raised to 0.4 power.  Is that number higher 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."
+
+output(wrap(q))
+print()
+
+am = AgentManager(prompt_client, tools, q)
+
+resp = am.invoke()
+
+output(resp, "< ")
+print()
+