Agent infra

agent/agent-demo7

@@ -0,0 +1,302 @@
+#!/usr/bin/env python3
+
+from trustgraph.clients.prompt_client import PromptClient
+import json
+import textwrap
+import ibis
+import textwrap
+import time
+import dataclasses
+
+class Input:
+    state: dict
+    plan_state: str
+    plan: str
+    iteration: int
+
+class Output:
+    state: dict
+    thought: str
+    next_state: str
+    plan: str
+    iteration: int
+    
+
+def output(text, prefix="> ", width=78):
+
+    out = textwrap.wrap(
+        text, initial_indent=prefix, subsequent_indent=prefix,
+        width=width
+    )
+    print("\n".join(out))
+
+pulsar_host = "pulsar://localhost:6650"
+
+prompt_client = PromptClient(pulsar_host=pulsar_host)
+
+class State:
+    pass
+
+class ActionWikipedia:
+
+    def __init__(self, prompt_client):
+        self.prompt_client = prompt_client
+
+    def act(self, state, topic):
+
+        pr = ibis.Template(
+            "Tell me what you think Wikipedia would say about {{topic}}"
+        )
+
+        input = pr.render({
+            "topic": topic,
+        })
+
+        resp = self.prompt_client.request(
+            "question",
+            {
+                "question": input
+            }
+        )
+
+        state.previous.append(
+            f"Question: {input.strip()}"
+        )
+
+        state.previous.append(
+            f"Answer: {resp.strip()}"
+        )
+
+        return state
+
+class ActionCalculate:
+
+    def __init__(self, prompt_client):
+        self.prompt_client = prompt_client
+
+    def act(self, state, calc):
+
+        pr = ibis.Template(
+            "Compute this: {{computation}}.  Just answer the question without explanation."
+        )
+
+        input = pr.render({
+            "computation": calc,
+        })
+
+        resp = self.prompt_client.request(
+            "question",
+            {
+                "question": input
+            }
+        )
+
+        print(resp)
+
+        state.previous.append(
+            "Calculate: " + calc.strip()
+        )
+
+        state.previous.append(
+            "Answer: " + resp.strip()
+        )
+
+        return state
+
+class ActionAnswer:
+
+    def __init__(self, prompt_client):
+        self.prompt_client = prompt_client
+
+    def act(self, state, question):
+
+        pr = ibis.Template(
+            "Answer this: {{question}}"
+        )
+
+        input = pr.render({
+            "question": question,
+        })
+
+        resp = self.prompt_client.request(
+            "question",
+            {
+                "question": input
+            }
+        )
+
+        state.facts.append(
+            "Answer to " + input + ":\n" +
+            resp
+        )
+
+        return state
+
+tools = {
+    "calculate": {
+        "description": "Takes a numeric computation and calculates the answer",
+        "implementation": ActionCalculate,
+    },
+    "cats-knowledge-store": {
+        "description": "Answer questions on Mark's cats using a knowledge store",
+        "implementation": ActionCalculate,
+    },
+    # "wikipedia": {
+    #     "description": "Takes a query and looks it up on Wikipedia",
+    #     "implementation": ActionWikipedia,
+    # },
+    # "answer": {
+    #     "description": "Take a simple question and use the LLM to provide the answer",
+    #     "implementation": ActionAnswer,
+    # },
+    # "space-shuttle-rag": {
+    #     "description": "Take a question about space shuttles and answer it using the TrustGraph GraphRAG service",
+    #     "implementation": ActionAnswer,
+    # }
+}
+
+class AgentManager:
+
+    determine_prompt = ibis.Template("""
+You have access to the following tools:
+{% for id, tool in tools.items() %}- tool-name: {{id}}
+  tool-description: {{tool.description}}
+{% endfor %}
+You operate in a loop.  For each iteration of the loop, you take the
+question and learnt knowledge.  If the knowledge is enough to answer
+the question, you will response with an answer.  If the knowledge is not
+enough, you will response with an action to be invoked to acquire more
+knowledge.
+
+Your output must ALWAYS be a well-formed JSON object.  Output ONLY a JSON
+object with no explanatory text or markup formatting.  The JSON object
+can have the following fields:
+- thought: Mandatory.  Your explanation of the current step, including the
+  goal and the reason for choosing this action.  A string.
+- action: An optional field, one of the tool-names above if you are using
+  an action.
+- action-input: An optional field, input to the selected tool if you are using
+  an action.
+- failure: a boolean.  Set to true if you cannot proceed.
+- final-answer: Your answer to the question as a string.
+
+When you know the final answer to the original question, emit a 'final-answer'
+JSON object with the following fields: thought and final-answer.
+
+When you don't know the answer but know an action to move towards the
+solution, emit an 'action' JSON object with the following fields:
+thought, action, action-input
+
+When you are stuck and aren't able to move torwards a solution, emit a
+'failure' JSON object the following fields: thought, failure
+
+{{previous}}""")
+
+    def __init__(self, tools, prompt_client):
+
+        self.tools = tools
+        self.prompt_client = prompt_client
+
+    def determine(self, state, thought=None):
+
+        input = __class__.determine_prompt.render({
+            "previous": "\n".join(state.previous),
+            "tools": self.tools,
+        })
+
+        print(input)
+
+        resp = self.prompt_client.request(
+            "question",
+            {
+                "question": input
+            }
+        )
+
+        resp = resp.replace("```json", "")
+        resp = resp.replace("```", "")
+        print(resp)
+    #     print(resp)
+        resp = json.loads(resp)
+    #     print(json.dumps(resp, indent=4))
+        return resp
+
+    def invoke(self, q, thought):
+
+        state = State()
+        state.question = q
+        state.previous = [
+            f"Question: {q}"
+        ]
+
+        while True:
+
+            resp = self.determine(state)
+
+            print(resp)
+
+            if "thought" in resp:
+                if thought:
+                    thought(resp["thought"])
+                state.previous.append(f"Thought: {resp['thought'].strip()}")
+
+            if "failure" in resp:
+                if resp["failure"]:
+                    print("Failed")
+                    return "Failed"
+
+            if "final-answer" in resp:
+                return resp["final-answer"]
+
+            if "action" not in resp:
+                raise RuntimeError("Didn't get final-answer/action response")
+
+            print(resp)
+
+            action = resp["action"]
+
+            print(action)
+
+            if action not in self.tools:
+                raise RuntimeError(f"Tool {action} not known")
+
+            if "implementation" not in self.tools[action]:
+                raise RuntimeError(f"Tool {action} not implemented")
+
+            if "action-input" in resp:
+                input = resp["action-input"]
+            else:
+                input = None
+
+            print("Action:", action)
+            print("Action input:", input)
+            print()
+
+            impl_class = self.tools[action]["implementation"]
+            impl = impl_class(self.prompt_client)
+            state = impl.act(state, input)
+
+        time.sleep(2)
+               
+def thought(t):
+    output("\U0001f914... " + t, prefix="? ")
+    print()
+
+q = """Take the square root of 1600.  That number is the number of a US
+President.  Who is that president and how do they relate to the space
+shuttle programme?
+"""
+
+q = "How many eggs do I have if I have 3 eggs and Fred gives me some more eggs.  The number of eggs he gives me is the square root of 64."
+
+q = "Find out the behavioral nature of Mark's cats"
+
+am = AgentManager(tools, prompt_client)
+
+output("Q: " + q)
+print()
+
+resp = am.invoke(q, thought)
+
+print(resp)
+