jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

feat: add functionality to stream or poll events from the build process (#5940)

Browse source 
* refactor: simplify TaskService by removing Celery integration

* refactor: enhance AnyIO backend task management and error handling

* refactor: restructure flow building process and enhance event handling

* feat: implement QueueService for managing job queues and tasks

* feat: Add QueueService for managing job queues and task lifecycle

* feat!: Enhance flow building with QueueService integration for job management

* revert changes to async session

* feat: Integrate QueueService into lifespan management for task handling

* refactor: Enhance QueueService with robust lifecycle management and cleanup mechanisms

* refactor: Simplify docstring for get_queue_service function

* refactor: Update import statements in queue factory for clarity

* test: Improve chat endpoint tests with comprehensive build flow scenarios

* refactor: Streamline flow build event generation and queue management

* refactor: Improve flow build process with modular event handling and URL parameter management

* test: Enhance loop component tests with flow build and event handling utilities

* fix: Improve transaction logging with null flow_id handling and debug logging

* fix: Remove unnecessary db.refresh() in transaction logging

* feat: Add polling mode for build events with optional streaming

* feat: Add event delivery configuration option to settings

* feat: Implement polling for build events with configuration-driven streaming

* refactor: Enhance queue cleanup with improved logging and error handling

* test: Improve build event stream assertion with detailed error reporting

* test: Add benchmark test for build flow polling mechanism

* refactor: Remove redundant end event in flow event generation

* test: Update test_component_tool_with_api_key to use async flow

* Convert test method to async
* Update graph start method to use async_start()
* Add client parameter to test method
* Modify tool retrieval to use async method

* test: Add session ID generation in agent component test

* fix: Enhance error handling in LCAgentComponent for ExceptionWithMessageError

Add null checks to prevent potential AttributeError when handling agent message deletion

* test: Refactor test assertion for model name options in agent component

Simplify the assertion for checking "sonnet" in model name options by extracting the options to a variable first

* refactor: Improve type hints and imports in AnyIOTaskResult

* Add `__future__` import for type annotations
* Use conditional import for `Callable` and `TracebackType`
* Explicitly type `_traceback` attribute with `TracebackType | None`

* fix: Add null checks for event task in build events streaming

* refactor: Extract build and disconnect handling into separate modules

This commit separates the build and disconnect handling logic from the chat API module into dedicated files:

- Created `langflow/api/build.py` to house the flow generation and vertex building logic
- Created `langflow/api/disconnect.py` to manage the custom streaming response with disconnect handling
- Removed duplicate code from `langflow/api/v1/chat.py`
- Improved code organization and modularity

* refactor: Extract flow build and event handling methods in build API

This commit introduces two new methods in the build API to improve code organization and reusability:

- `start_flow_build()`: Centralizes the logic for creating a queue and starting a flow build task
- `get_flow_events_response()`: Handles retrieving flow events for both streaming and polling modes

The changes simplify the chat API endpoint implementations and reduce code duplication

* refactor: Deprecate /task endpoint and add informative error message

This commit marks the /task endpoint as deprecated and raises an HTTP 400 error with a clear message indicating the endpoint will be removed in a future version. The changes improve API clarity and guide users towards the recommended /run endpoint.

* refactor: Update log_transaction function return type annotation

Modify the return type hint for log_transaction to allow for potential None return, improving type safety and clarity in the transactions logging method.

* feat: Add configurable event delivery streaming option

Introduce support for configurable event delivery mode (streaming/polling) across frontend components:

- Add `shouldStreamEvents()` method in NodeStatus and IOModal
- Update flowStore to accept a `stream` parameter
- Modify buildFlowVerticesWithFallback to use the stream parameter directly
- Remove hardcoded polling logic in buildUtils

* fix: Add null check for flow_id in log_transaction function

Prevent logging transactions without a flow_id by returning None, ensuring data integrity and avoiding potential database errors

* chore: Update changes-filter to include chat API path

* fix: Add error handling and null checks in flow transaction deletion

Improve robustness of transaction and vertex build deletion by:
- Adding a null check for flow_id in delete_transactions_by_flow_id
- Wrapping deletion operations in try-except blocks to prevent failures
- Logging debug messages for any deletion errors

* fix: Remove redundant commit in vertex builds deletion

Remove unnecessary db.commit() from delete_vertex_builds_by_flow_id to prevent duplicate commits and simplify database transaction handling

* fix: Improve transaction logging with debug message and return value

Move debug logging from CRUD layer to utils to provide more context about logged transactions and ensure proper error handling

* fix: Improve error handling and transaction cleanup in active_user fixture

Enhance user and transaction cleanup process in test fixture by:
- Adding separate try-except blocks for transaction/vertex build deletion and user deletion
- Adding debug logging for potential errors during cleanup
- Ensuring proper session commits for each operation

* propagate parent task CancelledError instead of supressing it

Co-authored-by: Christophe Bornet <cbornet@hotmail.com>

* refactor: Rename queue service to job queue service

Restructure job queue management by:
- Renaming QueueService to JobQueueService
- Moving queue-related files to a new job_queue directory
- Updating import paths and service type references
- Enhancing job queue service with more robust async job management

* fix: Correct ServiceType enum reference for job queue service

Update the service type constant to match the recently renamed JobQueueService, ensuring consistent service type referencing across the application.

* refactor: Enhance JobQueueService with comprehensive logging and documentation

Improve the JobQueueService implementation by:
- Adding detailed docstrings with clear explanations of methods and attributes
- Implementing comprehensive logging throughout the service
- Enhancing error handling and logging for queue and task management
- Providing more context in log messages for debugging and monitoring

* docs: Improve stop method docstring for JobQueueService

Enhance the documentation for the stop method by:
- Providing a more detailed and precise description of the shutdown process
- Clarifying the steps involved in gracefully stopping the service
- Improving the explanation of resource cleanup and task cancellation

* fix: Properly handle task cancellation and propagate exceptions in JobQueueService

* fix: Enhance error handling and raise appropriate exceptions in JobQueueService methods

* improve docstring

Co-authored-by: Christophe Bornet <cbornet@hotmail.com>

* fix: Remove redundant exception raises in JobQueueService methods and improve cleanup logic

* fix: Improve logging during job cleanup and handle exceptions more appropriately

* feat: add utility to run tests with multiple event delivery modes

* feat: integrate withEventDeliveryModes utility in multiple test files

* refactor: replace status code assertions with httpx codes for clarity

* remove noqa comment and change argument name

Co-authored-by: Christophe Bornet <cbornet@hotmail.com>

* refactor: streamline event polling logic in get_flow_events_response

* refactor: use getattr for safer session attribute access

Co-author: @cbornet

* feat: add is_started method to JobQueueService

* refactor: modify JobQueueService start method and main.py queue service initialization

* feat: add ready state and teardown method to JobQueueService

* refactor: simplify job queue cleanup logic in JobQueueService

* refactor: improve error logging in active_user fixture

* refactor: improve AnyIO task management with TaskGroup and CancelScope

* refactor: Implement LimitedBackgroundTasks for controlled vertex build logging (#6312)

* feat: implement LimitedBackgroundTasks for controlled vertex build logging

* refactor: replace BackgroundTasks with LimitedBackgroundTasks in build_flow endpoint

* refactor: improve task cancellation error handling in JobQueueService

* refactor: Rename LimitedBackgroundTasks to LimitVertexBuildBackgroundTasks

* feat: Add EventDeliveryType enum for event delivery methods

* feat: Add polling constants for endpoint and streaming status

* refactor: Update buildFlowVerticesWithFallback to use polling constants

* refactor: Update event delivery handling to use EventDeliveryType enum

* [autofix.ci] apply automated fixes

* fix: Improve error handling for cancelled build tasks

Refactor generate_flow_events to properly handle and propagate CancelledError
- Remove unnecessary task creation and manual cancellation
- Directly await _build_vertex instead of creating a separate task
- Improve exception logging for cancelled tasks
- Ensure CancelledError is raised instead of being silently handled

* fix: trigger event_manager.on_end after error handling in generate_flow_events

---------

Co-authored-by: Christophe Bornet <cbornet@hotmail.com>
Co-authored-by: anovazzi1 <otavio2204@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2025-02-18 11:24:39 -03:00 • committed by GitHub
commit 1530e61669
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
51 changed files with 1618 additions and 814 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/backend/tests/unit/base/tools/test_component_toolkit.py
View file

@ -35,17 +35,21 @@ def test_component_tool():
@pytest.mark.api_key_required
def test_component_tool_with_api_key():
@pytest.mark.usefixtures("client")
async def test_component_tool_with_api_key():
chat_output = ChatOutput()
openai_llm = OpenAIModelComponent()
openai_llm.set(api_key=os.environ["OPENAI_API_KEY"])
tool_calling_agent = ToolCallingAgentComponent()
tools = await chat_output.to_toolkit()
tool_calling_agent.set(
llm=openai_llm.build_model, tools=[chat_output], input_value="Which tools are available? Please tell its name."
llm=openai_llm.build_model,
tools=tools,
input_value="Which tools are available? Please tell its name.",
)
g = Graph(start=tool_calling_agent, end=tool_calling_agent)
assert g is not None
results = list(g.start())
results = [result async for result in g.async_start()]
assert len(results) == 4
assert "message_response" in tool_calling_agent._outputs_map["response"].value.get_text()

75
src/backend/tests/unit/build_utils.py Normal file
View file

@ -0,0 +1,75 @@
import json
from typing import Any
from uuid import UUID
from httpx import AsyncClient, codes
async def create_flow(client: AsyncClient, flow_data: str, headers: dict[str, str]) -> UUID:
"""Create a flow and return its ID."""
response = await client.post("api/v1/flows/", json=json.loads(flow_data), headers=headers)
assert response.status_code == codes.CREATED
return UUID(response.json()["id"])
async def build_flow(
client: AsyncClient, flow_id: UUID, headers: dict[str, str], json: dict[str, Any] | None = None
) -> dict[str, Any]:
"""Start a flow build and return the job_id."""
if json is None:
json = {}
response = await client.post(f"api/v1/build/{flow_id}/flow", json=json, headers=headers)
assert response.status_code == codes.OK
return response.json()
async def get_build_events(client: AsyncClient, job_id: str, headers: dict[str, str]):
"""Get events for a build job."""
return await client.get(f"api/v1/build/{job_id}/events", headers=headers)
async def consume_and_assert_stream(response, job_id):
"""Consume the event stream and assert the expected event structure."""
count = 0
lines = []
async for line in response.aiter_lines():
# Skip empty lines (ndjson uses double newlines)
if not line:
continue
lines.append(line)
parsed = json.loads(line)
if "job_id" in parsed:
assert parsed["job_id"] == job_id
continue
if count == 0:
# First event should be vertices_sorted
assert parsed["event"] == "vertices_sorted", (
"Invalid first event. Expected 'vertices_sorted'. Full event stream:\n" + "\n".join(lines)
)
ids = parsed["data"]["ids"]
ids.sort()
assert ids == ["ChatInput-CIGht"], "Invalid ids in first event. Full event stream:\n" + "\n".join(lines)
to_run = parsed["data"]["to_run"]
to_run.sort()
assert to_run == ["ChatInput-CIGht", "ChatOutput-QA7ej", "Memory-amN4Z", "Prompt-iWbCC"], (
"Invalid to_run list in first event. Full event stream:\n" + "\n".join(lines)
)
elif count > 0 and count < 5:
# Next events should be end_vertex events
assert parsed["event"] == "end_vertex", (
f"Invalid event at position {count}. Expected 'end_vertex'. Full event stream:\n" + "\n".join(lines)
)
assert parsed["data"]["build_data"] is not None, (
f"Missing build_data at position {count}. Full event stream:\n" + "\n".join(lines)
)
elif count == 5:
# Final event should be end
assert parsed["event"] == "end", "Invalid final event. Expected 'end'. Full event stream:\n" + "\n".join(
lines
)
else:
raise ValueError(f"Unexpected event at position {count}. Full event stream:\n" + "\n".join(lines))
count += 1

6
src/backend/tests/unit/components/agents/test_agent_component.py
View file

@ -78,9 +78,8 @@ class TestAgentComponent(ComponentTestBaseWithoutClient):
assert all(provider in updated_config["agent_llm"]["options"] for provider in MODEL_PROVIDERS_DICT)
assert "Anthropic" in updated_config["agent_llm"]["options"]
assert updated_config["agent_llm"]["input_types"] == []
assert any("sonnet" in option.lower() for option in updated_config["model_name"]["options"]), (
f"Options: {updated_config['model_name']['options']}"
)
options = updated_config["model_name"]["options"]
assert any("sonnet" in option.lower() for option in options), f"Options: {options}"
# Test updating build config for Custom
updated_config = await component.update_build_config(build_config, "Custom", "agent_llm")
@ -113,6 +112,7 @@ async def test_agent_component_with_calculator():
model_name="gpt-4o",
llm_type="OpenAI",
temperature=temperature,
_session_id=str(uuid4()),
)
response = await agent.message_response()

27
src/backend/tests/unit/components/logic/test_loop.py
View file

@ -1,14 +1,15 @@
from uuid import UUID
import orjson
import pytest
from httpx import AsyncClient
from langflow.components.logic.loop import LoopComponent
from langflow.memory import aget_messages
from langflow.schema.data import Data
from langflow.services.database.models.flow import FlowCreate
from orjson import orjson
from tests.base import ComponentTestBaseWithClient
from tests.unit.build_utils import build_flow, get_build_events
TEXT = (
"lorem ipsum dolor sit amet lorem ipsum dolor sit amet lorem ipsum dolor sit amet. "
@ -62,15 +63,25 @@ class TestLoopComponentWithAPI(ComponentTestBaseWithClient):
assert len(messages[1].text) > 0
async def test_build_flow_loop(self, client, json_loop_test, logged_in_headers):
# TODO: Add a test for the loop where the loop component gets updated even the component in json
"""Test building a flow with a loop component."""
# Create the flow
flow_id = await self._create_flow(client, json_loop_test, logged_in_headers)
async with client.stream("POST", f"api/v1/build/{flow_id}/flow", json={}, headers=logged_in_headers) as r:
async for line in r.aiter_lines():
# httpx split by \n, but ndjson sends two \n for each line
if line:
# Process the line if needed
pass
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers)
job_id = build_response["job_id"]
assert job_id is not None
# Get the events stream
events_response = await get_build_events(client, job_id, logged_in_headers)
assert events_response.status_code == 200
# Process the events stream
async for line in events_response.aiter_lines():
if not line: # Skip empty lines
continue
# Process events if needed
# We could add specific assertions here for loop-related events
await self.check_messages(flow_id)

265
src/backend/tests/unit/test_chat_endpoint.py
View file

@ -1,42 +1,63 @@
import json
import asyncio
import uuid
from uuid import UUID
import pytest
from httpx import codes
from langflow.memory import aget_messages
from langflow.services.database.models.flow import FlowCreate, FlowUpdate
from orjson import orjson
from langflow.services.database.models.flow import FlowUpdate
from tests.unit.build_utils import build_flow, consume_and_assert_stream, create_flow, get_build_events
@pytest.mark.benchmark
async def test_build_flow(client, json_memory_chatbot_no_llm, logged_in_headers):
flow_id = await _create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
"""Test the build flow endpoint with the new two-step process."""
# First create the flow
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
async with client.stream("POST", f"api/v1/build/{flow_id}/flow", json={}, headers=logged_in_headers) as r:
await consume_and_assert_stream(r)
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers)
job_id = build_response["job_id"]
assert job_id is not None
await check_messages(flow_id)
# Get the events stream
events_response = await get_build_events(client, job_id, logged_in_headers)
assert events_response.status_code == codes.OK
# Consume and verify the events
await consume_and_assert_stream(events_response, job_id)
@pytest.mark.benchmark
async def test_build_flow_from_request_data(client, json_memory_chatbot_no_llm, logged_in_headers):
flow_id = await _create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
response = await client.get("api/v1/flows/" + str(flow_id), headers=logged_in_headers)
"""Test building a flow from request data."""
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
response = await client.get(f"api/v1/flows/{flow_id}", headers=logged_in_headers)
flow_data = response.json()
async with client.stream(
"POST", f"api/v1/build/{flow_id}/flow", json={"data": flow_data["data"]}, headers=logged_in_headers
) as r:
await consume_and_assert_stream(r)
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers, json={"data": flow_data["data"]})
job_id = build_response["job_id"]
# Get the events stream
events_response = await get_build_events(client, job_id, logged_in_headers)
assert events_response.status_code == codes.OK
# Consume and verify the events
await consume_and_assert_stream(events_response, job_id)
await check_messages(flow_id)
async def test_build_flow_with_frozen_path(client, json_memory_chatbot_no_llm, logged_in_headers):
flow_id = await _create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
"""Test building a flow with a frozen path."""
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
response = await client.get("api/v1/flows/" + str(flow_id), headers=logged_in_headers)
response = await client.get(f"api/v1/flows/{flow_id}", headers=logged_in_headers)
flow_data = response.json()
flow_data["data"]["nodes"][0]["data"]["node"]["frozen"] = True
# Update the flow with frozen path
response = await client.patch(
f"api/v1/flows/{flow_id}",
json=FlowUpdate(name="Flow", description="description", data=flow_data["data"]).model_dump(),
@ -44,151 +65,131 @@ async def test_build_flow_with_frozen_path(client, json_memory_chatbot_no_llm, l
)
response.raise_for_status()
async with client.stream("POST", f"api/v1/build/{flow_id}/flow", json={}, headers=logged_in_headers) as r:
await consume_and_assert_stream(r)
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers)
job_id = build_response["job_id"]
# Get the events stream
events_response = await get_build_events(client, job_id, logged_in_headers)
assert events_response.status_code == codes.OK
# Consume and verify the events
await consume_and_assert_stream(events_response, job_id)
await check_messages(flow_id)
async def check_messages(flow_id):
messages = await aget_messages(flow_id=UUID(flow_id), order="ASC")
if isinstance(flow_id, str):
flow_id = UUID(flow_id)
messages = await aget_messages(flow_id=flow_id, order="ASC")
flow_id_str = str(flow_id)
assert len(messages) == 2
assert messages[0].session_id == flow_id
assert messages[0].session_id == flow_id_str
assert messages[0].sender == "User"
assert messages[0].sender_name == "User"
assert messages[0].text == ""
assert messages[1].session_id == flow_id
assert messages[1].session_id == flow_id_str
assert messages[1].sender == "Machine"
assert messages[1].sender_name == "AI"
async def consume_and_assert_stream(r):
count = 0
async for line in r.aiter_lines():
# httpx split by \n, but ndjson sends two \n for each line
if not line:
continue
parsed = json.loads(line)
if count == 0:
assert parsed["event"] == "vertices_sorted"
ids = parsed["data"]["ids"]
ids.sort()
assert ids == ["ChatInput-CIGht"]
to_run = parsed["data"]["to_run"]
to_run.sort()
assert to_run == ["ChatInput-CIGht", "ChatOutput-QA7ej", "Memory-amN4Z", "Prompt-iWbCC"]
elif count > 0 and count < 5:
assert parsed["event"] == "end_vertex"
assert parsed["data"]["build_data"] is not None
elif count == 5:
assert parsed["event"] == "end"
else:
msg = f"Unexpected line: {line}"
raise ValueError(msg)
count += 1
@pytest.mark.benchmark
async def test_build_flow_invalid_job_id(client, logged_in_headers):
"""Test getting events for an invalid job ID."""
invalid_job_id = str(uuid.uuid4())
response = await get_build_events(client, invalid_job_id, logged_in_headers)
assert response.status_code == codes.NOT_FOUND
assert "No queue found for job_id" in response.json()["detail"]
async def _create_flow(client, json_memory_chatbot_no_llm, logged_in_headers):
vector_store = orjson.loads(json_memory_chatbot_no_llm)
data = vector_store["data"]
vector_store = FlowCreate(name="Flow", description="description", data=data, endpoint_name="f")
response = await client.post("api/v1/flows/", json=vector_store.model_dump(), headers=logged_in_headers)
response.raise_for_status()
return response.json()["id"]
@pytest.mark.benchmark
async def test_build_flow_invalid_flow_id(client, logged_in_headers):
"""Test starting a build with an invalid flow ID."""
invalid_flow_id = uuid.uuid4()
response = await client.post(f"api/v1/build/{invalid_flow_id}/flow", json={}, headers=logged_in_headers)
assert response.status_code == codes.NOT_FOUND
# TODO: Fix this test
# async def test_multiple_runs_with_no_payload_generate_max_vertex_builds(
# client, json_memory_chatbot_no_llm, logged_in_headers
# ):
# """Test that multiple builds of a flow generate the correct number of vertex builds."""
# # Create the initial flow
# flow_id = await _create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
@pytest.mark.benchmark
async def test_build_flow_start_only(client, json_memory_chatbot_no_llm, logged_in_headers):
"""Test only the build flow start endpoint."""
# First create the flow
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
# # Get the flow data to count nodes before making requests
# response = await client.get(f"api/v1/flows/{flow_id}", headers=logged_in_headers)
# flow_data = response.json()
# num_nodes = len(flow_data["data"]["nodes"])
# max_vertex_builds = get_settings_service().settings.max_vertex_builds_per_vertex
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers)
# logger.debug(f"Starting test with {num_nodes} nodes, max_vertex_builds={max_vertex_builds}")
# Assert response structure
assert "job_id" in build_response
assert isinstance(build_response["job_id"], str)
# Verify it's a valid UUID
assert uuid.UUID(build_response["job_id"])
# # Make multiple build requests - ensure we exceed max_vertex_builds significantly
# num_requests = max_vertex_builds * 3 # Triple the max to ensure rotation
# for i in range(num_requests):
# # Generate a random session ID for each request
# session_id = session_id_generator()
# payload = {"inputs": {"session": session_id, "type": "chat", "input_value": f"Test message {i + 1}"}}
# async with client.stream("POST", f"api/v1/build/{flow_id}/flow",
# json=payload, headers=logged_in_headers) as r:
# await consume_and_assert_stream(r)
@pytest.mark.benchmark
async def test_build_flow_start_with_inputs(client, json_memory_chatbot_no_llm, logged_in_headers):
"""Test the build flow start endpoint with input data."""
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
# # Add a small delay between requests to ensure proper ordering
# await asyncio.sleep(0.1)
# Start build with some input data
test_inputs = {"inputs": {"session": "test_session", "input_value": "test message"}}
# # Track builds after each request
# async with session_scope() as session:
# builds = await get_vertex_builds_by_flow_id(db=session, flow_id=flow_id)
# by_vertex = {}
# for build in builds:
# build_dict = build.model_dump()
# vertex_id = build_dict.get("id")
# by_vertex.setdefault(vertex_id, []).append(build_dict)
build_response = await build_flow(client, flow_id, logged_in_headers, json=test_inputs)
# # Log state of each vertex with more details
# for vertex_id, vertex_builds in by_vertex.items():
# vertex_builds.sort(key=lambda x: x.get("timestamp"))
# logger.debug(
# f"Request {i + 1} (session={session_id}) - Vertex {vertex_id}: {len(vertex_builds)} builds "
# f"(max allowed: {max_vertex_builds}), "
# f"build_ids: {[b.get('build_id') for b in vertex_builds]}"
# )
assert "job_id" in build_response
assert isinstance(build_response["job_id"], str)
assert uuid.UUID(build_response["job_id"])
# # Wait a bit before final verification to ensure all DB operations complete
# await asyncio.sleep(0.5)
# # Final verification with detailed logging
# async with session_scope() as session:
# vertex_builds = await get_vertex_builds_by_flow_id(db=session, flow_id=flow_id)
# assert len(vertex_builds) > 0, "No vertex builds found"
@pytest.mark.benchmark
async def test_build_flow_polling(client, json_memory_chatbot_no_llm, logged_in_headers):
"""Test the build flow endpoint with polling (non-streaming)."""
# First create the flow
flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers)
# builds_by_vertex = {}
# for build in vertex_builds:
# build_dict = build.model_dump()
# vertex_id = build_dict.get("id")
# builds_by_vertex.setdefault(vertex_id, []).append(build_dict)
# Start the build and get job_id
build_response = await build_flow(client, flow_id, logged_in_headers)
job_id = build_response["job_id"]
assert job_id is not None
# # Log detailed final state
# logger.debug(f"\nFinal state after {num_requests} requests:")
# for vertex_id, builds in builds_by_vertex.items():
# builds.sort(key=lambda x: x.get("timestamp"))
# logger.debug(
# f"Vertex {vertex_id}: {len(builds)} builds "
# f"(oldest: {builds[0].get('timestamp')}, "
# f"newest: {builds[-1].get('timestamp')}), "
# f"build_ids: {[b.get('build_id') for b in builds]}"
# )
# Create a response object that mimics a streaming response but uses polling
class PollingResponse:
def __init__(self, client, job_id, headers):
self.client = client
self.job_id = job_id
self.headers = headers
self.status_code = codes.OK
# # Log individual build details for debugging
# for build in builds:
# logger.debug(
# f" - Build {build.get('build_id')}: timestamp={build.get('timestamp')}, "
# f"valid={build.get('valid')}"
# )
async def aiter_lines(self):
try:
sleeps = 0
max_sleeps = 100
while True:
response = await self.client.get(
f"api/v1/build/{self.job_id}/events?stream=false", headers=self.headers
)
assert response.status_code == codes.OK
data = response.json()
# # Verify each vertex has correct number of builds
# for vertex_id, vertex_builds_list in builds_by_vertex.items():
# assert len(vertex_builds_list) == max_vertex_builds, (
# f"Vertex {vertex_id} has {len(vertex_builds_list)} builds, expected {max_vertex_builds}"
# )
if data["event"] is None:
# No event available, add delay to prevent tight polling
await asyncio.sleep(0.1)
sleeps += 1
continue
# # Verify total number of builds
# total_builds = len(vertex_builds)
# expected_total = max_vertex_builds * num_nodes
# assert total_builds == expected_total, (
# f"Total builds ({total_builds}) doesn't match expected "
# f"({max_vertex_builds} builds/vertex * {num_nodes} nodes = {expected_total})"
# )
# assert all(vertex_build.get("valid") for vertex_build in vertex_builds)
yield data["event"]
# If this was the end event, stop polling
if '"end"' in data["event"]:
break
if sleeps > max_sleeps:
msg = "Build event polling timed out."
raise TimeoutError(msg)
except asyncio.TimeoutError as e:
msg = "Build event polling timed out."
raise TimeoutError(msg) from e
polling_response = PollingResponse(client, job_id, logged_in_headers)
# Use the same consume_and_assert_stream function to verify the events
await consume_and_assert_stream(polling_response, job_id)