aidial-sdk 0.44.0.dev5__tar.gz → 0.45.0.dev1__tar.gz

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aidial-sdk
-Version: 0.44.0.dev5
+Version: 0.45.0.dev1
 Summary: Framework to create applications and model adapters for AI DIAL
 License-Expression: Apache-2.0
 License-File: LICENSE

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/aidial_sdk/chat_completion/choice.py

@@ -195,12 +195,25 @@ class Choice(ChoiceBase):
         self.send_chunk(FormSchemaChunk(self._index, form_schema))
 
     def create_stage(self, name: str | None = None) -> Stage:
+        return self._create_stage(name, parent=None)
+
+    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)
+        stage = Stage(
+            self,
+            self._last_stage_index,
+            name,
+            parent,
+        )
         self._last_stage_index += 1
 
         return stage

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/aidial_sdk/chat_completion/chunks.py

@@ -211,29 +211,34 @@ 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": [
                 {
                     "index": self.choice_index,
                     "finish_reason": None,
-                    "delta": {
-                        "custom_content": {
-                            "stages": [
-                                {
-                                    "index": self.stage_index,
-                                    "name": self.name,
-                                    "status": None,
-                                }
-                            ]
-                        }
-                    },
+                    "delta": {"custom_content": {"stages": [stage]}},
                 }
             ],
             "usage": None,

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/aidial_sdk/chat_completion/request.py

@@ -46,6 +46,7 @@ class Stage(ExtraAllowModel, IgnoreIndex):
     status: Status
     content: StrictStr | None = None
     attachments: list[Attachment] | None = None
+    parent_stage_index: StrictInt | None = None
 
 
 class CustomContent(ExtraAllowModel):

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/aidial_sdk/chat_completion/stage.py

@@ -1,5 +1,5 @@
 from types import TracebackType
-from typing import overload
+from typing import TYPE_CHECKING, overload
 
 from aidial_sdk._pydantic import ValidationError
 from aidial_sdk.chat_completion._types import ChunkQueue
@@ -15,31 +15,52 @@ from aidial_sdk.chat_completion.request import Attachment
 from aidial_sdk.utils._attachment import create_attachment
 from aidial_sdk.utils._content_stream import ContentStream
 from aidial_sdk.utils.errors import runtime_error
+from aidial_sdk.utils.logging import log_warning
+
+if TYPE_CHECKING:
+    from aidial_sdk.chat_completion.choice import Choice
 
 
 class Stage:
+    _choice: "Choice"
     _queue: ChunkQueue
     _choice_index: int
     _stage_index: int
     _name: str | None
+    _parent: "Stage | None"
+    _children: list["Stage"]
     _last_attachment_index: int
     _closed: bool
     _opened: bool
 
     def __init__(
         self,
-        queue: ChunkQueue,
-        choice_index: int,
+        choice: "Choice",
         stage_index: int,
         name: str | None = None,
+        parent: "Stage | None" = None,
     ):
-        self._queue = queue
-        self._choice_index = choice_index
+        if parent is not None and parent._choice is not choice:
+            raise runtime_error(
+                "Trying to create a stage whose parent stage belongs to another choice"
+            )
+
+        self._choice = choice
+        self._queue = choice._queue
+        self._choice_index = choice.index
         self._stage_index = stage_index
         self._last_attachment_index = 0
         self._opened = False
         self._closed = False
         self._name = name
+        self._parent = parent
+        self._children = []
+
+        if parent is not None:
+            parent._children.append(self)
+
+    def create_stage(self, name: str | None = None) -> "Stage":
+        return self._choice._create_stage(name, parent=self)
 
     def __enter__(self):
         self.open()
@@ -119,9 +140,25 @@ class Stage:
         if self._opened:
             raise runtime_error("The stage is already open")
 
+        parent = self._parent
+        if parent is not None:
+            if not parent._opened:
+                raise runtime_error(
+                    "Trying to open a stage whose parent stage is not open"
+                )
+            if parent._closed:
+                raise runtime_error(
+                    "Trying to open a stage whose parent stage is already closed"
+                )
+
         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,
+                parent._stage_index if parent is not None else None,
+            )
         )
 
     def close(self, status: Status = Status.COMPLETED):
@@ -130,6 +167,17 @@ class Stage:
         if self._closed:
             raise runtime_error("The stage is already closed")
 
+        open_children = [
+            child._stage_index
+            for child in self._children
+            if child._opened and not child._closed
+        ]
+        if open_children:
+            log_warning(
+                f"Closing the stage {self._stage_index} which still has "
+                f"open child stages: {open_children}"
+            )
+
         self._closed = True
         self._queue.put_nowait(
             FinishStageChunk(self._choice_index, self._stage_index, status)

{aidial_sdk-0.44.0.dev5 → aidial_sdk-0.45.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-sdk"
-version = "0.44.0.dev5"
+version = "0.45.0.dev1"
 description = "Framework to create applications and model adapters for AI DIAL"
 authors = [{ name = "EPAM RAIL", email = "SpecialEPM-DIALDevTeam@epam.com" }]
 license = "Apache-2.0"