New template output parser

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):