diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index dd3b1e88..274d4834 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -1,15 +1,20 @@ import json import asyncio +import logging from . request_response_spec import RequestResponse, RequestResponseSpec from .. schema import PromptRequest, PromptResponse +logger = logging.getLogger(__name__) + class PromptClient(RequestResponse): async def prompt(self, id, variables, timeout=600, streaming=False, chunk_callback=None): + logger.info(f"DEBUG prompt_client: prompt called, id={id}, streaming={streaming}, chunk_callback={chunk_callback is not None}") if not streaming: + logger.info("DEBUG prompt_client: Non-streaming path") # Non-streaming path resp = await self.request( PromptRequest( @@ -31,44 +36,59 @@ class PromptClient(RequestResponse): return json.loads(resp.object) else: + logger.info("DEBUG prompt_client: Streaming path") # Streaming path - collect all chunks full_text = "" full_object = None async def collect_chunks(resp): nonlocal full_text, full_object + logger.info(f"DEBUG prompt_client: collect_chunks called, resp.text={resp.text[:50] if resp.text else None}, end_of_stream={getattr(resp, 'end_of_stream', False)}") if resp.error: + logger.error(f"DEBUG prompt_client: Error in response: {resp.error.message}") raise RuntimeError(resp.error.message) if resp.text: full_text += resp.text + logger.info(f"DEBUG prompt_client: Accumulated {len(full_text)} chars") # Call chunk callback if provided if chunk_callback: + logger.info(f"DEBUG prompt_client: Calling chunk_callback") if asyncio.iscoroutinefunction(chunk_callback): await chunk_callback(resp.text) else: chunk_callback(resp.text) elif resp.object: + logger.info(f"DEBUG prompt_client: Got object response") full_object = resp.object - return getattr(resp, 'end_of_stream', False) + end_stream = getattr(resp, 'end_of_stream', False) + logger.info(f"DEBUG prompt_client: Returning end_of_stream={end_stream}") + return end_stream + logger.info("DEBUG prompt_client: Creating PromptRequest") + req = PromptRequest( + id = id, + terms = { + k: json.dumps(v) + for k, v in variables.items() + }, + streaming = True + ) + logger.info(f"DEBUG prompt_client: About to call self.request with recipient, timeout={timeout}") await self.request( - PromptRequest( - id = id, - terms = { - k: json.dumps(v) - for k, v in variables.items() - }, - streaming = True - ), + req, recipient=collect_chunks, timeout=timeout ) + logger.info(f"DEBUG prompt_client: self.request returned, full_text has {len(full_text)} chars") - if full_text: return full_text + if full_text: + logger.info("DEBUG prompt_client: Returning full_text") + return full_text + logger.info("DEBUG prompt_client: Returning parsed full_object") return json.loads(full_object) async def extract_definitions(self, text, timeout=600): diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py index 03d5d379..9311a57b 100644 --- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py +++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py @@ -220,32 +220,72 @@ class AgentManager: logger.info(f"prompt: {variables}") + logger.info(f"DEBUG: streaming={streaming}, think={think is not None}") + # Streaming path - use StreamingReActParser if streaming and think: + logger.info("DEBUG: Entering streaming path") from .streaming_parser import StreamingReActParser - # Create parser with streaming callbacks - # Thought chunks go to think(), answer chunks go to answer() + logger.info("DEBUG: Creating StreamingReActParser") + + # Collect chunks to send via async callbacks + thought_chunks = [] + answer_chunks = [] + + # Create parser with synchronous callbacks that just collect chunks 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)), + on_thought_chunk=lambda chunk: thought_chunks.append(chunk), + on_answer_chunk=lambda chunk: answer_chunks.append(chunk), ) + logger.info("DEBUG: StreamingReActParser created") - # Create async chunk callback that feeds parser + # Create async chunk callback that feeds parser and sends collected chunks async def on_chunk(text): - parser.feed(text) + logger.info(f"DEBUG: on_chunk called with {len(text)} chars") + # Track what we had before + prev_thought_count = len(thought_chunks) + prev_answer_count = len(answer_chunks) + + # Feed the parser (synchronous) + logger.info(f"DEBUG: About to call parser.feed") + parser.feed(text) + logger.info(f"DEBUG: parser.feed returned") + + # Send any new thought chunks + for i in range(prev_thought_count, len(thought_chunks)): + logger.info(f"DEBUG: Sending thought chunk {i}") + await think(thought_chunks[i]) + + # Send any new answer chunks + for i in range(prev_answer_count, len(answer_chunks)): + logger.info(f"DEBUG: Sending answer chunk {i}") + if answer: + await answer(answer_chunks[i]) + else: + await think(answer_chunks[i]) + + logger.info("DEBUG: Getting prompt-request client from context") + client = context("prompt-request") + logger.info(f"DEBUG: Got client: {client}") + + logger.info("DEBUG: About to call agent_react with streaming=True") # Get streaming response - response_text = await context("prompt-request").agent_react( + response_text = await client.agent_react( variables=variables, streaming=True, chunk_callback=on_chunk ) + logger.info(f"DEBUG: agent_react returned, got {len(response_text) if response_text else 0} chars") # Finalize parser + logger.info("DEBUG: Finalizing parser") parser.finalize() + logger.info("DEBUG: Parser finalized") # Get result + logger.info("DEBUG: Getting result from parser") result = parser.get_result() if result is None: raise RuntimeError("Parser failed to produce a result") @@ -254,11 +294,18 @@ class AgentManager: return result else: + logger.info("DEBUG: Entering NON-streaming path") # Non-streaming path - get complete text and parse - response_text = await context("prompt-request").agent_react( + logger.info("DEBUG: Getting prompt-request client from context") + client = context("prompt-request") + logger.info(f"DEBUG: Got client: {client}") + + logger.info("DEBUG: About to call agent_react with streaming=False") + response_text = await client.agent_react( variables=variables, streaming=False ) + logger.info(f"DEBUG: agent_react returned, got response") logger.debug(f"Response text:\n{response_text}") diff --git a/trustgraph-flow/trustgraph/agent/react/streaming_parser.py b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py index 76192e92..3f1d856d 100644 --- a/trustgraph-flow/trustgraph/agent/react/streaming_parser.py +++ b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py @@ -118,7 +118,11 @@ class StreamingReActParser: self.line_buffer = re.sub(r'\n```$', '', self.line_buffer) # Process based on current state + # Track previous state to detect if we're making progress while self.line_buffer and self.state != ParserState.COMPLETE: + prev_buffer_len = len(self.line_buffer) + prev_state = self.state + if self.state == ParserState.INITIAL: self._process_initial() elif self.state == ParserState.THOUGHT: @@ -130,6 +134,11 @@ class StreamingReActParser: elif self.state == ParserState.FINAL_ANSWER: self._process_final_answer() + # If no progress was made (buffer unchanged AND state unchanged), break + # to avoid infinite loop. We'll process more when the next chunk arrives. + if len(self.line_buffer) == prev_buffer_len and self.state == prev_state: + break + def _process_initial(self) -> None: """Process INITIAL state - looking for 'Thought:' delimiter""" idx = self.line_buffer.find(self.THOUGHT_DELIMITER)