diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py new file mode 100644 index 00000000..6d292e11 --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 754dc194..db93e3b7 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/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__) diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py new file mode 100644 index 00000000..918279db --- /dev/null +++ b/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") } + ) + diff --git a/trustgraph-flow/trustgraph/agent/react/types.py b/trustgraph-flow/trustgraph/agent/react/types.py new file mode 100644 index 00000000..efc2ef6c --- /dev/null +++ b/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 +