From 66a74abfda1c1de14830f96cf7b456d4dd94f7c3 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 26 Nov 2025 00:17:17 +0000 Subject: [PATCH] Streaming parser for agent responses --- .../trustgraph/base/prompt_client.py | 12 +- .../trustgraph/agent/react/agent_manager.py | 76 +++- .../trustgraph/agent/react/service.py | 40 ++- .../agent/react/streaming_parser.py | 339 ++++++++++++++++++ 4 files changed, 445 insertions(+), 22 deletions(-) create mode 100644 trustgraph-flow/trustgraph/agent/react/streaming_parser.py diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index ccdf8e9c..dd3b1e88 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -1,12 +1,13 @@ import json +import asyncio from . request_response_spec import RequestResponse, RequestResponseSpec from .. schema import PromptRequest, PromptResponse class PromptClient(RequestResponse): - async def prompt(self, id, variables, timeout=600, streaming=False): + async def prompt(self, id, variables, timeout=600, streaming=False, chunk_callback=None): if not streaming: # Non-streaming path @@ -42,6 +43,12 @@ class PromptClient(RequestResponse): if resp.text: full_text += resp.text + # Call chunk callback if provided + if chunk_callback: + if asyncio.iscoroutinefunction(chunk_callback): + await chunk_callback(resp.text) + else: + chunk_callback(resp.text) elif resp.object: full_object = resp.object @@ -108,12 +115,13 @@ class PromptClient(RequestResponse): timeout = timeout, ) - async def agent_react(self, variables, timeout=600, streaming=False): + async def agent_react(self, variables, timeout=600, streaming=False, chunk_callback=None): return await self.prompt( id = "agent-react", variables = variables, timeout = timeout, streaming = streaming, + chunk_callback = chunk_callback, ) async def question(self, question, timeout=600): diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py index 70b7597b..35df2c59 100644 --- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py +++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py @@ -2,6 +2,7 @@ import logging import json import re +import asyncio from . types import Action, Final @@ -169,7 +170,7 @@ class AgentManager: raise ValueError(f"Could not parse response: {text}") - async def reason(self, question, history, context, streaming=False): + async def reason(self, question, history, context, streaming=False, think=None, observe=None, answer=None): logger.debug(f"calling reason: {question}") @@ -219,25 +220,59 @@ class AgentManager: logger.info(f"prompt: {variables}") - # Get text response from prompt service - response_text = await context("prompt-request").agent_react(variables, streaming=streaming) + # Streaming path - use StreamingReActParser + if streaming and think: + from .streaming_parser import StreamingReActParser - logger.debug(f"Response text:\n{response_text}") + # Create parser with streaming callbacks + # Thought chunks go to think(), answer chunks go to answer() + parser = StreamingReActParser( + on_thought_chunk=lambda chunk: asyncio.create_task(think(chunk)), + on_answer_chunk=lambda chunk: asyncio.create_task(answer(chunk) if answer else think(chunk)), + ) - logger.info(f"response: {response_text}") + # Create async chunk callback that feeds parser + async def on_chunk(text): + parser.feed(text) + + # Get streaming response + response_text = await context("prompt-request").agent_react( + variables, + streaming=True, + chunk_callback=on_chunk + ) + + # Finalize parser + parser.finalize() + + # Get result + result = parser.get_result() + if result is None: + raise RuntimeError("Parser failed to produce a result") - # 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, streaming=False): + else: + # Non-streaming path - get complete text and parse + response_text = await context("prompt-request").agent_react(variables, streaming=False) + + logger.debug(f"Response text:\n{response_text}") + + logger.info(f"response: {response_text}") + + # 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, streaming=False, answer=None): logger.info(f"question: {question}") @@ -246,17 +281,26 @@ class AgentManager: history = history, context = context, streaming = streaming, + think = think, + observe = observe, + answer = answer, ) logger.info(f"act: {act}") if isinstance(act, Final): - await think(act.thought) + # In non-streaming mode, send complete thought + # In streaming mode, thoughts were already sent as chunks + if not streaming: + await think(act.thought) return act else: - await think(act.thought) + # In non-streaming mode, send complete thought + # In streaming mode, thoughts were already sent as chunks + if not streaming: + await think(act.thought) logger.debug(f"ACTION: {act.name}") diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 8edb6a9b..2fb5b9c9 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -270,6 +270,34 @@ class Processor(AgentService): await respond(r) + async def answer(x): + + logger.debug(f"Answer: {x}") + + if streaming: + # Streaming format + r = AgentResponse( + chunk_type="answer", + content=x, + end_of_message=False, # More chunks may follow + end_of_dialog=False, + # Legacy fields for backward compatibility + answer=None, + error=None, + thought=None, + observation=None, + ) + else: + # Legacy format - shouldn't be called in non-streaming mode + r = AgentResponse( + answer=x, + error=None, + thought=None, + observation=None, + ) + + await respond(r) + # Apply tool filtering based on request groups and state filtered_tools = filter_tools_by_group_and_state( tools=self.agent.tools, @@ -306,6 +334,7 @@ class Processor(AgentService): history = history, think = think, observe = observe, + answer = answer, context = UserAwareContext(flow, request.user), streaming = streaming, ) @@ -322,10 +351,11 @@ class Processor(AgentService): f = json.dumps(act.final) if streaming: - # Streaming format - mark as final dialog + # Streaming format - send end-of-dialog marker + # Answer chunks were already sent via think() callback during parsing r = AgentResponse( chunk_type="answer", - content=f, + content="", # Empty content, just marking end of dialog end_of_message=True, end_of_dialog=True, # Legacy fields for backward compatibility @@ -334,7 +364,7 @@ class Processor(AgentService): thought=None, ) else: - # Legacy format + # Legacy format - send complete answer r = AgentResponse( answer=act.final, error=None, @@ -369,7 +399,9 @@ class Processor(AgentService): observation=h.observation ) for h in history - ] + ], + user=request.user, + streaming=streaming, ) await next(r) diff --git a/trustgraph-flow/trustgraph/agent/react/streaming_parser.py b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py new file mode 100644 index 00000000..76192e92 --- /dev/null +++ b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py @@ -0,0 +1,339 @@ +""" +Streaming parser for ReAct responses. + +This parser handles text chunks from LLM streaming responses and parses them +into ReAct format (Thought/Action/Args or Thought/Final Answer). It maintains +state across chunk boundaries to handle cases where delimiters or JSON are split. + +Key challenges: +- Delimiters may be split across chunks: "Tho" + "ught:" or "Final An" + "swer:" +- JSON arguments may be split: '{"loc' + 'ation": "NYC"}' +- Need to emit thought/answer chunks as they arrive for streaming +""" + +import json +import logging +import re +from enum import Enum +from typing import Optional, Callable, Any +from . types import Action, Final + +logger = logging.getLogger(__name__) + + +class ParserState(Enum): + """States for the streaming ReAct parser state machine""" + INITIAL = "initial" # Waiting for first content + THOUGHT = "thought" # Accumulating thought content + ACTION = "action" # Found "Action:", collecting action name + ARGS = "args" # Found "Args:", collecting JSON arguments + FINAL_ANSWER = "final_answer" # Found "Final Answer:", collecting answer + COMPLETE = "complete" # Parsing complete, object ready + + +class StreamingReActParser: + """ + Stateful parser for streaming ReAct responses. + + 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] + + Usage: + parser = StreamingReActParser( + on_thought_chunk=lambda chunk: print(f"Thought: {chunk}"), + on_answer_chunk=lambda chunk: print(f"Answer: {chunk}"), + ) + + for chunk in llm_stream: + parser.feed(chunk) + if parser.is_complete(): + result = parser.get_result() + break + """ + + # Delimiters we're looking for + THOUGHT_DELIMITER = "Thought:" + ACTION_DELIMITER = "Action:" + ARGS_DELIMITER = "Args:" + FINAL_ANSWER_DELIMITER = "Final Answer:" + + # Maximum buffer size for delimiter detection (longest delimiter + safety margin) + MAX_DELIMITER_BUFFER = 20 + + def __init__( + self, + on_thought_chunk: Optional[Callable[[str], Any]] = None, + on_answer_chunk: Optional[Callable[[str], Any]] = None, + ): + """ + Initialize streaming parser. + + Args: + on_thought_chunk: Callback for thought text chunks as they arrive + on_answer_chunk: Callback for final answer text chunks as they arrive + """ + self.on_thought_chunk = on_thought_chunk + self.on_answer_chunk = on_answer_chunk + + # Parser state + self.state = ParserState.INITIAL + + # Buffers for accumulating content + self.line_buffer = "" # For detecting delimiters across chunk boundaries + self.thought_buffer = "" # Accumulated thought text + self.action_buffer = "" # Action name + self.args_buffer = "" # JSON arguments text + self.answer_buffer = "" # Final answer text + + # JSON parsing state for Args + self.brace_count = 0 + self.args_started = False + + # Result object (Action or Final) + self.result = None + + def feed(self, chunk: str) -> None: + """ + Feed a text chunk to the parser. + + Args: + chunk: Text chunk from LLM stream + """ + if self.state == ParserState.COMPLETE: + return # Already complete, ignore further chunks + + # Add chunk to line buffer for delimiter detection + self.line_buffer += chunk + + # Remove markdown code blocks if present + self.line_buffer = re.sub(r'^```[^\n]*\n', '', self.line_buffer) + self.line_buffer = re.sub(r'\n```$', '', self.line_buffer) + + # Process based on current state + while self.line_buffer and self.state != ParserState.COMPLETE: + if self.state == ParserState.INITIAL: + self._process_initial() + elif self.state == ParserState.THOUGHT: + self._process_thought() + elif self.state == ParserState.ACTION: + self._process_action() + elif self.state == ParserState.ARGS: + self._process_args() + elif self.state == ParserState.FINAL_ANSWER: + self._process_final_answer() + + def _process_initial(self) -> None: + """Process INITIAL state - looking for 'Thought:' delimiter""" + idx = self.line_buffer.find(self.THOUGHT_DELIMITER) + + if idx >= 0: + # Found thought delimiter + # Discard any content before it + self.line_buffer = self.line_buffer[idx + len(self.THOUGHT_DELIMITER):] + self.state = ParserState.THOUGHT + elif len(self.line_buffer) >= self.MAX_DELIMITER_BUFFER: + # Buffer getting too large, probably junk before thought + # Keep only the tail that might contain partial delimiter + self.line_buffer = self.line_buffer[-self.MAX_DELIMITER_BUFFER:] + + def _process_thought(self) -> None: + """Process THOUGHT state - accumulating thought content""" + # Check for Action or Final Answer delimiter + action_idx = self.line_buffer.find(self.ACTION_DELIMITER) + final_idx = self.line_buffer.find(self.FINAL_ANSWER_DELIMITER) + + # Find which delimiter comes first (if any) + next_delimiter_idx = -1 + next_state = None + + if action_idx >= 0 and (final_idx < 0 or action_idx < final_idx): + next_delimiter_idx = action_idx + next_state = ParserState.ACTION + delimiter_len = len(self.ACTION_DELIMITER) + elif final_idx >= 0: + next_delimiter_idx = final_idx + next_state = ParserState.FINAL_ANSWER + delimiter_len = len(self.FINAL_ANSWER_DELIMITER) + + if next_delimiter_idx >= 0: + # Found next delimiter + thought_chunk = self.line_buffer[:next_delimiter_idx].strip() + if thought_chunk: + self.thought_buffer += thought_chunk + if self.on_thought_chunk: + self.on_thought_chunk(thought_chunk) + + self.line_buffer = self.line_buffer[next_delimiter_idx + delimiter_len:] + self.state = next_state + else: + # No delimiter found yet + # Keep tail in buffer (might contain partial delimiter) + # Emit the rest as thought chunk + if len(self.line_buffer) > self.MAX_DELIMITER_BUFFER: + emittable = self.line_buffer[:-self.MAX_DELIMITER_BUFFER] + self.thought_buffer += emittable + if self.on_thought_chunk: + self.on_thought_chunk(emittable) + self.line_buffer = self.line_buffer[-self.MAX_DELIMITER_BUFFER:] + + def _process_action(self) -> None: + """Process ACTION state - collecting action name""" + # Action name is on one line (or at least until newline or Args:) + newline_idx = self.line_buffer.find('\n') + args_idx = self.line_buffer.find(self.ARGS_DELIMITER) + + # Find which comes first + if args_idx >= 0 and (newline_idx < 0 or args_idx < newline_idx): + # Args delimiter found first + self.action_buffer = self.line_buffer[:args_idx].strip().strip('"') + self.line_buffer = self.line_buffer[args_idx + len(self.ARGS_DELIMITER):] + self.state = ParserState.ARGS + elif newline_idx >= 0: + # Newline found, action name complete + self.action_buffer = self.line_buffer[:newline_idx].strip().strip('"') + self.line_buffer = self.line_buffer[newline_idx + 1:] + # Stay in ACTION state or move to ARGS if we find delimiter + # Actually, check if next line has Args: + if self.line_buffer.lstrip().startswith(self.ARGS_DELIMITER): + args_start = self.line_buffer.find(self.ARGS_DELIMITER) + self.line_buffer = self.line_buffer[args_start + len(self.ARGS_DELIMITER):] + self.state = ParserState.ARGS + else: + # Not enough content yet, keep buffering + # But if buffer is getting large, action name is probably complete + if len(self.line_buffer) > 100: + self.action_buffer = self.line_buffer.strip().strip('"') + self.line_buffer = "" + # Assume Args comes next, but we need more content + self.state = ParserState.ARGS + + def _process_args(self) -> None: + """Process ARGS state - collecting JSON arguments""" + # Process character by character to track brace matching + i = 0 + while i < len(self.line_buffer): + char = self.line_buffer[i] + self.args_buffer += char + + if char == '{': + self.brace_count += 1 + self.args_started = True + elif char == '}': + self.brace_count -= 1 + + # Check if JSON is complete + if self.args_started and self.brace_count == 0: + # JSON complete, try to parse + try: + args_dict = json.loads(self.args_buffer.strip()) + # Success! Create Action result + self.result = Action( + thought=self.thought_buffer.strip(), + name=self.action_buffer, + arguments=args_dict, + observation="" + ) + self.state = ParserState.COMPLETE + self.line_buffer = "" # Clear buffer + return + except json.JSONDecodeError as e: + logger.error(f"Failed to parse JSON args: {self.args_buffer}") + raise ValueError(f"Invalid JSON in Args: {e}") + + i += 1 + + # Consumed entire buffer, clear it and wait for more chunks + self.line_buffer = "" + + def _process_final_answer(self) -> None: + """Process FINAL_ANSWER state - collecting final answer""" + # For final answer, we consume everything until we decide we're done + # In streaming mode, we can't know when answer is complete until stream ends + # So we emit chunks and accumulate + + # Check if this might be JSON + is_json = self.answer_buffer.strip().startswith('{') or \ + self.line_buffer.strip().startswith('{') + + if is_json: + # Handle JSON final answer + self.answer_buffer += self.line_buffer + + # Count braces to detect completion + brace_count = self.answer_buffer.count('{') - self.answer_buffer.count('}') + + if brace_count == 0 and '{' in self.answer_buffer: + # JSON might be complete + # Note: We can't be 100% sure without trying to parse + # But in streaming mode, we'll finish when stream ends + pass + + # Emit chunk + if self.on_answer_chunk: + self.on_answer_chunk(self.line_buffer) + + self.line_buffer = "" + else: + # Regular text answer - emit everything + if self.line_buffer: + self.answer_buffer += self.line_buffer + if self.on_answer_chunk: + self.on_answer_chunk(self.line_buffer) + self.line_buffer = "" + + def finalize(self) -> None: + """ + Call this when the stream is complete to finalize parsing. + This handles any remaining buffered content. + """ + if self.state == ParserState.COMPLETE: + return + + # Flush any remaining thought chunks + if self.state == ParserState.THOUGHT and self.line_buffer: + self.thought_buffer += self.line_buffer + if self.on_thought_chunk: + self.on_thought_chunk(self.line_buffer) + self.line_buffer = "" + + # Finalize final answer + if self.state == ParserState.FINAL_ANSWER: + # Flush any remaining answer content + if self.line_buffer: + self.answer_buffer += self.line_buffer + if self.on_answer_chunk: + self.on_answer_chunk(self.line_buffer) + self.line_buffer = "" + + # Create Final result + self.result = Final( + thought=self.thought_buffer.strip(), + final=self.answer_buffer.strip() + ) + self.state = ParserState.COMPLETE + + # If we're in other states, something went wrong + if self.state not in [ParserState.COMPLETE, ParserState.FINAL_ANSWER]: + if self.thought_buffer: + raise ValueError( + f"Stream ended in {self.state.value} state with incomplete parsing. " + f"Thought: {self.thought_buffer[:100]}..." + ) + else: + raise ValueError(f"Stream ended in {self.state.value} state with no content") + + def is_complete(self) -> bool: + """Check if parsing is complete""" + return self.state == ParserState.COMPLETE + + def get_result(self) -> Optional[Action | Final]: + """Get the parsed result (Action or Final)""" + return self.result