trustgraph/agent/agent-hacking11

#!/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

def wrap(text, width=75):

    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"
pulsar_host = "pulsar://localhost:6650"

prompt_client = PromptClient(pulsar_host=pulsar_host)
rag_client = GraphRagClient(pulsar_host=pulsar_host)

def graph_rag_query(q):
    resp = rag_client.request(q)
    return resp

tools = [
    {
        "function": "cats-kb",
        "description": "Query a knowledge base with information about Mark's cats.  The query should be a simple natural language question",
        "arguments": [
            {
                "name": "query",
                "type": "string",
                "description": "The search query string"
            }
        ]
    },
    {
        "function": "shuttle-kb",
        "description": "Query a knowledge base with information about the space shuttle.  The query should be a simple natural language question",
        "arguments": [
            {
                "name": "query",
                "type": "string",
                "description": "The search query string"
            }
        ]
    },
    {
        "function": "compute",
        "description": "A computation engine which can answer questions about maths",
        "arguments": [
            {
                "name": "query",
                "type": "string",
                "description": "The computation"
            }
        ]
    }
]

template="""Answer the following questions as best you can. You have access to the following tools:

{% for tool in tools %}{
    "function": "{{ tool.function }}",
    "description": "{{ tool.description }}",
    "arguments": [
{% for arg in tool.arguments %}        {
            "name": "{{ arg.name }}",
            "type": "{{ arg.type }}",
            "description": "{{ arg.description }}",
        }
{% endfor %}
    ]
}
{% endfor %}
To call a function, respond - immediately and only - with a JSON object of the following format:
{
    "action": "function_name",
    "arguments": {
        "argument1": "argument_value",
        "argument2": "argument_value"
    }
}

Each step has the following format in your output:

{
  "thought": "you should always think about what to do",
  "action": "the action to take, should be one of [{{tool_names}}]",
  "arguments": {
      "argument1": action argument,
      "argument2": action argument2
  },
  "observation": "the result of the action",
}
... (this JSON object can repeat N times)
{
  "thought": "I now know the final answer",
  "final-answer": "the final answer to the original input question"
}

Respond by describing either one single thought/action/arguments or the final-answer.  Pause after providing each action.

Begin!

Question: {{question}}

Input:
{% for h in history %}{{h}}

{% endfor %}"""

tpl = ibis.Template(template)

history = []

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()

for iter in range(0, 10):

    prompt = tpl.render({
        "tools": tools,
        "question": q,
        "tool_names": ",".join([t["function"] for t in tools]),
        "history": [json.dumps(h, indent=2) for h in history],
    })

    print("-" * 76)
    print(prompt)
    print("-" * 76)

    resp = prompt_client.request(
        "question",
        {
            "question": prompt
        }
    )

    print("-" * 76)
    print(resp)
    print("-" * 76)

    resp = resp.replace("```json", "")
    resp = resp.replace("```", "")

    obj = json.loads(resp)

    output(obj["thought"], "? ")
    print()

    if "final-answer" in obj:
        history.append(obj)
        break

    if "action" in obj:
        if obj["action"] == "cats-kb":
            ans = graph_rag_query(obj["arguments"]["query"])
            obj["observation"] = ans.strip()
        elif obj["action"] == "compute":
            if "196" in obj["arguments"]["query"]:
                obj["observation"] = "196 > 1.319501555"
            else:
                obj["observation"] = "The answer is 1.319501155"
        elif obj["action"] == "shuttle-kb":
            ans = graph_rag_query(obj["arguments"]["query"])
            obj["observation"] = ans.strip()
        else:
            raise RuntimeError("Unknown action:", obj["action"])

        output(obj["observation"], "! ")
        print()

        history.append(obj)

        print("-" * 76)
        print(history)
        print("-" * 76)
        
        continue

    raise RuntimeError("Iteration output is not action or final-answer")

output(obj["final-answer"], "< ")