feat: supported writable stages/choices (epam#122)

aidial_sdk/chat_completion/choice.py

@@ -19,6 +19,7 @@
 from aidial_sdk.chat_completion.stage import Stage
 from aidial_sdk.pydantic_v1 import ValidationError
 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_debug
 
@@ -91,6 +92,10 @@ def append_content(self, content: str) -> None:
         self.send_chunk(ContentChunk(content, self._index))
         self._last_finish_reason = FinishReason.STOP
 
+    @property
+    def content_stream(self) -> ContentStream:
+        return ContentStream(self)
+
     def create_function_tool_call(
         self, id: str, name: str, arguments: Optional[str] = None
     ) -> FunctionToolCall:

aidial_sdk/chat_completion/stage.py

@@ -13,6 +13,7 @@
 from aidial_sdk.chat_completion.request import Attachment
 from aidial_sdk.pydantic_v1 import ValidationError
 from aidial_sdk.utils._attachment import create_attachment
+from aidial_sdk.utils._content_stream import ContentStream
 from aidial_sdk.utils.errors import runtime_error
 
 
@@ -68,6 +69,10 @@ def append_content(self, content: str):
             ContentStageChunk(self._choice_index, self._stage_index, content)
         )
 
+    @property
+    def content_stream(self) -> ContentStream:
+        return ContentStream(self)
+
     def append_name(self, name: str):
         if not self._opened:
             raise runtime_error("Trying to append name to an unopened stage")

aidial_sdk/utils/_content_stream.py

@@ -0,0 +1,40 @@
+from typing import Protocol
+
+
+class ContentReceiver(Protocol):
+    def append_content(self, content: str) -> None: ...
+
+
+class ContentStream:
+    """
+    The ContentStream class allows using the receiver in contexts where typing.SupportsWrite[str] is expected.
+    For example:
+
+    1. Redirecting print statements:
+
+        print("Hello, world", file=content_stream)
+
+    2. Using with tqdm for progress bars:
+
+        import tqdm
+        for item in tqdm(items, file=content_stream):
+            process(item)
+
+    3. Redirecting logs to the content stream:
+
+        import logging
+        logging_handler = logging.StreamHandler(stream=content_stream)
+
+    4. Writing CSV data:
+
+        import csv
+        csv.writer(content_stream).writerows(data)
+    """
+
+    _receiver: ContentReceiver
+
+    def __init__(self, receiver: ContentReceiver) -> None:
+        self._receiver = receiver
+
+    def write(self, s: str) -> None:
+        self._receiver.append_content(s)