mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-02-10 23:05:17 -05:00
1a1fe7c0b7
Currently it's only possible to open latest graph from monitor and see the node execution results only when manually running. This PR adds ability to open running and finished graphs in builder. ### Changes 🏗️ Builder now handles graph version and execution ID in addition to graph ID when opening a graph. When an execution ID is provided, node execution results are fetched and subscribed to in real time. This makes it possible to open a graph that is already executing and see both existing node execution data and real-time updates (if it's still running). - Use graph version and execution id on the builder page and in `useAgentGraph` - Use graph version on the `execute_graph` endpoint - Use graph version on the websockets to distinguish between versions - Move `formatEdgeID` to utils; it's used in `useAgentGraph.ts` and in `Flow.tsx` ### Checklist 📋 #### For code changes: - [x] I have clearly listed my changes in the PR description - [x] I have made a test plan - [x] I have tested my changes according to the test plan: - [x] Opening finished execution restores node results - [x] Opening running execution restores results and continues to run properly - [x] Results are separate for each graph across multiple tabs #### For configuration changes: - [ ] `.env.example` is updated or already compatible with my changes - [ ] `docker-compose.yml` is updated or already compatible with my changes - [ ] I have included a list of my configuration changes in the PR description (under **Changes**) <details> <summary>Examples of configuration changes</summary> - Changing ports - Adding new services that need to communicate with each other - Secrets or environment variable changes - New or infrastructure changes such as databases </details> --------- Co-authored-by: Zamil Majdy <zamil.majdy@agpt.co>
119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
from datetime import datetime, timezone
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from fastapi import WebSocket
|
|
|
|
from backend.data.execution import ExecutionResult, ExecutionStatus
|
|
from backend.server.conn_manager import ConnectionManager
|
|
from backend.server.model import Methods, WsMessage
|
|
|
|
|
|
@pytest.fixture
|
|
def connection_manager() -> ConnectionManager:
|
|
return ConnectionManager()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_websocket() -> AsyncMock:
|
|
websocket: AsyncMock = AsyncMock(spec=WebSocket)
|
|
websocket.send_text = AsyncMock()
|
|
return websocket
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connect(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
await connection_manager.connect(mock_websocket)
|
|
assert mock_websocket in connection_manager.active_connections
|
|
mock_websocket.accept.assert_called_once()
|
|
|
|
|
|
def test_disconnect(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
connection_manager.active_connections.add(mock_websocket)
|
|
connection_manager.subscriptions["test_graph_1"] = {mock_websocket}
|
|
|
|
connection_manager.disconnect(mock_websocket)
|
|
|
|
assert mock_websocket not in connection_manager.active_connections
|
|
assert mock_websocket not in connection_manager.subscriptions["test_graph_1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subscribe(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
await connection_manager.subscribe("test_graph", 1, mock_websocket)
|
|
assert mock_websocket in connection_manager.subscriptions["test_graph_1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unsubscribe(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
connection_manager.subscriptions["test_graph_1"] = {mock_websocket}
|
|
|
|
await connection_manager.unsubscribe("test_graph", 1, mock_websocket)
|
|
|
|
assert "test_graph" not in connection_manager.subscriptions
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_execution_result(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
connection_manager.subscriptions["test_graph_1"] = {mock_websocket}
|
|
result: ExecutionResult = ExecutionResult(
|
|
graph_id="test_graph",
|
|
graph_version=1,
|
|
graph_exec_id="test_exec_id",
|
|
node_exec_id="test_node_exec_id",
|
|
node_id="test_node_id",
|
|
block_id="test_block_id",
|
|
status=ExecutionStatus.COMPLETED,
|
|
input_data={"input1": "value1"},
|
|
output_data={"output1": ["result1"]},
|
|
add_time=datetime.now(tz=timezone.utc),
|
|
queue_time=None,
|
|
start_time=datetime.now(tz=timezone.utc),
|
|
end_time=datetime.now(tz=timezone.utc),
|
|
)
|
|
|
|
await connection_manager.send_execution_result(result)
|
|
|
|
mock_websocket.send_text.assert_called_once_with(
|
|
WsMessage(
|
|
method=Methods.EXECUTION_EVENT,
|
|
channel="test_graph_1",
|
|
data=result.model_dump(),
|
|
).model_dump_json()
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_execution_result_no_subscribers(
|
|
connection_manager: ConnectionManager, mock_websocket: AsyncMock
|
|
) -> None:
|
|
result: ExecutionResult = ExecutionResult(
|
|
graph_id="test_graph",
|
|
graph_version=1,
|
|
graph_exec_id="test_exec_id",
|
|
node_exec_id="test_node_exec_id",
|
|
node_id="test_node_id",
|
|
block_id="test_block_id",
|
|
status=ExecutionStatus.COMPLETED,
|
|
input_data={"input1": "value1"},
|
|
output_data={"output1": ["result1"]},
|
|
add_time=datetime.now(),
|
|
queue_time=None,
|
|
start_time=datetime.now(),
|
|
end_time=datetime.now(),
|
|
)
|
|
|
|
await connection_manager.send_execution_result(result)
|
|
|
|
mock_websocket.send_text.assert_not_called()
|