From 7a52e947797ae1f7a4390d23dd29bca4e2c41c5a Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Sun, 10 Nov 2024 10:52:28 +0000
Subject: [PATCH] Now working with the tool descriptor arguments

---
 .../trustgraph/agent/react/agent_manager.py   | 36 ++++++++++++-------
 .../trustgraph/agent/react/service.py         | 27 ++++++--------
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py
index 094b329b..34964ff5 100644
--- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py
+++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py
@@ -13,10 +13,10 @@ class AgentManager:
 access to the following functions:
 
 {% for tool in tools %}{
-    "function": "{{ tool.tool.name }}",
-    "description": "{{ tool.tool.description }}",
+    "function": "{{ tool.name }}",
+    "description": "{{ tool.description }}",
     "arguments": [
-{% for arg in tool.tool.arguments %}        {
+{% for arg in tool.arguments %}        {
             "name": "{{ arg.name }}",
             "type": "{{ arg.type }}",
             "description": "{{ arg.description }}",
@@ -91,11 +91,25 @@ Input:
         tools = self.tools
 
         tool_names = ",".join([
-            t.tool.name for t in self.tools
+            t for t in self.tools.keys()
         ])
 
         prompt = tpl.render({
-            "tools": tools,
+            "tools": [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "arguments": [
+                        {
+                            "name": arg.name,
+                            "type": arg.type,
+                            "description": arg.description
+                        }
+                        for arg in tool.arguments.values()
+                    ]
+                }
+                for tool in self.tools.values()
+            ],
             "question": question,
             "tool_names": tool_names,
             "history": [
@@ -159,16 +173,12 @@ Input:
 
             think(act.thought)
 
-            action = None
-
-            for tool in self.tools:
-                if tool.tool.name == act.name:
-                    action = tool
-
-            if action is None:
+            if act.name in self.tools:
+                action = self.tools[act.name]
+            else:
                 raise RuntimeError(f"No action for {act.name}!")
 
-            resp = action.invoke(**act.arguments)
+            resp = action.implementation.invoke(**act.arguments)
 
             resp = resp.strip()
 
diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py
index e1bd75c6..8e26a035 100755
--- a/trustgraph-flow/trustgraph/agent/react/service.py
+++ b/trustgraph-flow/trustgraph/agent/react/service.py
@@ -24,7 +24,8 @@ from ... clients.graph_rag_client import GraphRagClient
 
 from . tools import KnowledgeQueryImpl, TextCompletionImpl
 from . agent_manager import AgentManager
-from . types import Final, Action, Tool
+
+from . types import Final, Action, Tool, Argument
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -55,9 +56,9 @@ class Processor(ConsumerProducer):
                     )
 
                 if ttoks[0] == "knowledge-query":
-                    impl = KnowledgeQueryImpl
+                    impl = KnowledgeQueryImpl(self)
                 elif ttoks[0] == "text-completion":
-                    impl = TextCompletionImpl
+                    impl = TextCompletionImpl(self)
                 else:
                     raise RuntimeError(
                         f"Tool-kind {ttoks[0]} not known"
@@ -111,15 +112,13 @@ class Processor(ConsumerProducer):
                     )
                 if toks[0] not in tool_base:
                     raise RuntimeError(f"Description, tool {toks[0]} not known")
-                tool_base[toks[0]].arguments[ttoks[0]] = {
-                    "name": ttoks[0],
-                    "type": ttoks[1],
-                    "description": ttoks[2],
-                }
+                tool_base[toks[0]].arguments[ttoks[0]] = Argument(
+                    name = ttoks[0],
+                    type = ttoks[1],
+                    description = ttoks[2]
+                )
 
-        print(json.dumps({k: str(v) for k, v in tool_base.items()}, indent=4))
-
-        sys.exit(0)
+#         print(json.dumps({k: str(v) for k, v in tool_base.items()}, indent=4))
 
         input_queue = params.get("input_queue", default_input_queue)
         output_queue = params.get("output_queue", default_output_queue)
@@ -186,11 +185,7 @@ class Processor(ConsumerProducer):
             schema=JsonSchema(AgentRequest),
         )
 
-        tools = [
-#            CatsKb(self), ShuttleKb(self), Compute(self),
-        ]
-
-        self.agent = AgentManager(self, tools)
+        self.agent = AgentManager(self, tool_base)
 
     def parse_json(self, text):
         json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)