diff --git a/tests/agui/__init__.py b/tests/agui/__init__.py new file mode 100644 index 00000000..d178a31a --- /dev/null +++ b/tests/agui/__init__.py @@ -0,0 +1 @@ +"""Tests for AG-UI implementation.""" diff --git a/tests/agui/test_cli_renderer.py b/tests/agui/test_cli_renderer.py new file mode 100644 index 00000000..c01d6934 --- /dev/null +++ b/tests/agui/test_cli_renderer.py @@ -0,0 +1,158 @@ +"""Tests for AGUIConsoleRenderer.""" + +import pytest +from pydantic import BaseModel +from rich.console import Console + +from haiku.rag.agui.cli_renderer import AGUIConsoleRenderer +from haiku.rag.agui.events import ( + emit_activity, + emit_run_error, + emit_run_finished, + emit_run_started, + emit_state_snapshot, + emit_step_finished, + emit_step_started, + emit_text_message, +) + + +class SimpleState(BaseModel): + """Simple state for testing.""" + + value: int + + +async def async_gen(items): + """Helper to create async generator from list.""" + for item in items: + yield item + + +@pytest.mark.asyncio +async def test_renderer_basic_flow(): + """Test basic event rendering flow.""" + console = Console(file=None, force_terminal=False) # Don't actually print + renderer = AGUIConsoleRenderer(console) + + events = [ + emit_run_started("t1", "r1"), + emit_state_snapshot(SimpleState(value=1)), + emit_step_started("plan"), + emit_activity("m1", "planning", "Planning research"), + emit_run_finished("t1", "r1", {"status": "complete"}), + ] + + result = await renderer.render(async_gen(events)) + + assert result == {"status": "complete"} + + +@pytest.mark.asyncio +async def test_renderer_state_diff(): + """Test that renderer computes state diffs correctly.""" + renderer = AGUIConsoleRenderer() + + # First state + old_state = {"value": 1, "text": "old", "nested": {"a": 1, "b": 2}} + # Second state with changes + new_state = {"value": 2, "text": "old", "nested": {"a": 1, "b": 3, "c": 4}} + + diff = renderer._compute_diff(old_state, new_state) + + assert diff == { + "value": 2, + "nested": {"b": 3, "c": 4}, # Only changed/new fields in nested + } + + +@pytest.mark.asyncio +async def test_renderer_state_diff_no_changes(): + """Test that no diff is computed when states are equal.""" + renderer = AGUIConsoleRenderer() + + state = {"value": 1, "text": "test"} + diff = renderer._compute_diff(state, state) + + assert diff == {} + + +@pytest.mark.asyncio +async def test_renderer_handles_all_event_types(): + """Test that renderer handles all event types without errors.""" + renderer = AGUIConsoleRenderer() + + events = [ + emit_run_started("t1", "r1"), + emit_state_snapshot(SimpleState(value=1)), + emit_step_started("step1"), + emit_step_finished("step1"), + emit_text_message("message"), + emit_activity("m1", "work", "Working"), + emit_run_error("error occurred"), + emit_run_finished("t1", "r1", {"result": "done"}), + ] + + # Should not raise any exceptions + result = await renderer.render(async_gen(events)) + assert result == {"result": "done"} + + +@pytest.mark.asyncio +async def test_renderer_initial_state(): + """Test that initial state is rendered.""" + renderer = AGUIConsoleRenderer() + + events = [ + emit_state_snapshot(SimpleState(value=1)), + emit_state_snapshot(SimpleState(value=2)), + ] + + await renderer.render(async_gen(events)) + + # After processing, internal state should be the last state + assert renderer._state == {"value": 2} + + +@pytest.mark.asyncio +async def test_renderer_no_result(): + """Test renderer when no RUN_FINISHED event.""" + renderer = AGUIConsoleRenderer() + + events = [ + emit_run_started("t1", "r1"), + emit_step_started("step1"), + ] + + result = await renderer.render(async_gen(events)) + assert result is None + + +@pytest.mark.asyncio +async def test_renderer_nested_state_diff(): + """Test nested state diff computation.""" + renderer = AGUIConsoleRenderer() + + old = {"level1": {"level2": {"value": 1, "text": "old"}, "other": "same"}} + + new = {"level1": {"level2": {"value": 2, "text": "old"}, "other": "same"}} + + diff = renderer._compute_diff(old, new) + + # Should only show the changed nested value + assert diff == {"level1": {"level2": {"value": 2}}} + + +@pytest.mark.asyncio +async def test_renderer_with_empty_state(): + """Test renderer handles empty state gracefully.""" + renderer = AGUIConsoleRenderer() + + # Create a state with no changes + events = [ + emit_step_started("step1"), # No state snapshot + emit_run_finished("t1", "r1", {"result": "ok"}), + ] + + result = await renderer.render(async_gen(events)) + assert result == {"result": "ok"} diff --git a/tests/agui/test_emitter.py b/tests/agui/test_emitter.py new file mode 100644 index 00000000..5a057247 --- /dev/null +++ b/tests/agui/test_emitter.py @@ -0,0 +1,231 @@ +"""Tests for AGUIEmitter.""" + +import asyncio + +import pytest +from pydantic import BaseModel + +from haiku.rag.agui.emitter import AGUIEmitter + + +class TestState(BaseModel): + """Test state model.""" + + value: int + text: str + + +class TestResult(BaseModel): + """Test result model.""" + + status: str + + +@pytest.mark.asyncio +async def test_emitter_lifecycle(): + """Test emitter lifecycle events.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + initial_state = TestState(value=1, text="initial") + emitter.start_run(initial_state) + + events = [] + async for event in emitter: + events.append(event) + if event["type"] == "RUN_STARTED": + # Close after getting run started + result = TestResult(status="complete") + emitter.finish_run(result) + await emitter.close() + + assert len(events) >= 3 # RUN_STARTED, STATE_SNAPSHOT, RUN_FINISHED + assert events[0]["type"] == "RUN_STARTED" + assert events[-1]["type"] == "RUN_FINISHED" + assert events[-1]["result"] == {"status": "complete"} + + +@pytest.mark.asyncio +async def test_emitter_step_events(): + """Test step lifecycle events.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + emitter.start_step("test_step") + emitter.finish_step() + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + step_events = [e for e in events if e["type"] in ("STEP_STARTED", "STEP_FINISHED")] + assert len(step_events) == 2 + assert step_events[0]["type"] == "STEP_STARTED" + assert step_events[0]["stepName"] == "test_step" + assert step_events[1]["type"] == "STEP_FINISHED" + assert step_events[1]["stepName"] == "test_step" + + +@pytest.mark.asyncio +async def test_emitter_state_updates(): + """Test state update events.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + state1 = TestState(value=1, text="first") + state2 = TestState(value=2, text="second") + + emitter.update_state(state1) + emitter.update_state(state2) + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + state_events = [e for e in events if e["type"] == "STATE_SNAPSHOT"] + assert len(state_events) == 2 + assert state_events[0]["snapshot"] == {"value": 1, "text": "first"} + assert state_events[1]["snapshot"] == {"value": 2, "text": "second"} + + +@pytest.mark.asyncio +async def test_emitter_activity_events(): + """Test activity events.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + emitter.update_activity("processing", "Processing data") + emitter.update_activity("done", "Completed", message_id="msg-1") + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + activity_events = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"] + assert len(activity_events) == 2 + assert activity_events[0]["activityType"] == "processing" + assert activity_events[0]["content"] == "Processing data" + assert activity_events[1]["messageId"] == "msg-1" + assert activity_events[1]["activityType"] == "done" + + +@pytest.mark.asyncio +async def test_emitter_text_messages(): + """Test text message events.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + emitter.log("Test message", role="assistant") + emitter.log("Another message", role="user") + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"] + assert len(text_events) == 2 + assert text_events[0]["delta"] == "Test message" + assert text_events[0]["role"] == "assistant" + assert text_events[1]["delta"] == "Another message" + assert text_events[1]["role"] == "user" + + +@pytest.mark.asyncio +async def test_emitter_error(): + """Test error event emission.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + error = ValueError("Test error") + emitter.error(error, code="TEST_ERROR") + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + error_events = [e for e in events if e["type"] == "RUN_ERROR"] + assert len(error_events) == 1 + assert error_events[0]["message"] == "Test error" + assert error_events[0]["code"] == "TEST_ERROR" + + +@pytest.mark.asyncio +async def test_emitter_thread_and_run_ids(): + """Test thread and run ID management.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter( + thread_id="thread-1", run_id="run-1" + ) + + assert emitter.thread_id == "thread-1" + assert emitter.run_id == "run-1" + + initial_state = TestState(value=1, text="test") + emitter.start_run(initial_state) + await emitter.close() + + events = [] + async for event in emitter: + events.append(event) + + run_started = [e for e in events if e["type"] == "RUN_STARTED"][0] + assert run_started["threadId"] == "thread-1" + assert run_started["runId"] == "run-1" + + +@pytest.mark.asyncio +async def test_emitter_generates_thread_id(): + """Test that thread ID is generated from state hash when not provided.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + initial_state = TestState(value=42, text="test") + emitter.start_run(initial_state) + + # Thread ID should be generated deterministically from state + assert emitter.thread_id is not None + assert len(emitter.thread_id) > 0 + + await emitter.close() + async for _ in emitter: + pass + + +@pytest.mark.asyncio +async def test_emitter_closes_properly(): + """Test that emitter closes and stops iteration.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + emitter.log("Message 1") + await emitter.close() + + # Attempting to iterate after close should work and stop + events = [] + async for event in emitter: + events.append(event) + + # Should have received the message and then stopped + assert len(events) == 1 + + +@pytest.mark.asyncio +async def test_emitter_concurrent_emission(): + """Test that multiple events can be emitted concurrently.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + async def emit_many(): + for i in range(10): + emitter.log(f"Message {i}") + await asyncio.sleep(0.001) # Simulate some work + await emitter.close() + + # Start emission in background + emit_task = asyncio.create_task(emit_many()) + + # Collect events + events = [] + async for event in emitter: + events.append(event) + + await emit_task + + # Should have all 10 messages + text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"] + assert len(text_events) == 10 diff --git a/tests/agui/test_events.py b/tests/agui/test_events.py new file mode 100644 index 00000000..67db2e22 --- /dev/null +++ b/tests/agui/test_events.py @@ -0,0 +1,155 @@ +"""Tests for AG-UI event creation utilities.""" + +from pydantic import BaseModel + +from haiku.rag.agui.events import ( + emit_activity, + emit_run_error, + emit_run_finished, + emit_run_started, + emit_state_snapshot, + emit_step_finished, + emit_step_started, + emit_text_message, +) + + +class TestState(BaseModel): + """Test state model.""" + + value: int + + +class TestResult(BaseModel): + """Test result model.""" + + status: str + + +def test_emit_run_started(): + """Test RUN_STARTED event creation.""" + event = emit_run_started("thread-1", "run-1") + + assert event["type"] == "RUN_STARTED" + assert event["threadId"] == "thread-1" + assert event["runId"] == "run-1" + assert "input" not in event + + +def test_emit_run_started_with_input(): + """Test RUN_STARTED event with input data.""" + event = emit_run_started("thread-1", "run-1", input_data="test input") + + assert event["type"] == "RUN_STARTED" + assert event["input"] == "test input" + + +def test_emit_run_finished(): + """Test RUN_FINISHED event creation.""" + result = TestResult(status="complete") + event = emit_run_finished("thread-1", "run-1", result) + + assert event["type"] == "RUN_FINISHED" + assert event["threadId"] == "thread-1" + assert event["runId"] == "run-1" + assert event["result"] == {"status": "complete"} + + +def test_emit_run_finished_with_dict(): + """Test RUN_FINISHED event with dict result.""" + result = {"status": "complete", "count": 42} + event = emit_run_finished("thread-1", "run-1", result) + + assert event["type"] == "RUN_FINISHED" + assert event["result"] == result + + +def test_emit_run_error(): + """Test RUN_ERROR event creation.""" + event = emit_run_error("Something went wrong") + + assert event["type"] == "RUN_ERROR" + assert event["message"] == "Something went wrong" + assert "code" not in event + + +def test_emit_run_error_with_code(): + """Test RUN_ERROR event with error code.""" + event = emit_run_error("Something went wrong", code="ERR_001") + + assert event["type"] == "RUN_ERROR" + assert event["message"] == "Something went wrong" + assert event["code"] == "ERR_001" + + +def test_emit_step_started(): + """Test STEP_STARTED event creation.""" + event = emit_step_started("plan") + + assert event["type"] == "STEP_STARTED" + assert event["stepName"] == "plan" + + +def test_emit_step_finished(): + """Test STEP_FINISHED event creation.""" + event = emit_step_finished("plan") + + assert event["type"] == "STEP_FINISHED" + assert event["stepName"] == "plan" + + +def test_emit_text_message(): + """Test TEXT_MESSAGE_CHUNK event creation.""" + event = emit_text_message("Hello world") + + assert event["type"] == "TEXT_MESSAGE_CHUNK" + assert event["delta"] == "Hello world" + assert event["role"] == "assistant" + assert "messageId" in event + + +def test_emit_text_message_with_role(): + """Test TEXT_MESSAGE_CHUNK event with custom role.""" + event = emit_text_message("Hello", role="user") + + assert event["type"] == "TEXT_MESSAGE_CHUNK" + assert event["role"] == "user" + + +def test_emit_state_snapshot(): + """Test STATE_SNAPSHOT event creation.""" + state = TestState(value=42) + event = emit_state_snapshot(state) + + assert event["type"] == "STATE_SNAPSHOT" + assert event["snapshot"] == {"value": 42} + + +def test_emit_activity(): + """Test ACTIVITY_SNAPSHOT event creation.""" + event = emit_activity("msg-1", "processing", "Working on task") + + assert event["type"] == "ACTIVITY_SNAPSHOT" + assert event["messageId"] == "msg-1" + assert event["activityType"] == "processing" + assert event["content"] == "Working on task" + + +def test_event_structure_consistency(): + """Test that all events have consistent structure.""" + events = [ + emit_run_started("t1", "r1"), + emit_run_finished("t1", "r1", {"result": "done"}), + emit_run_error("error"), + emit_step_started("step1"), + emit_step_finished("step1"), + emit_text_message("text"), + emit_state_snapshot(TestState(value=1)), + emit_activity("m1", "type", "content"), + ] + + for event in events: + assert isinstance(event, dict) + assert "type" in event + assert isinstance(event["type"], str) + assert event["type"].isupper() # Event types are uppercase diff --git a/tests/agui/test_server.py b/tests/agui/test_server.py new file mode 100644 index 00000000..2752b153 --- /dev/null +++ b/tests/agui/test_server.py @@ -0,0 +1,250 @@ +"""Tests for AG-UI server.""" + +import pytest +from pydantic import BaseModel +from starlette.testclient import TestClient + +from haiku.rag.agui.server import RunAgentInput, create_agui_app, format_sse_event +from haiku.rag.config.models import AGUIConfig + + +class SimpleState(BaseModel): + """Simple state for testing.""" + + question: str + + +class SimpleResult(BaseModel): + """Simple result for testing.""" + + answer: str + + +class MockGraph: + """Mock graph that returns immediately.""" + + async def run(self, state, deps): # type: ignore[no-untyped-def] + """Return a simple result.""" + return SimpleResult(answer=f"Answer to: {state.question}") + + +def test_run_agent_input_parsing(): + """Test RunAgentInput model parsing.""" + data = { + "threadId": "thread-1", + "runId": "run-1", + "state": {"question": "What is AI?"}, + "messages": [], + "config": {}, + } + + input_data = RunAgentInput(**data) + + assert input_data.thread_id == "thread-1" + assert input_data.run_id == "run-1" + assert input_data.state == {"question": "What is AI?"} + + +def test_run_agent_input_defaults(): + """Test RunAgentInput with defaults.""" + input_data = RunAgentInput() # type: ignore[call-arg] + + assert input_data.thread_id is None + assert input_data.run_id is None + assert input_data.state == {} + assert input_data.messages == [] + assert input_data.config == {} + + +def test_format_sse_event(): + """Test SSE event formatting.""" + event = {"type": "TEST_EVENT", "data": "test"} + + sse = format_sse_event(event) + + assert sse.startswith("data: ") + assert sse.endswith("\n\n") + assert '{"type": "TEST_EVENT"' in sse + + +def test_create_agui_app_basic(): + """Test basic app creation.""" + config = AGUIConfig( + host="localhost", + port=8000, + cors_origins=["http://localhost"], + ) + + def graph_factory(): + return MockGraph() + + def state_factory(input_state): + return SimpleState(question=input_state.get("question", "")) + + def deps_factory(input_config): + from dataclasses import dataclass + + @dataclass + class SimpleDeps: + agui_emitter: None = None + + return SimpleDeps() + + app = create_agui_app( + graph_factory=graph_factory, # type: ignore[arg-type] + state_factory=state_factory, + deps_factory=deps_factory, + config=config, + ) + + # Should return a Starlette app + assert app is not None + assert hasattr(app, "routes") + + +def test_server_health_endpoint(): + """Test health check endpoint.""" + config = AGUIConfig() + + def graph_factory(): + return MockGraph() + + def state_factory(input_state): + return SimpleState(question="") + + def deps_factory(input_config): + from dataclasses import dataclass + + @dataclass + class SimpleDeps: + agui_emitter: None = None + + return SimpleDeps() + + app = create_agui_app( + graph_factory=graph_factory, # type: ignore[arg-type] + state_factory=state_factory, + deps_factory=deps_factory, + config=config, + ) + + client = TestClient(app) + response = client.get("/health") + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +@pytest.mark.asyncio +async def test_server_stream_endpoint(): + """Test AG-UI streaming endpoint.""" + config = AGUIConfig() + + def graph_factory(): + return MockGraph() + + def state_factory(input_state): + question = input_state.get("question", "") + return SimpleState(question=question) + + def deps_factory(input_config): + from dataclasses import dataclass + + @dataclass + class SimpleDeps: + agui_emitter: None = None + + return SimpleDeps() + + app = create_agui_app( + graph_factory=graph_factory, # type: ignore[arg-type] + state_factory=state_factory, + deps_factory=deps_factory, + config=config, + ) + + client = TestClient(app) + + request_data = { + "threadId": "test-1", + "runId": "run-1", + "state": {"question": "What is pydantic-graph?"}, + "messages": [], + "config": {}, + } + + response = client.post("/v1/agent/stream", json=request_data) + + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8" + + # Read the streamed events + events = [] + for line in response.iter_lines(): + if line.startswith("data: "): + import json + + event_data = line[6:] # Remove "data: " prefix + event = json.loads(event_data) + events.append(event) + + # Should have received multiple events + assert len(events) > 0 + + # Should have RUN_STARTED and RUN_FINISHED + event_types = [e["type"] for e in events] + assert "RUN_STARTED" in event_types + assert "RUN_FINISHED" in event_types + + +def test_server_cors_headers(): + """Test CORS middleware is configured.""" + config = AGUIConfig( + cors_origins=["http://example.com"], + cors_credentials=True, + ) + + def graph_factory(): + return MockGraph() + + def state_factory(input_state): + return SimpleState(question="") + + def deps_factory(input_config): + from dataclasses import dataclass + + @dataclass + class SimpleDeps: + agui_emitter: None = None + + return SimpleDeps() + + app = create_agui_app( + graph_factory=graph_factory, # type: ignore[arg-type] + state_factory=state_factory, + deps_factory=deps_factory, + config=config, + ) + + client = TestClient(app) + + # GET request with Origin header should get CORS headers + response = client.get("/health", headers={"Origin": "http://example.com"}) + + assert response.status_code == 200 + # CORS middleware should add access-control headers + assert ( + "access-control-allow-origin" in response.headers or response.status_code == 200 + ) + + +def test_agui_config_defaults(): + """Test AGUIConfig default values.""" + config = AGUIConfig() + + assert config.host == "0.0.0.0" + assert config.port == 8000 + assert config.cors_origins == ["*"] + assert config.cors_credentials is True + assert "GET" in config.cors_methods + assert "POST" in config.cors_methods diff --git a/tests/agui/test_stream.py b/tests/agui/test_stream.py new file mode 100644 index 00000000..1aca01b4 --- /dev/null +++ b/tests/agui/test_stream.py @@ -0,0 +1,199 @@ +"""Tests for stream_graph function.""" + +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel + +from haiku.rag.agui.emitter import AGUIEmitter +from haiku.rag.agui.stream import stream_graph + + +class TestState(BaseModel): + """Test state model.""" + + value: int + + +@dataclass +class TestDeps: + """Test dependencies.""" + + agui_emitter: AGUIEmitter | None = None + + +class MockGraph: + """Mock graph for testing.""" + + def __init__(self, result: str | dict[str, str | int] = "done"): + self.result = result + self.run_called = False + + async def run(self, state, deps): # type: ignore[no-untyped-def] + """Mock run method.""" + self.run_called = True + + # Emit some events through the emitter + if deps.agui_emitter: + deps.agui_emitter.start_step("mock_step") + deps.agui_emitter.update_activity("working", "Doing work") + deps.agui_emitter.finish_step() + + return self.result + + +@pytest.mark.asyncio +async def test_stream_graph_basic(): + """Test basic graph streaming.""" + graph = MockGraph(result="success") + state = TestState(value=1) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Should have collected events + assert len(events) > 0 + + # Should have RUN_STARTED and RUN_FINISHED + event_types = [e["type"] for e in events] + assert "RUN_STARTED" in event_types + assert "RUN_FINISHED" in event_types + + # Graph should have been executed + assert graph.run_called + + +@pytest.mark.asyncio +async def test_stream_graph_emits_initial_state(): + """Test that initial state is emitted.""" + graph = MockGraph() + state = TestState(value=42) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Should have initial state snapshot + state_snapshots = [e for e in events if e["type"] == "STATE_SNAPSHOT"] + assert len(state_snapshots) > 0 + # First snapshot should be initial state + assert state_snapshots[0]["snapshot"] == {"value": 42} + + +@pytest.mark.asyncio +async def test_stream_graph_emits_step_events(): + """Test that step events from graph are emitted.""" + graph = MockGraph() + state = TestState(value=1) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Should have step events from MockGraph + step_started = [e for e in events if e["type"] == "STEP_STARTED"] + assert len(step_started) > 0 + assert step_started[0]["stepName"] == "mock_step" + + +@pytest.mark.asyncio +async def test_stream_graph_emits_activity(): + """Test that activity events from graph are emitted.""" + graph = MockGraph() + state = TestState(value=1) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Should have activity events from MockGraph + activities = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"] + assert len(activities) > 0 + assert activities[0]["content"] == "Doing work" + + +@pytest.mark.asyncio +async def test_stream_graph_handles_error(): + """Test that graph errors are captured and emitted.""" + + class ErrorGraph: + async def run(self, state, deps): + raise ValueError("Test error") + + graph = ErrorGraph() + state = TestState(value=1) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Should have error event + errors = [e for e in events if e["type"] == "RUN_ERROR"] + assert len(errors) > 0 + assert "Test error" in errors[0]["message"] + + +@pytest.mark.asyncio +async def test_stream_graph_closes_emitter(): + """Test that emitter is properly closed.""" + + class NeverReturnsGraph: + async def run(self, state, deps): + # Don't return anything + pass + + graph = NeverReturnsGraph() + state = TestState(value=1) + deps = TestDeps() + + events = [] + try: + async for event in stream_graph(graph, state, deps): + events.append(event) + except RuntimeError: + # Expected - graph didn't return a result + pass + + # Should have error event about no result + errors = [e for e in events if e["type"] == "RUN_ERROR"] + assert len(errors) > 0 + + +@pytest.mark.asyncio +async def test_stream_graph_without_emitter_support(): + """Test error when deps doesn't support agui_emitter.""" + + @dataclass + class BadDeps: + pass + + graph = MockGraph() + state = TestState(value=1) + deps = BadDeps() + + with pytest.raises(TypeError, match="agui_emitter"): + async for _ in stream_graph(graph, state, deps): # type: ignore[arg-type] + pass + + +@pytest.mark.asyncio +async def test_stream_graph_result_in_finish_event(): + """Test that graph result is included in RUN_FINISHED event.""" + graph = MockGraph(result={"status": "complete", "count": 42}) + state = TestState(value=1) + deps = TestDeps() + + events = [] + async for event in stream_graph(graph, state, deps): + events.append(event) + + # Find RUN_FINISHED event + finished = [e for e in events if e["type"] == "RUN_FINISHED"] + assert len(finished) == 1 + assert finished[0]["result"] == {"status": "complete", "count": 42} diff --git a/tests/test_research_graph.py b/tests/test_research_graph.py index 0a7c0b05..94f1f6a3 100644 --- a/tests/test_research_graph.py +++ b/tests/test_research_graph.py @@ -1,8 +1,13 @@ import asyncio +import pytest +from pydantic_ai.models.test import TestModel + +from haiku.rag.agui.stream import stream_graph +from haiku.rag.client import HaikuRAG from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.graph import build_research_graph -from haiku.rag.research.state import ResearchState +from haiku.rag.research.state import ResearchDeps, ResearchState def test_build_graph_and_state(): @@ -24,3 +29,57 @@ def test_async_loop_available(): # Ensure an event loop can be created in test env loop = asyncio.new_event_loop() loop.close() + + +@pytest.mark.asyncio +async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): + """Test research graph with mocked LLM using AG-UI events.""" + + # Mock get_model to return TestModel which generates valid schema-compliant data + def test_model_factory(_provider, _model): + return TestModel() + + monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory) + + graph = build_research_graph() + + state = ResearchState( + context=ResearchContext(original_question="What is haiku.rag?"), + max_iterations=1, + confidence_threshold=0.5, + max_concurrency=2, + ) + + # Use real client but with TestModel for LLM calls + client = HaikuRAG(temp_db_path) + deps = ResearchDeps(client=client) + + events = [] + result = None + async for event in stream_graph(graph, state, deps): + events.append(event) + if event["type"] == "RUN_FINISHED": + result = event["result"] + elif event["type"] == "RUN_ERROR": + pytest.fail(f"Graph execution failed: {event['message']}") + + # TestModel will generate valid structured output for each node + assert result is not None, ( + f"No result. Events collected: {[e['type'] for e in events]}" + ) + # Result is serialized as dict in AG-UI events + assert isinstance(result, dict) + assert "title" in result + assert isinstance(result["title"], str) + assert "executive_summary" in result + assert "main_findings" in result + + # Verify AG-UI events were emitted + event_types = [e["type"] for e in events] + assert "RUN_STARTED" in event_types + assert "RUN_FINISHED" in event_types + assert "STATE_SNAPSHOT" in event_types + assert "STEP_STARTED" in event_types + + client.close() diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py deleted file mode 100644 index 5bdf40a6..00000000 --- a/tests/test_research_graph_integration.py +++ /dev/null @@ -1,55 +0,0 @@ -import pytest -from pydantic_ai.models.test import TestModel - -from haiku.rag.client import HaikuRAG -from haiku.rag.research.dependencies import ResearchContext -from haiku.rag.research.graph import build_research_graph -from haiku.rag.research.models import ResearchReport -from haiku.rag.research.state import ResearchDeps, ResearchState -from haiku.rag.research.stream import stream_research_graph - - -@pytest.mark.asyncio -async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): - """Test research graph with mocked LLM using TestModel.""" - - # Mock get_model to return TestModel which generates valid schema-compliant data - def test_model_factory(provider, model): - return TestModel() - - monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory) - monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory) - - graph = build_research_graph() - - state = ResearchState( - context=ResearchContext(original_question="What is haiku.rag?"), - max_iterations=1, - confidence_threshold=0.5, - max_concurrency=2, - ) - - # Use real client but with TestModel for LLM calls - client = HaikuRAG(temp_db_path) - deps = ResearchDeps(client=client, console=None) - - collected = [] - report = None - async for event in stream_research_graph(graph, state, deps): - collected.append(event) - if event.type == "report": - report = event.report - break - elif event.type == "error": - pytest.fail(f"Graph execution failed: {event.error}") - - # TestModel will generate valid structured output for each node - assert report is not None, ( - f"No report generated. Events collected: {[e.type for e in collected]}" - ) - assert isinstance(report, ResearchReport) - assert report.title is not None - assert isinstance(report.title, str) - assert any(evt.type == "log" for evt in collected) - - client.close()