From 084d77aff9412dae3a3b23200ff77269f400d2d3 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Sun, 10 Nov 2024 11:36:56 +0000
Subject: [PATCH] Add additional context

---
 .../trustgraph/agent/react/README.md          |  3 ++
 .../trustgraph/agent/react/agent_manager.py   | 10 +++++--
 .../trustgraph/agent/react/service.py         | 29 ++++++++++++-------
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/trustgraph-flow/trustgraph/agent/react/README.md b/trustgraph-flow/trustgraph/agent/react/README.md
index 9965194a..dd5cbea5 100644
--- a/trustgraph-flow/trustgraph/agent/react/README.md
+++ b/trustgraph-flow/trustgraph/agent/react/README.md
@@ -14,3 +14,6 @@ agent-manager-react \
         shuttle="query:string:The search query string" \
         compute="computation:string:The computation to solve"
 
+
+    --context 'The space shuttle challenger final mission was 58-L'
+
diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py
index 34964ff5..11b258a6 100644
--- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py
+++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py
@@ -62,7 +62,9 @@ format in your output:
 
 Respond by describing either one single thought/action/arguments or
 the final-answer.  Pause after providing one action or final-answer.
-Begin!
+
+{% if context %}Additional context has been provided:
+{{context}}{% endif %}
 
 Question: {{question}}
 
@@ -80,9 +82,10 @@ Input:
 }
 {% endfor %}"""
 
-    def __init__(self, context, tools):
+    def __init__(self, context, tools, additional_context=None):
         self.context = context
         self.tools = tools
+        self.additional_context = additional_context
 
     def reason(self, question, history):
 
@@ -110,6 +113,7 @@ Input:
                 }
                 for tool in self.tools.values()
             ],
+            "context": self.additional_context,
             "question": question,
             "tool_names": tool_names,
             "history": [
@@ -123,6 +127,8 @@ Input:
             ],
         })
 
+        print(prompt)
+
         logger.info(f"prompt: {prompt}")
 
         resp = self.context.prompt.request(
diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py
index 8e26a035..f367d10f 100755
--- a/trustgraph-flow/trustgraph/agent/react/service.py
+++ b/trustgraph-flow/trustgraph/agent/react/service.py
@@ -37,7 +37,9 @@ class Processor(ConsumerProducer):
 
     def __init__(self, **params):
 
-        tool_base = {}
+        additional = params.get("context", None)
+
+        tools = {}
 
         # Parsing the prompt information to the prompt configuration
         # structure
@@ -65,7 +67,7 @@ class Processor(ConsumerProducer):
                     )
 
                 if len(ttoks) == 1:
-                    tool_base[toks[0]] = Tool(
+                    tools[toks[0]] = Tool(
                         name = ttoks[0],
                         description = "",
                         implementation = impl,
@@ -73,7 +75,7 @@ class Processor(ConsumerProducer):
                         arguments = {},
                     )
                 else:
-                    tool_base[toks[0]] = Tool(
+                    tools[toks[0]] = Tool(
                         name = ttoks[0],
                         description = "",
                         implementation = impl,
@@ -91,9 +93,9 @@ class Processor(ConsumerProducer):
                     raise runtimeerror(
                         f"tool-type string not well-formed: {t}"
                     )
-                if toks[0] not in tool_base:
+                if toks[0] not in tools:
                     raise runtimeerror(f"description, tool {toks[0]} not known")
-                tool_base[toks[0]].description = toks[1]
+                tools[toks[0]].description = toks[1]
 
         # Parsing the prompt information to the prompt configuration
         # structure
@@ -110,16 +112,14 @@ class Processor(ConsumerProducer):
                     raise RuntimeError(
                         f"Tool argument string not well-formed: {t}"
                     )
-                if toks[0] not in tool_base:
+                if toks[0] not in tools:
                     raise RuntimeError(f"Description, tool {toks[0]} not known")
-                tool_base[toks[0]].arguments[ttoks[0]] = Argument(
+                tools[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))
-
         input_queue = params.get("input_queue", default_input_queue)
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
@@ -185,7 +185,11 @@ class Processor(ConsumerProducer):
             schema=JsonSchema(AgentRequest),
         )
 
-        self.agent = AgentManager(self, tool_base)
+        self.agent = AgentManager(
+            context=self,
+            tools=tools,
+            additional_context=additional
+        )
 
     def parse_json(self, text):
         json_match = re.search(r'```(?:json)?(.*?)```', text, re.DOTALL)
@@ -385,6 +389,11 @@ it is read by the LLM and used to determine how to use the argument.
 description.'''
         )
 
+        parser.add_argument(
+            '--context', 
+            help=f'Optional, specifies additional context text for the LLM.'
+        )
+
 def run():
 
     Processor.start(module, __doc__)