From b4b39f1e2a8838a2a822003a12d5f77e55db71f8 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sun, 20 Jul 2025 11:02:21 +0100 Subject: [PATCH] New template output parser --- .../trustgraph/agent/react/agent_manager.py | 185 +++++++++++++++--- 1 file changed, 163 insertions(+), 22 deletions(-) diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py index 391f188b..24282031 100644 --- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py +++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py @@ -1,6 +1,7 @@ import logging import json +import re from . types import Action, Final @@ -12,6 +13,154 @@ class AgentManager: self.tools = tools self.additional_context = additional_context + def parse_react_response(self, text): + """Parse text-based ReAct response format. + + Expected format: + Thought: [reasoning about what to do next] + Action: [tool_name] + Args: { + "param": "value" + } + + OR + + Thought: [reasoning about the final answer] + Final Answer: [the answer] + """ + if not isinstance(text, str): + raise ValueError(f"Expected string response, got {type(text)}") + + # Remove any markdown code blocks that might wrap the response + text = re.sub(r'^```[^\n]*\n', '', text.strip()) + text = re.sub(r'\n```$', '', text.strip()) + + lines = text.strip().split('\n') + + thought = None + action = None + args = None + final_answer = None + + i = 0 + while i < len(lines): + line = lines[i].strip() + + # Parse Thought + if line.startswith("Thought:"): + thought = line[8:].strip() + # Handle multi-line thoughts + i += 1 + while i < len(lines): + next_line = lines[i].strip() + if next_line.startswith(("Action:", "Final Answer:", "Args:")): + break + thought += " " + next_line + i += 1 + continue + + # Parse Final Answer + if line.startswith("Final Answer:"): + final_answer = line[13:].strip() + # Handle multi-line final answers (including JSON) + i += 1 + + # Check if the answer might be JSON + if final_answer.startswith('{') or (i < len(lines) and lines[i].strip().startswith('{')): + # Collect potential JSON answer + json_text = final_answer if final_answer.startswith('{') else "" + brace_count = json_text.count('{') - json_text.count('}') + + while i < len(lines) and (brace_count > 0 or not json_text): + current_line = lines[i].strip() + if current_line.startswith(("Thought:", "Action:")) and brace_count == 0: + break + json_text += ("\n" if json_text else "") + current_line + brace_count += current_line.count('{') - current_line.count('}') + i += 1 + + # Try to parse as JSON + try: + final_answer = json.loads(json_text) + except json.JSONDecodeError: + # Not valid JSON, treat as regular text + final_answer = json_text + else: + # Regular text answer + while i < len(lines): + next_line = lines[i].strip() + if next_line.startswith(("Thought:", "Action:")): + break + final_answer += " " + next_line + i += 1 + + # If we have a final answer, return Final object + return Final( + thought=thought or "", + final=final_answer + ) + + # Parse Action + if line.startswith("Action:"): + action = line[7:].strip() + + # Parse Args + if line.startswith("Args:"): + # Check if JSON starts on the same line + args_on_same_line = line[5:].strip() + if args_on_same_line: + args_text = args_on_same_line + brace_count = args_on_same_line.count('{') - args_on_same_line.count('}') + else: + args_text = "" + brace_count = 0 + + # Collect all lines that form the JSON arguments + i += 1 + started = bool(args_on_same_line and '{' in args_on_same_line) + + while i < len(lines) and (not started or brace_count > 0): + current_line = lines[i] + args_text += ("\n" if args_text else "") + current_line + + # Count braces to determine when JSON is complete + for char in current_line: + if char == '{': + brace_count += 1 + started = True + elif char == '}': + brace_count -= 1 + + # If we've started and braces are balanced, we're done + if started and brace_count == 0: + break + + i += 1 + + # Parse the JSON arguments + try: + args = json.loads(args_text.strip()) + except json.JSONDecodeError as e: + logger.error(f"Failed to parse JSON arguments: {args_text}") + raise ValueError(f"Invalid JSON in Args: {e}") + + i += 1 + + # If we have an action, return Action object + if action: + return Action( + thought=thought or "", + name=action, + arguments=args or {}, + observation="" + ) + + # If we only have a thought but no action or final answer + if thought and not action and not final_answer: + raise ValueError(f"Response has thought but no action or final answer: {text}") + + raise ValueError(f"Could not parse response: {text}") + async def reason(self, question, history, context): print(f"calling reason: {question}", flush=True) @@ -62,31 +211,23 @@ class AgentManager: logger.info(f"prompt: {variables}") - obj = await context("prompt-request").agent_react(variables) + # Get text response from prompt service + response_text = await context("prompt-request").agent_react(variables) - print(json.dumps(obj, indent=4), flush=True) + print(f"Response text:\n{response_text}", flush=True) - logger.info(f"response: {obj}") + logger.info(f"response: {response_text}") - 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 + # Parse the text response + try: + result = self.parse_react_response(response_text) + logger.info(f"Parsed result: {result}") + return result + except ValueError as e: + logger.error(f"Failed to parse response: {e}") + # Try to provide a helpful error message + logger.error(f"Response was: {response_text}") + raise RuntimeError(f"Failed to parse agent response: {e}") async def react(self, question, history, think, observe, context):