Fixed the tests

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',
                     "}"
                 ]

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"

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"

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: