mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-21 11:11:03 +02:00
302 lines
7.5 KiB
Python
Executable file
302 lines
7.5 KiB
Python
Executable file
#!/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)
|
|
|