From 5a2f04a8aaaee6891fada01812fdff9cab65f310 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 26 Nov 2025 19:40:03 +0000 Subject: [PATCH] Fixed the tests --- .../test_agent_streaming_integration.py | 10 ++-- .../test_prompt_streaming_integration.py | 7 +-- ...t_text_completion_streaming_integration.py | 47 ++++++++++--------- .../agent/react/streaming_parser.py | 8 +++- 4 files changed, 39 insertions(+), 33 deletions(-) diff --git a/tests/integration/test_agent_streaming_integration.py b/tests/integration/test_agent_streaming_integration.py index ea7c890d..2b619098 100644 --- a/tests/integration/test_agent_streaming_integration.py +++ b/tests/integration/test_agent_streaming_integration.py @@ -36,13 +36,13 @@ Args: { }""" if streaming and chunk_callback: - # Send chunks that avoid parser state bug - # Parser bug: if "Args:" starts a new chunk while in ACTION state, - # it overwrites action_buffer with empty string (line_buffer[:0]) - # Solution: Keep "Action:" line and "Args:" together in same chunk + # Send realistic line-by-line chunks + # This tests that the parser properly handles "Args:" starting a new chunk + # (which previously caused a bug where action_buffer was overwritten) chunks = [ "Thought: I need to search for information about machine learning.\n", - "Action: knowledge_query\nArgs: {\n", # Keep Action and Args together! + "Action: knowledge_query\n", + "Args: {\n", # This used to trigger bug - Args: at start of chunk ' "question": "What is machine learning?"\n', "}" ] diff --git a/tests/integration/test_prompt_streaming_integration.py b/tests/integration/test_prompt_streaming_integration.py index 65b3d68a..1aef42bf 100644 --- a/tests/integration/test_prompt_streaming_integration.py +++ b/tests/integration/test_prompt_streaming_integration.py @@ -212,6 +212,7 @@ class TestPromptStreaming: async def test_prompt_streaming_error_handling(self, prompt_processor_streaming): """Test error handling during streaming""" # Arrange + from trustgraph.schema import Error context = MagicMock() # Mock text completion client that raises an error @@ -219,10 +220,10 @@ class TestPromptStreaming: async def failing_request(request, recipient=None, timeout=600): if recipient: - # Send error response + # Send error response with proper Error schema error_response = TextCompletionResponse( response="", - error=MagicMock(message="Text completion error"), + error=Error(message="Text completion error", type="processing_error"), end_of_stream=True ) await recipient(error_response) @@ -386,4 +387,4 @@ class TestPromptStreaming: # Verify chunks concatenate to expected result full_text = "".join(chunk_texts) - assert full_text == "Machine learning is a field of artificial intelligence." + assert full_text == "Machine learning is a field of artificial intelligence" diff --git a/tests/integration/test_text_completion_streaming_integration.py b/tests/integration/test_text_completion_streaming_integration.py index 7c9ff18a..a70afb4c 100644 --- a/tests/integration/test_text_completion_streaming_integration.py +++ b/tests/integration/test_text_completion_streaming_integration.py @@ -29,30 +29,31 @@ class TestTextCompletionStreaming: def create_streaming_completion(**kwargs): """Generator that yields streaming chunks""" - if kwargs.get('stream', False): - # Simulate OpenAI streaming response - chunks_text = [ - "Machine", " learning", " is", " a", " subset", - " of", " AI", " that", " enables", " computers", - " to", " learn", " from", " data", "." - ] - - for text in chunks_text: - delta = ChoiceDelta(content=text, role=None) - choice = StreamChoice(index=0, delta=delta, finish_reason=None) - chunk = ChatCompletionChunk( - id="chatcmpl-streaming", - choices=[choice], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion.chunk" - ) - yield chunk - else: - # Non-streaming response (shouldn't happen in these tests) + # Check if streaming is enabled + if not kwargs.get('stream', False): raise ValueError("Expected streaming mode") - client.chat.completions.create.return_value = create_streaming_completion() + # Simulate OpenAI streaming response + chunks_text = [ + "Machine", " learning", " is", " a", " subset", + " of", " AI", " that", " enables", " computers", + " to", " learn", " from", " data", "." + ] + + for text in chunks_text: + delta = ChoiceDelta(content=text, role=None) + choice = StreamChoice(index=0, delta=delta, finish_reason=None) + chunk = ChatCompletionChunk( + id="chatcmpl-streaming", + choices=[choice], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + yield chunk + + # Return a new generator each time create is called + client.chat.completions.create.side_effect = lambda **kwargs: create_streaming_completion(**kwargs) return client @pytest.fixture @@ -284,7 +285,7 @@ class TestTextCompletionStreaming: ) yield chunk - mock_streaming_openai_client.chat.completions.create.return_value = create_streaming_with_empties() + mock_streaming_openai_client.chat.completions.create.side_effect = lambda **kwargs: create_streaming_with_empties(**kwargs) processor = MagicMock() processor.default_model = "gpt-3.5-turbo" diff --git a/trustgraph-flow/trustgraph/agent/react/streaming_parser.py b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py index b5f87186..1cdada11 100644 --- a/trustgraph-flow/trustgraph/agent/react/streaming_parser.py +++ b/trustgraph-flow/trustgraph/agent/react/streaming_parser.py @@ -202,12 +202,16 @@ class StreamingReActParser: # 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('"') + # Only set action_buffer if not already set (to avoid overwriting with empty string) + if not self.action_buffer: + self.action_buffer = self.line_buffer[:args_idx].strip().strip('"') self.line_buffer = self.line_buffer[args_idx + len(self.ARGS_DELIMITER):].lstrip() self.state = ParserState.ARGS elif newline_idx >= 0: # Newline found, action name complete - self.action_buffer = self.line_buffer[:newline_idx].strip().strip('"') + # Only set action_buffer if not already set + if not self.action_buffer: + 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: