"""Integration tests: stream Agent output through StreamLLMDemoStep. Requires LLM_API_KEY (and optionally LLM_BASE_URL / LLM_MODEL_NAME) in the environment or a .env file at the repo root. Hits the real LLM API. """ import asyncio import sys from pathlib import Path INTEGRATION_DIR = Path(__file__).resolve().parent sys.path.insert(0, str(INTEGRATION_DIR)) # pylint: disable=wrong-import-position from _vault_fixture import vault_env # noqa: E402 from reme4.enumeration import ChunkEnum # noqa: E402 from reme4.schema import StreamChunk # noqa: E402 from reme4.steps.common.stream_llm_demo import StreamLLMDemoStep # noqa: E402 from reme4.utils.common_utils import execute_stream_task # noqa: E402 async def _test_stream_llm_basic_chat(): """StreamLLMDemoStep streams text chunks via add_stream_string.""" with vault_env() as env: app = await env.make_app() try: step = StreamLLMDemoStep(app_context=app.context) queue: asyncio.Queue = asyncio.Queue() chunks: list[StreamChunk] = [] task = asyncio.create_task( step( stream_queue=queue, query="Explain step by step how to compute 1 + 1, and give the final answer.", ), ) print("\n[stream_basic] streaming output:") async for raw in execute_stream_task(queue, task, output_format="chunk"): chunk: StreamChunk = raw # type: ignore[assignment] chunks.append(chunk) if chunk.chunk_type == ChunkEnum.CONTENT: sys.stdout.write(chunk.chunk) sys.stdout.flush() response = task.result() # Should have received multiple CONTENT chunks for a longer response content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] print(f"\n\n[stream_basic] got {len(content_chunks)} CONTENT chunks") assert len(content_chunks) > 1, "Expected multiple CONTENT chunks for streaming" # Final answer should be populated text = (response.answer or "").strip() assert text, "Empty assistant response" assert "2" in text, f"Expected '2' in response, got: {text!r}" # Concatenated stream text should match the final answer streamed_text = "".join(c.chunk for c in content_chunks) assert streamed_text.strip() == text, f"Stream text mismatch: {streamed_text!r} vs {text!r}" print("✓ test_stream_llm_basic_chat passed") finally: await env.close_all() async def _test_stream_llm_with_tool(): """StreamLLMDemoStep streams tool call events when tools are used.""" with vault_env() as env: app = await env.make_app() try: step = StreamLLMDemoStep(app_context=app.context) queue: asyncio.Queue = asyncio.Queue() chunks: list[StreamChunk] = [] task = asyncio.create_task( step( stream_queue=queue, query="Use the add tool to compute 21 + 21 and report the result.", sys_prompt="Use the `add` tool whenever the user asks to add numbers.", ), ) print("\n[stream_tool] streaming output:") async for raw in execute_stream_task(queue, task, output_format="chunk"): chunk: StreamChunk = raw # type: ignore[assignment] chunks.append(chunk) if chunk.chunk_type == ChunkEnum.CONTENT: sys.stdout.write(chunk.chunk) sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.TOOL_CALL: sys.stdout.write(f"\033[33m{chunk.chunk}\033[0m") sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.TOOL_RESULT: sys.stdout.write(f"\033[32m{chunk.chunk}\033[0m") sys.stdout.flush() response = task.result() tool_call_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.TOOL_CALL] tool_result_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.TOOL_RESULT] content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] print(f"\n\n[stream_tool] TOOL_CALL chunks: {len(tool_call_chunks)}") print(f"[stream_tool] TOOL_RESULT chunks: {len(tool_result_chunks)}") print(f"[stream_tool] CONTENT chunks: {len(content_chunks)}") assert len(tool_call_chunks) > 0, "Expected TOOL_CALL chunks" assert len(tool_result_chunks) > 0, "Expected TOOL_RESULT chunks" text = (response.answer or "").strip() print(f"[stream_tool] final answer: {text!r}") assert "42" in text, f"Expected '42' in response, got: {text!r}" print("✓ test_stream_llm_with_tool passed") finally: await env.close_all() async def _test_stream_llm_fallback_no_stream(): """Without stream_queue, still uses streaming under the hood for real-time output.""" with vault_env() as env: app = await env.make_app() try: step = StreamLLMDemoStep(app_context=app.context) queue: asyncio.Queue = asyncio.Queue() chunks: list[StreamChunk] = [] task = asyncio.create_task( step( stream_queue=queue, query="Explain step by step how to compute 1 + 1, and give the final answer.", ), ) print("\n[fallback_stream] streaming output:") async for raw in execute_stream_task(queue, task, output_format="chunk"): chunk: StreamChunk = raw # type: ignore[assignment] chunks.append(chunk) if chunk.chunk_type == ChunkEnum.CONTENT: sys.stdout.write(chunk.chunk) sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.THINK: sys.stdout.write(f"\033[2m{chunk.chunk}\033[0m") sys.stdout.flush() response = task.result() text = (response.answer or "").strip() content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] print(f"\n\n[fallback_stream] got {len(content_chunks)} CONTENT chunks") assert text, "Empty assistant response" assert "2" in text, f"Expected '2' in response, got: {text!r}" print("✓ test_stream_llm_fallback_no_stream passed") finally: await env.close_all() def test_stream_llm_basic_chat(): """StreamLLMDemoStep streams text chunks via add_stream_string.""" asyncio.run(_test_stream_llm_basic_chat()) def test_stream_llm_with_tool(): """StreamLLMDemoStep streams tool call events when tools are used.""" asyncio.run(_test_stream_llm_with_tool()) def test_stream_llm_fallback_no_stream(): """Without stream_queue, falls back to non-streaming reply.""" asyncio.run(_test_stream_llm_fallback_no_stream()) async def _demo_stream_print(): """Real-time streaming print demo — ask a longer question to see chunked output.""" with vault_env() as env: app = await env.make_app() try: step = StreamLLMDemoStep(app_context=app.context) queue: asyncio.Queue = asyncio.Queue() query = ( "Please explain in detail how neural networks learn through backpropagation. " "Include the chain rule, gradient descent, and give a concrete example with numbers." ) task = asyncio.create_task( step( stream_queue=queue, query=query, sys_prompt="You are a knowledgeable AI teacher. Explain concepts thoroughly.", ), ) async for raw in execute_stream_task(queue, task, output_format="chunk"): chunk: StreamChunk = raw # type: ignore[assignment] if chunk.chunk_type == ChunkEnum.CONTENT: sys.stdout.write(chunk.chunk) sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.THINK: sys.stdout.write(f"\033[2m{chunk.chunk}\033[0m") sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.TOOL_CALL: sys.stdout.write(f"\n\033[33m[tool_call] {chunk.chunk}\033[0m") sys.stdout.flush() elif chunk.chunk_type == ChunkEnum.TOOL_RESULT: sys.stdout.write(f"\033[32m{chunk.chunk}\033[0m") sys.stdout.flush() print() finally: await env.close_all() async def _run_all(): print("=== StreamLLMDemoStep integration tests ===") await _test_stream_llm_basic_chat() await _test_stream_llm_with_tool() await _test_stream_llm_fallback_no_stream() print("\nAll stream integration tests passed!") if __name__ == "__main__": if len(sys.argv) > 1 and sys.argv[1] == "demo": asyncio.run(_demo_stream_print()) else: asyncio.run(_run_all())