diff --git a/aidial_sdk/chat_completion/choice.py b/aidial_sdk/chat_completion/choice.py index d72aae95..7aacc056 100644 --- a/aidial_sdk/chat_completion/choice.py +++ b/aidial_sdk/chat_completion/choice.py @@ -181,13 +181,25 @@ def set_form_schema(self, form_schema: dict) -> None: self._schema_submitted = True self.send_chunk(FormSchemaChunk(self._index, form_schema)) - def create_stage(self, name: str | None = None) -> Stage: + def create_stage( + self, + name: str | None = None, + *, + parent: "Stage | None" = None, + ) -> Stage: if not self._opened: raise runtime_error("Trying to create stage to an unopened choice") if self._closed: raise runtime_error("Trying to create stage to a closed choice") - stage = Stage(self._queue, self._index, self._last_stage_index, name) + parent_stage_index = parent.stage_index if parent is not None else None + stage = Stage( + self._queue, + self._index, + self._last_stage_index, + name, + parent_stage_index, + ) self._last_stage_index += 1 return stage diff --git a/aidial_sdk/chat_completion/chunks.py b/aidial_sdk/chat_completion/chunks.py index 5c6a0404..d19f0d71 100644 --- a/aidial_sdk/chat_completion/chunks.py +++ b/aidial_sdk/chat_completion/chunks.py @@ -190,13 +190,28 @@ class StartStageChunk(BaseChunk): choice_index: int stage_index: int name: str | None + parent_stage_index: int | None - def __init__(self, choice_index: int, stage_index: int, name: str | None): + def __init__( + self, + choice_index: int, + stage_index: int, + name: str | None, + parent_stage_index: int | None = None, + ): self.choice_index = choice_index self.stage_index = stage_index self.name = name + self.parent_stage_index = parent_stage_index def to_dict(self): + stage: dict[str, Any] = { + "index": self.stage_index, + "name": self.name, + "status": None, + } + if self.parent_stage_index is not None: + stage["parent_stage_index"] = self.parent_stage_index return { "choices": [ { @@ -204,13 +219,7 @@ def to_dict(self): "finish_reason": None, "delta": { "custom_content": { - "stages": [ - { - "index": self.stage_index, - "name": self.name, - "status": None, - } - ] + "stages": [stage] } }, } diff --git a/aidial_sdk/chat_completion/stage.py b/aidial_sdk/chat_completion/stage.py index 55a8d39c..ade75b74 100644 --- a/aidial_sdk/chat_completion/stage.py +++ b/aidial_sdk/chat_completion/stage.py @@ -22,6 +22,7 @@ class Stage: _choice_index: int _stage_index: int _name: str | None + _parent_stage_index: int | None _last_attachment_index: int _closed: bool _opened: bool @@ -32,6 +33,7 @@ def __init__( choice_index: int, stage_index: int, name: str | None = None, + parent_stage_index: int | None = None, ): self._queue = queue self._choice_index = choice_index @@ -40,6 +42,11 @@ def __init__( self._opened = False self._closed = False self._name = name + self._parent_stage_index = parent_stage_index + + @property + def stage_index(self) -> int: + return self._stage_index def __enter__(self): self.open() @@ -121,7 +128,12 @@ def open(self): self._opened = True self._queue.put_nowait( - StartStageChunk(self._choice_index, self._stage_index, self._name) + StartStageChunk( + self._choice_index, + self._stage_index, + self._name, + self._parent_stage_index, + ) ) def close(self, status: Status = Status.COMPLETED): diff --git a/tests/test_stage_parent_index.py b/tests/test_stage_parent_index.py new file mode 100644 index 00000000..1a2fc89a --- /dev/null +++ b/tests/test_stage_parent_index.py @@ -0,0 +1,135 @@ +"""Tests for parent_stage_index in StartStageChunk, Stage.stage_index, and Choice.create_stage(parent=...).""" + +import asyncio + +from aidial_sdk.chat_completion import Choice +from aidial_sdk.chat_completion.chunks import StartStageChunk + + +def _drain(queue: asyncio.Queue) -> list: + items = [] + while not queue.empty(): + items.append(queue.get_nowait()) + return items + + +def _stage_dicts_from_queue(queue: asyncio.Queue) -> list: + results = [] + for item in _drain(queue): + d = item.to_dict() + for ch in d.get("choices", []): + for s in ch.get("delta", {}).get("custom_content", {}).get("stages", []): + results.append(s) + return results + + +class TestStartStageChunkParentIndex: + def test_omits_parent_stage_index_when_none(self): + chunk = StartStageChunk(choice_index=0, stage_index=0, name="step") + d = chunk.to_dict() + stage = d["choices"][0]["delta"]["custom_content"]["stages"][0] + assert "parent_stage_index" not in stage + + def test_includes_parent_stage_index_when_set(self): + chunk = StartStageChunk( + choice_index=0, stage_index=1, name="child", parent_stage_index=0 + ) + d = chunk.to_dict() + stage = d["choices"][0]["delta"]["custom_content"]["stages"][0] + assert stage["parent_stage_index"] == 0 + + def test_status_always_none_in_start_chunk(self): + chunk = StartStageChunk( + choice_index=0, stage_index=0, name="x", parent_stage_index=5 + ) + d = chunk.to_dict() + stage = d["choices"][0]["delta"]["custom_content"]["stages"][0] + assert stage["status"] is None + + def test_existing_fields_unchanged(self): + chunk = StartStageChunk(choice_index=0, stage_index=3, name="foo") + d = chunk.to_dict() + stage = d["choices"][0]["delta"]["custom_content"]["stages"][0] + assert stage["index"] == 3 + assert stage["name"] == "foo" + assert stage["status"] is None + + +class TestStageIndex: + def test_stage_index_property_returns_correct_value(self): + choice = Choice(asyncio.Queue(), 0) + choice.open() + s0 = choice.create_stage("first") + s1 = choice.create_stage("second") + assert s0.stage_index == 0 + assert s1.stage_index == 1 + + def test_stage_index_matches_allocation_order(self): + choice = Choice(asyncio.Queue(), 0) + choice.open() + stages = [choice.create_stage(f"s{i}") for i in range(5)] + for i, stage in enumerate(stages): + assert stage.stage_index == i + + +class TestCreateStageWithParent: + def test_create_stage_without_parent_emits_no_parent_stage_index(self): + q = asyncio.Queue() + choice = Choice(q, 0) + choice.open() + _drain(q) # discard StartChoiceChunk + stage = choice.create_stage("root") + stage.open() + stage_dicts = _stage_dicts_from_queue(q) + root_start = next(s for s in stage_dicts if s.get("index") == 0) + assert "parent_stage_index" not in root_start + + def test_create_stage_with_parent_emits_parent_stage_index(self): + q = asyncio.Queue() + choice = Choice(q, 0) + choice.open() + _drain(q) + parent = choice.create_stage("parent") + parent.open() + child = choice.create_stage("child", parent=parent) + child.open() + stage_dicts = _stage_dicts_from_queue(q) + child_start = next(s for s in stage_dicts if s.get("index") == 1) + assert child_start["parent_stage_index"] == 0 + + def test_parent_stage_index_references_parent_index(self): + """When parent has index N, child's parent_stage_index must equal N.""" + q = asyncio.Queue() + choice = Choice(q, 0) + choice.open() + _drain(q) + choice.create_stage("skipped") # index 0 + parent = choice.create_stage("parent") # index 1 + parent.open() + child = choice.create_stage("child", parent=parent) # index 2 + child.open() + stage_dicts = _stage_dicts_from_queue(q) + child_start = next(s for s in stage_dicts if s.get("index") == 2) + assert child_start["parent_stage_index"] == 1 + + def test_stage_open_emits_parent_stage_index_in_wire_format(self): + """End-to-end: opening a child stage produces correct wire-format dict.""" + q = asyncio.Queue() + choice = Choice(q, 0) + choice.open() + _drain(q) + parent = choice.create_stage("parent") + parent.open() + child = choice.create_stage("child", parent=parent) + child.open() + chunks = _drain(q) + # Find the StartStageChunk for child (index=1) + start_dicts = [] + for chunk in chunks: + d = chunk.to_dict() + for ch in d.get("choices", []): + for s in ch.get("delta", {}).get("custom_content", {}).get("stages", []): + if s.get("index") == 1 and s.get("status") is None and "name" in s: + start_dicts.append(s) + assert len(start_dicts) == 1 + assert start_dicts[0]["parent_stage_index"] == 0