diff --git a/docs/designs/sub-stage-propagation.md b/docs/designs/sub-stage-propagation.md new file mode 100644 index 00000000..b10a3bff --- /dev/null +++ b/docs/designs/sub-stage-propagation.md @@ -0,0 +1,160 @@ +# Sub-Stage Propagation + +**Status:** Approved +**Approved:** 2026-09-04 +**Author:** Aleksei Korota + +--- + +## Problem + +When a QuickApp calls another QuickApp as a tool, the sub-app's UI stages are silently discarded — the user sees only an empty "Calling X" stage with no indication of what the sub-app is actually doing. + +--- + +## Root Cause + +`DialCompletionService._consume_stream()` constructs `ChatStreamConfig` without a `destination` and without `propagate_stages=True`. As a result, `ChoiceUiSink` is instantiated but completely inert — its `on_delta()` returns immediately when `destination is None`. The sub-app's stage deltas are accumulated into `ChatStreamAccumulator.stages` by `AccumulationSink`, but `complete_request_async()` never reads that field, so the stages are effectively thrown away. + +``` +BaseDeploymentTool + └─ DialCompletionService._consume_stream() + └─ ChatStreamConfig(stage_wrapper=...) ← no destination, no propagate_stages + └─ ChoiceUiSink(destination=None) ← inert, sub-app stages discarded +``` + +The sub-app's **text content** does reach the parent (via `StageWrapperUiSink`), so the "Calling X" stage renders the response text. Only stage deltas from the sub-app's `custom_content.stages` are lost. + +--- + +## Approach 1 — Flat Propagation (short-term) + +Sub-app stages are re-emitted into the parent's flat stage list with a name prefix that indicates their origin. No SDK or UI changes required. + +### Mechanism + +1. Inject `Choice` (already request-scoped) into `DialCompletionService`. +2. Pass it as `destination` and set `propagate_stages=True` in `_consume_stream`. +3. Thread a `sub_stage_prefix` string (e.g. `"[WeatherApp]"`) from `BaseDeploymentTool` through `ChatStreamConfig` to `ChoiceUiSink._stream_stage_delta()`, where it is prepended to the stage name on creation. + +For deep nesting (A → B → C), C's stages reach A's stream as siblings of B's stages, all with their respective prefixes. The feature is gated by a new `PreviewField` on `ApplicationConfig` (`orchestrator.propagate_sub_stages`, default `true`) and by the existing `ENABLE_PREVIEW_FEATURES` env switch. + +### Stage index safety + +`ChoiceUiSink` tracks sub-app stage indices in a local `_stages_by_index` dict. Actual allocation is done through `destination.create_stage()`, which increments `Choice._last_stage_index` independently — no collision with the parent's own stages. + +### Files changed + +| File | Change | +|------|--------| +| `config/application.py` | `PreviewField` `orchestrator.propagate_sub_stages: bool` | +| `dial_deployment_tooling/dial_completion_service.py` | Inject `Choice`; pass `destination` + prefix to `_consume_stream` | +| `dial_deployment_tooling/base_deployment_tool.py` | Pass `application_name` as prefix to `complete_request_async` | +| `common/chat_completion_stream/handler.py` | Add `sub_stage_prefix: str | None` to `ChatStreamConfig` | +| `common/chat_completion_stream/chat_stream_sink_factory.py` | Forward `sub_stage_prefix` to `ChoiceUiSink` | +| `common/chat_completion_stream/choice_ui_stream_sink.py` | Prepend prefix in `_stream_stage_delta` | + +### Trade-offs + +**Pros:** No SDK bump, no UI changes, deployable today. +**Cons:** Flat list without structural hierarchy; deep nesting produces a long undifferentiated list. + +--- + +## Approach 2 — Nested Stages (target architecture) + +Sub-app stages are re-emitted as children of the open "Calling X" stage, preserving the call hierarchy at any depth. Requires changes to `aidial_sdk` and a wire-format contract that the UI team must implement. + +### Wire format change (`aidial_sdk`) + +A new optional field `parent_stage_index` is added to `StartStageChunk`. When present, the UI renders the stage as a child of the referenced stage rather than a top-level sibling. + +**Before:** +```json +{"custom_content": {"stages": [ + {"index": 0, "name": "Calling WeatherApp", "status": null}, + {"index": 1, "name": "Fetching forecast", "status": null} +]}} +``` + +**After:** +```json +{"custom_content": {"stages": [ + {"index": 0, "name": "Calling WeatherApp", "status": null}, + {"index": 1, "name": "Fetching forecast", "parent_stage_index": 0, "status": null} +]}} +``` + +`parent_stage_index` is omitted for top-level stages (backwards-compatible). `FinishStageChunk`, `ContentStageChunk`, `NameStageChunk`, and `AttachmentStageChunk` are identified by `index` alone and require no changes. + +### SDK API + +Two additions to `aidial_sdk`: + +```python +# Option A — factory method on Stage +child = parent_stage.create_child_stage("Fetching forecast") + +# Option B — optional parameter on Choice.create_stage +child = choice.create_stage("Fetching forecast", parent=parent_stage) +``` + +`Stage._stage_index` is exposed as a read-only property `Stage.stage_index`. `StartStageChunk.__init__` accepts `parent_stage_index: int | None = None` and includes it in `to_dict()` via `remove_nones`. + +### Mechanism (backend) + +1. `BaseStageWrapper` adds `@property stage -> Stage` to expose its private `self.__stage`. +2. `BaseDeploymentTool._run_in_stage_async()` passes `stage_wrapper.stage` to `complete_request_async()` as a new `parent_stage: Stage | None` parameter. +3. `DialCompletionService._consume_stream()` forwards `parent_stage` in `ChatStreamConfig`. +4. `ChoiceUiSink` receives `parent_stage` and uses it when re-emitting sub-app stages: + - Top-level sub-stages (no `parent_stage_index` in delta): created with `parent=parent_stage`. + - Nested sub-stages (`parent_stage_index` present): remapped through the local `_stages_by_index` dict to find the `Stage` object created in the parent context, then passed as `parent=`. + +The `_stages_by_index` dict (already maintained by `ChoiceUiSink`) naturally handles arbitrary recursion depth: each level's index space is remapped independently as it is processed. + +The feature is gated identically to Approach 1 (`PreviewField` + `ENABLE_PREVIEW_FEATURES`). + +### UI contract (for the UI team) + +The UI must: +- Accept `parent_stage_index: number | undefined` on each stage object in `custom_content.stages`. +- Build a tree from stage deltas as they arrive in the stream: when `parent_stage_index` is set, attach the stage as a child of the already-opened stage at that index. +- Render nested stages as collapsible sections inside their parent stage. +- Treat absent `parent_stage_index` as a top-level stage (no behaviour change for existing flat streams). + +### Files changed + +**`aidial_sdk`:** + +| File | Change | +|------|--------| +| `chat_completion/chunks.py` | `StartStageChunk`: add `parent_stage_index: int | None`; update `to_dict` | +| `chat_completion/stage.py` | Add `parent_stage_index` param; expose `stage_index` property | +| `chat_completion/choice.py` | `create_stage()`: add `parent: Stage | None = None` | + +**`quickapps-backend`:** + +| File | Change | +|------|--------| +| `config/application.py` | `PreviewField` `orchestrator.propagate_sub_stages: bool` | +| `common/_stage_delta_types.py` | Add `parent_stage_index` to `StageDeltaItem` TypedDict | +| `common/base_stage_wrapper.py` | Add `@property stage -> Stage` | +| `dial_deployment_tooling/base_deployment_tool.py` | Pass `stage_wrapper.stage` as `parent_stage` | +| `dial_deployment_tooling/dial_completion_service.py` | Add `parent_stage` param; forward to `_consume_stream` | +| `common/chat_completion_stream/handler.py` | Add `parent_stage: Stage | None` to `ChatStreamConfig` | +| `common/chat_completion_stream/chat_stream_sink_factory.py` | Forward `parent_stage` to `ChoiceUiSink` | +| `common/chat_completion_stream/choice_ui_stream_sink.py` | Index remapping + `parent=` on `create_stage` | + +### Trade-offs + +**Pros:** Correct structural representation; scales cleanly to arbitrary nesting depth. +**Cons:** Requires SDK version bump and UI-team implementation of tree rendering. + +--- + +## Scale of changes + +| | `aidial_sdk` | `quickapps-backend` | +|---|---|---| +| Approach 1 | — | 6 files, ~40 lines | +| Approach 2 | 3 files, ~40 lines | 8 files, ~80 lines | diff --git a/docs/generated-app-schema.json b/docs/generated-app-schema.json index 9e2b66c7..94518a54 100644 --- a/docs/generated-app-schema.json +++ b/docs/generated-app-schema.json @@ -3188,6 +3188,20 @@ "$ref": "#/$defs/StageDisplayLevel", "default": "info", "description": "Threshold for stage visibility. none=no stages at all; errors=failures only; info=user-facing (default); debug=all." + }, + "propagate_sub_stages": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": null, + "description": "When True (default when preview features are enabled), stages emitted by sub-apps called as tools are propagated as nested children of the 'Calling X' stage. Set to False to disable.", + "title": "Propagate Sub Stages", + "x-preview": true } }, "title": "StageDisplayConfig", diff --git a/src/quickapp/common/_stage_delta_types.py b/src/quickapp/common/_stage_delta_types.py index 8d659732..f7ef6b83 100644 --- a/src/quickapp/common/_stage_delta_types.py +++ b/src/quickapp/common/_stage_delta_types.py @@ -17,6 +17,7 @@ class StageDeltaItem(TypedDict, total=False): content: str attachments: list[dict[str, Any]] status: str + parent_stage_index: int def as_stage_delta(item: dict[str, Any]) -> StageDeltaItem: diff --git a/src/quickapp/common/base_stage_wrapper.py b/src/quickapp/common/base_stage_wrapper.py index 91841f93..f7c5909a 100644 --- a/src/quickapp/common/base_stage_wrapper.py +++ b/src/quickapp/common/base_stage_wrapper.py @@ -31,6 +31,10 @@ def __init__( extract_parameters_config_map(tool_config) ) + @property + def stage(self) -> Stage: + return self.__stage + def __enter__(self) -> "BaseStageWrapper": if self.__already_open: # Stage was opened while tool-call arguments streamed; skip re-open. diff --git a/src/quickapp/common/chat_completion_stream/accumulation_stream_sink.py b/src/quickapp/common/chat_completion_stream/accumulation_stream_sink.py index 1ef91f1f..9998d988 100644 --- a/src/quickapp/common/chat_completion_stream/accumulation_stream_sink.py +++ b/src/quickapp/common/chat_completion_stream/accumulation_stream_sink.py @@ -9,9 +9,8 @@ class AccumulationSink(ChatStreamSink): """Always active: builds the in-memory stream result for history / execute / logs.""" - def __init__(self, accumulator: ChatStreamAccumulator, *, stream_content: bool = True) -> None: + def __init__(self, accumulator: ChatStreamAccumulator) -> None: self._accumulator = accumulator - self._stream_content = stream_content def on_stream_start(self) -> None: return @@ -25,7 +24,7 @@ def on_delta(self, delta: NormalizedChoiceDelta) -> None: if norm.state is not None: self._accumulator.merge_state(norm.state) - if delta.content and self._stream_content: + if delta.content: self._accumulator.append_content(delta.content) for tool_call in delta.tool_calls: diff --git a/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py b/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py index f03f9cb0..2b34bfe2 100644 --- a/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py +++ b/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py @@ -81,13 +81,14 @@ def create(self, config: ChatStreamConfig) -> ChatStreamPipeline: return ChatStreamPipeline( accumulator, [ - AccumulationSink(accumulator, stream_content=config.stream_content), + AccumulationSink(accumulator), ChoiceUiSink( accumulator, destination=config.destination, stream_content=config.stream_content, propagate_stages=config.propagate_stages, tools_by_name=_tools_by_name(self._tools()), + parent_stage=config.parent_stage, ), StageWrapperUiSink( stage_wrapper=config.stage_wrapper, diff --git a/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py b/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py index 0f0da2d7..60706d60 100644 --- a/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py +++ b/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py @@ -59,6 +59,7 @@ def __init__( stream_content: bool = True, propagate_stages: bool = False, tools_by_name: dict[str, StagedBaseTool] | None = None, + parent_stage: Stage | None = None, ) -> None: self._accumulator = accumulator self._destination = destination @@ -68,6 +69,7 @@ def __init__( self._stages_by_index: dict[int, Stage] = {} self._tool_stages_by_index: dict[int, _StreamingToolStageState] = {} self._suppressed_tool_indexes: set[int] = set() + self._parent_stage: Stage | None = parent_stage def on_stream_start(self) -> None: destination = self._destination @@ -114,7 +116,7 @@ def on_stream_failure(self) -> None: def _apply_custom(self, norm: NormalizedCustomContent) -> None: destination = self._destination assert destination is not None - if norm.attachments: + if norm.attachments and self._parent_stage is None: self._add_attachments(destination, norm.attachments) for position, raw in norm.stage_entries: stage_delta = as_stage_delta(raw) @@ -314,6 +316,23 @@ def _publish_adopted_tool_stages(self) -> None: ) self._tool_stages_by_index.pop(index, None) + def _resolve_parent(self, delta: StageDeltaItem) -> Stage | None: + """Resolve the parent Stage for a new sub-stage, or None for flat behaviour.""" + if self._parent_stage is None: + return None + sub_parent_idx = delta.get("parent_stage_index") + if sub_parent_idx is not None: + parent = self._stages_by_index.get(sub_parent_idx) + if parent is None: + logger.warning( + "parent_stage_index=%s not yet seen when creating index=%s; " + "falling back to top-level parent", + sub_parent_idx, + delta.get("index"), + ) + return parent if parent is not None else self._parent_stage + return self._parent_stage + def _stream_stage_delta(self, delta: StageDeltaItem, position: int) -> None: destination = self._destination assert destination is not None @@ -331,8 +350,9 @@ def _stream_stage_delta(self, delta: StageDeltaItem, position: int) -> None: ) log_payload(logger, "Stage delta with missing name: %s", delta) return + parent = self._resolve_parent(delta) try: - stage = destination.create_stage(stage_name) + stage = destination.create_stage(stage_name, parent=parent) # type: ignore[call-arg] stage.open() self._stages_by_index[idx] = stage just_created = True diff --git a/src/quickapp/common/chat_completion_stream/handler.py b/src/quickapp/common/chat_completion_stream/handler.py index 1847e4f6..8fcf8754 100644 --- a/src/quickapp/common/chat_completion_stream/handler.py +++ b/src/quickapp/common/chat_completion_stream/handler.py @@ -1,7 +1,7 @@ import logging from collections.abc import AsyncIterable -from aidial_sdk.chat_completion import Choice +from aidial_sdk.chat_completion import Choice, Stage from injector import inject from openai import APIError, BadRequestError from openai.types.chat import ChatCompletionChunk @@ -36,6 +36,7 @@ class ChatStreamConfig(BaseModel): stage_wrapper: BaseStageWrapper | None = None stream_content: bool = True propagate_stages: bool = False + parent_stage: Stage | None = None class ChatCompletionStreamHandler: diff --git a/src/quickapp/config/application.py b/src/quickapp/config/application.py index e1b7b280..d23b9a9d 100644 --- a/src/quickapp/config/application.py +++ b/src/quickapp/config/application.py @@ -42,6 +42,14 @@ class StageDisplayConfig(BaseModel): default=StageDisplayLevel.INFO, description="Threshold for stage visibility. none=no stages at all; errors=failures only; info=user-facing (default); debug=all.", ) + propagate_sub_stages: bool | None = PreviewField( # type: ignore[assignment] + default=None, + description=( + "When True (default when preview features are enabled), stages emitted by " + "sub-apps called as tools are propagated as nested children of the " + "'Calling X' stage. Set to False to disable." + ), + ) def get_max_iterations() -> int: diff --git a/src/quickapp/dial_deployment_tooling/base_deployment_tool.py b/src/quickapp/dial_deployment_tooling/base_deployment_tool.py index be37cd88..ba675639 100644 --- a/src/quickapp/dial_deployment_tooling/base_deployment_tool.py +++ b/src/quickapp/dial_deployment_tooling/base_deployment_tool.py @@ -20,6 +20,7 @@ from quickapp.common.payload_logging import log_payload from quickapp.common.perf_timer.perf_timer import PerformanceTimer from quickapp.common.utils import to_plain_dict +from quickapp.config.application import ApplicationConfig from quickapp.config.dial_deployment import DialDeploymentParameters from quickapp.config.tools.base import ConfigurableSchemaSimpleType, JsonTypeEnum, OpenAiToolConfig from quickapp.config.tools.deployment import ContentPropagation, DialDeploymentTool @@ -50,6 +51,7 @@ def __init__( messages_mixin: MessagesMixin, perf_timer: PerformanceTimer, stage_wrapper_builder: AssistedBuilder[DeploymentStageWrapper], + app_config: ApplicationConfig, argument_transformers: list[ToolArgumentTransformer] | None = None, **kwargs: Any, ): @@ -65,6 +67,7 @@ def __init__( self.__dial_completion_service: DialCompletionService = dial_completion_service self.__attachment_resolver: AttachmentResolver = attachment_resolver self.__content_propagation: ContentPropagation | None = content_propagation + self.__app_config: ApplicationConfig = app_config if content_propagation and content_propagation.propagate_history: logger.warning( "The 'propagate_history' parameter is deprecated and will be removed in a future release. " @@ -123,6 +126,11 @@ async def _run_in_stage_async( tool_config = cast(DialDeploymentTool, self.tool_config) session_id, is_first_call = self._setup_session(kwargs, tool_config, tool_call_id) history = await self._resolve_history(tool_config, session_id) + stage_display = ( + self.__app_config.features.stage_display if self.__app_config.features else None + ) + propagate = stage_display.propagate_sub_stages is not False if stage_display else True + parent_stage = stage_wrapper.stage if (propagate and stage_wrapper is not None) else None result = await self.__dial_completion_service.complete_request_async( kwargs, self.__application_id, @@ -131,6 +139,7 @@ async def _run_in_stage_async( attachment_urls, history=history, supports_url_attachments=tool_config.supports_url_attachments, + parent_stage=parent_stage, ) if is_first_call and session_id: result.content = result.content + f"\n\n[session_id: {session_id}]" diff --git a/src/quickapp/dial_deployment_tooling/deployment_tool.py b/src/quickapp/dial_deployment_tooling/deployment_tool.py index ca2c4731..f60a2b39 100644 --- a/src/quickapp/dial_deployment_tooling/deployment_tool.py +++ b/src/quickapp/dial_deployment_tooling/deployment_tool.py @@ -4,7 +4,7 @@ from quickapp.common.abstract.base_tool_argument_transformer import ToolArgumentTransformer from quickapp.common.messages_mixin import MessagesMixin from quickapp.common.perf_timer.perf_timer import PerformanceTimer -from quickapp.config.application import StageDisplayLevel +from quickapp.config.application import ApplicationConfig, StageDisplayLevel from quickapp.config.tools.deployment import ContentPropagation, DialDeploymentTool from ._attachment_resolver import AttachmentResolver @@ -32,6 +32,7 @@ def __init__( messages_mixin: MessagesMixin, stage_wrapper_builder: AssistedBuilder[DeploymentStageWrapper], perf_timer: PerformanceTimer, + app_config: ApplicationConfig, stage_display_level: StageDisplayLevel = StageDisplayLevel.INFO, argument_transformers: list[ToolArgumentTransformer] | None = None, ): @@ -46,6 +47,7 @@ def __init__( stage_wrapper_builder=stage_wrapper_builder, description=description, perf_timer=perf_timer, + app_config=app_config, stage_display_level=stage_display_level, argument_transformers=argument_transformers, ) diff --git a/src/quickapp/dial_deployment_tooling/dial_completion_service.py b/src/quickapp/dial_deployment_tooling/dial_completion_service.py index 1dba85d8..24263c44 100644 --- a/src/quickapp/dial_deployment_tooling/dial_completion_service.py +++ b/src/quickapp/dial_deployment_tooling/dial_completion_service.py @@ -7,6 +7,7 @@ CustomContentParam, UserMessageParam, ) +from aidial_sdk.chat_completion import Choice, Stage from injector import inject from openai.types.chat import ChatCompletionChunk @@ -42,12 +43,14 @@ def __init__( stream_handler: ChatCompletionStreamHandler, timeout_resolver: ToolTimeoutResolver, attachment_resolver: AttachmentResolver, + choice: Choice, ) -> None: self.__azure_client = azure_client self.__forwarded_headers: ForwardedHeaders = forwarded_headers self.__stream_handler = stream_handler self.__timeout_resolver: ToolTimeoutResolver = timeout_resolver self.__attachment_resolver = attachment_resolver + self.__choice = choice async def complete_request_async( self, @@ -58,6 +61,7 @@ async def complete_request_async( relative_attachment_urls: list[str] | None = None, history: list[UserMessageParam | AssistantMessageParam] | None = None, supports_url_attachments: bool = False, + parent_stage: Stage | None = None, ) -> ToolCallResult: # Expect params to be pre-processed by BaseDeploymentTool._pre_process_params content = params.get(CONTENT_PARAM, "") @@ -77,7 +81,7 @@ async def complete_request_async( params, deployment_id, messages, self.__forwarded_headers ) chunks = await self.__azure_client.chat.completions.create(**chat_params) - result = await self._consume_stream(chunks, stage_wrapper) + result = await self._consume_stream(chunks, stage_wrapper, parent_stage) return ToolCallResult( content=result.content, @@ -121,11 +125,22 @@ async def _consume_stream( self, chunks: AsyncIterable[ChatCompletionChunk], stage_wrapper: BaseStageWrapper | None, + parent_stage: Stage | None = None, ) -> ChatStreamAccumulator: + if parent_stage is not None: + config = ChatStreamConfig( + stage_wrapper=stage_wrapper, + destination=self.__choice, + propagate_stages=True, + stream_content=False, + parent_stage=parent_stage, + ) + else: + config = ChatStreamConfig(stage_wrapper=stage_wrapper) try: return await self.__stream_handler.process_stream( chunks=chunks, - config=ChatStreamConfig(stage_wrapper=stage_wrapper), + config=config, ) except ChatStreamHandlerError: logger.exception("Deployment stream handling failed.") diff --git a/src/tests/unit_tests/dial_deployment_tooling_tests/test_base_deployment_tool.py b/src/tests/unit_tests/dial_deployment_tooling_tests/test_base_deployment_tool.py index 60110552..86df32a2 100644 --- a/src/tests/unit_tests/dial_deployment_tooling_tests/test_base_deployment_tool.py +++ b/src/tests/unit_tests/dial_deployment_tooling_tests/test_base_deployment_tool.py @@ -105,6 +105,7 @@ def _build_tool( perf_timer=MagicMock(), stage_wrapper_builder=MagicMock(), stage_display_level=StageDisplayLevel.INFO, + app_config=MagicMock(), ) @@ -406,6 +407,7 @@ def _build_tool_with_config( perf_timer=MagicMock(), stage_wrapper_builder=MagicMock(), stage_display_level=StageDisplayLevel.INFO, + app_config=MagicMock(), ) @@ -518,6 +520,7 @@ def _build_tool_with_propagation( perf_timer=MagicMock(), stage_wrapper_builder=MagicMock(), stage_display_level=StageDisplayLevel.INFO, + app_config=MagicMock(), ) @@ -778,6 +781,7 @@ def _build_tool_with_content_propagation( perf_timer=MagicMock(), stage_wrapper_builder=MagicMock(), stage_display_level=StageDisplayLevel.INFO, + app_config=MagicMock(), ) diff --git a/src/tests/unit_tests/dial_deployment_tooling_tests/test_completion_service.py b/src/tests/unit_tests/dial_deployment_tooling_tests/test_completion_service.py index bdd9e1cd..9191a901 100644 --- a/src/tests/unit_tests/dial_deployment_tooling_tests/test_completion_service.py +++ b/src/tests/unit_tests/dial_deployment_tooling_tests/test_completion_service.py @@ -51,6 +51,7 @@ def completion_service(azure_client, attachment_resolver): stream_handler=ChatCompletionStreamHandler.with_default_sinks(), timeout_resolver=noop_timeout_resolver(), attachment_resolver=attachment_resolver, + choice=MagicMock(), ) @@ -294,6 +295,7 @@ async def test_forwarded_x_headers_passed_to_chat_completion( stream_handler=ChatCompletionStreamHandler.with_default_sinks(), timeout_resolver=noop_timeout_resolver(), attachment_resolver=attachment_resolver, + choice=MagicMock(), ) await service.complete_request_async( diff --git a/src/tests/unit_tests/stream_test_doubles.py b/src/tests/unit_tests/stream_test_doubles.py index b27f2d69..b0d40ff7 100644 --- a/src/tests/unit_tests/stream_test_doubles.py +++ b/src/tests/unit_tests/stream_test_doubles.py @@ -37,7 +37,7 @@ def set_state(self, state: Any) -> None: self.set_state_calls.append(state) return super().set_state(state) - def create_stage(self, name: str | None = None) -> Stage: + def create_stage(self, name: str | None = None, *, parent: Stage | None = None) -> Stage: stage = super().create_stage(name) self.created_stages.append(stage) return stage