chore: chunk processing refactoring

pyproject.toml

@@ -154,6 +154,8 @@ priority = "primary"
 
 [tool.pytest.ini_options]
 testpaths = ["src/tests"]
+pythonpath = ["src"]
+env_files = [".env"]
 asyncio_default_fixture_loop_scope = "session"
 addopts = "-ra"
 markers = [

src/quickapp/agent/agent_module.py

@@ -9,13 +9,18 @@
 from quickapp.agent._prompt_providers import ConfigBasedPromptProvider
 from quickapp.agent.agent_settings import AgentSettings
 from quickapp.agent.assistant_invoker import AssistantInvoker
-from quickapp.agent.chunk_processor import ChunkProcessor
 from quickapp.agent.models import OpenAiToolConfigDict
 from quickapp.agent.orchestrator import Orchestrator
-from quickapp.common import DIAL_API_KEY, ForwardedHeaders, StagedBaseTool
+from quickapp.common import (
+    DIAL_API_KEY,
+    ORCHESTRATOR_AZURE_CLIENT,
+    ForwardedHeaders,
+    StagedBaseTool,
+)
 from quickapp.common.abstract.base_prompt_provider import PromptPartProvider
 from quickapp.common.abstract.base_transformer import MessagesTransformer, PreInvocationTransformer
 from quickapp.common.abstract.completion_result_enricher import CompletionResultEnricher
+from quickapp.common.chat_completion_stream.handler import ChatCompletionStreamHandler
 from quickapp.common.dial_settings import DialSettings
 from quickapp.common.state_holder import StateHolder
 from quickapp.common.utils import sanitize_toolname
@@ -57,7 +62,7 @@ def configure(self, binder: Binder) -> None:
         binder.bind(Orchestrator, to=Orchestrator)  # type: ignore[type-abstract]
         binder.bind(StateHolder, to=StateHolder, scope=request_scope)
         binder.bind(AssistantInvoker, to=AssistantInvoker, scope=NoScope)
-        binder.bind(ChunkProcessor, to=ChunkProcessor, scope=NoScope)
+        binder.bind(ChatCompletionStreamHandler, to=ChatCompletionStreamHandler, scope=NoScope)
         binder.bind(_AttachmentFilter, to=_AttachmentFilter, scope=request_scope)
         binder.bind(
             _AddSystemPromptTransformer, to=_AddSystemPromptTransformer, scope=request_scope
@@ -72,7 +77,7 @@ def provide_openai_client(
         api_key: DIAL_API_KEY,
         config: ApplicationConfig,
         forwarded_headers: ForwardedHeaders,
-    ) -> AsyncAzureOpenAI:
+    ) -> ORCHESTRATOR_AZURE_CLIENT:
         azure_client = AsyncAzureOpenAI(
             azure_endpoint=dial_settings.url,
             api_key=api_key.get_secret_value(),

src/quickapp/agent/assistant_invoker.py

@@ -6,13 +6,12 @@
 from aidial_sdk.chat_completion.request import Message
 from injector import inject
 from openai import AsyncStream
-from openai.lib.azure import AsyncAzureOpenAI
 from openai.types.chat import ChatCompletionChunk
 
 from quickapp.agent.agent_settings import AgentSettings
 from quickapp.agent.message_logger import format_openai_message_pipe_tree
 from quickapp.agent.models import STATE_KEY_ORCHESTRATOR, OpenAiToolConfigDict
-from quickapp.common import RESPONSE_FORMAT, ForwardedHeaders
+from quickapp.common import ORCHESTRATOR_AZURE_CLIENT, RESPONSE_FORMAT, ForwardedHeaders
 from quickapp.common.abstract.base_transformer import PreInvocationTransformer
 from quickapp.common.presentation_settings import PresentationSettings
 from quickapp.config.application import ApplicationConfig
@@ -28,7 +27,7 @@ def __init__(
         config: ApplicationConfig,
         messages: list[Message],
         choice: Choice,
-        azure_client: AsyncAzureOpenAI,
+        azure_client: ORCHESTRATOR_AZURE_CLIENT,
         response_format: RESPONSE_FORMAT,
         pre_invocation_transformers: list[PreInvocationTransformer],
         presentation_settings: PresentationSettings,

src/quickapp/agent/orchestrator.py

@@ -5,12 +5,20 @@
 from aidial_sdk.chat_completion import Choice
 from aidial_sdk.chat_completion.request import CustomContent, Message, Role
 from injector import ProviderOf, inject
+from openai import AsyncStream
+from openai.types.chat import ChatCompletionChunk
 
 from quickapp.agent.assistant_invoker import AssistantInvoker
-from quickapp.agent.chunk_processor import ChunkProcessor
 from quickapp.agent.models import STATE_KEY_ORCHESTRATOR, TOOL_EXECUTION_HISTORY
 from quickapp.agent.tool_executor import ToolExecutor
 from quickapp.common import DeploymentUsage
+from quickapp.common.chat_completion_stream.exceptions import ChatStreamHandlerError
+from quickapp.common.chat_completion_stream.handler import (
+    ChatCompletionStreamHandler,
+    ChatStreamConfig,
+)
+from quickapp.common.chat_completion_stream.stream_result import ChatStreamAccumulator
+from quickapp.common.chat_completion_stream.tool_call import AccumulatedToolCall
 from quickapp.common.exceptions import OrchestratorExceedMaxIterationsException
 from quickapp.common.messages_mixin import MessagesMixin
 from quickapp.common.perf_timer.perf_timer import PerformanceTimer
@@ -19,8 +27,6 @@
 from quickapp.config.application import ApplicationConfig
 from quickapp.usage_statistics.usage_statistics_service import UsageStatisticsService
 
-from ._models import AccumulatedToolCall
-
 logger = logging.getLogger(__name__)
 
 
@@ -36,7 +42,7 @@ def __init__(
         usage_statistics_service: UsageStatisticsService,
         tool_executor: ToolExecutor,
         assistant_invoker_provider: ProviderOf[AssistantInvoker],
-        chunk_processor_provider: ProviderOf[ChunkProcessor],
+        stream_handler_provider: ProviderOf[ChatCompletionStreamHandler],
         app_config: ApplicationConfig,
         perf_timer: PerformanceTimer,
     ) -> None:
@@ -47,7 +53,7 @@ def __init__(
         self.__SHOW_USAGE_STATISTICS = presentation_settings.show_usage_statistics
         self.__tool_executor = tool_executor
         self.__assistant_invoker_provider = assistant_invoker_provider
-        self.__chunk_processor_provider = chunk_processor_provider
+        self.__stream_handler_provider = stream_handler_provider
         self.__iterations_counter = 0
         self.__MAX_ITERATIONS_COUNT = app_config.orchestrator.max_iterations
         self.__orchestrator_deployment_name = app_config.orchestrator.deployment.name
@@ -96,24 +102,18 @@ async def _run_iteration(self) -> bool:
 
         assistant_invoker = self.__assistant_invoker_provider.get()
         chat_completion_stream = await assistant_invoker.invoke()
-        assistant_call_result = await self.__chunk_processor_provider.get().process_chunks(
-            chat_completion=chat_completion_stream,
-            destination=self.__choice,
-            propagate_orchestrator_stages=self.__propagate_orchestrator_stages,
-        )
-        if assistant_call_result is None:
-            raise RuntimeError("Assistant invocation returned no result.")
+        stream_result = await self.accumulate_stream(chat_completion_stream)
 
-        tool_calls = assistant_call_result.tool_calls
+        tool_calls = stream_result.tool_calls
 
         # Thinking stages stay in custom_content (streamed to choice), not in state.
-        stream_state = dict(assistant_call_result.state or {})
+        response_state = dict(stream_result.state or {})
         state: dict[str, object] | None = (
-            {STATE_KEY_ORCHESTRATOR: stream_state} if stream_state else None
+            {STATE_KEY_ORCHESTRATOR: response_state} if response_state else None
         )
 
         custom_content_kwargs: dict[str, object] = {
-            "attachments": assistant_call_result.attachments,
+            "attachments": stream_result.attachments,
         }
         if state:
             custom_content_kwargs["state"] = state
@@ -124,18 +124,18 @@ async def _run_iteration(self) -> bool:
         self.__messages_context.append_message(
             Message(
                 role=Role.ASSISTANT,
-                content=assistant_call_result.content or " ",
+                content=stream_result.content or " ",
                 custom_content=CustomContent(**custom_content_kwargs),
                 tool_calls=AccumulatedToolCall.to_sdk_tool_calls(tool_calls),
             )
         )
 
-        if assistant_call_result.usage and self.__SHOW_USAGE_STATISTICS:
+        if stream_result.usage and self.__SHOW_USAGE_STATISTICS:
             self.__usage_statistics_list.append(
                 DeploymentUsage(
                     model_name=self.__orchestrator_deployment_name,
-                    prompt_tokens=assistant_call_result.usage.prompt_tokens,
-                    completion_tokens=assistant_call_result.usage.completion_tokens,
+                    prompt_tokens=stream_result.usage.prompt_tokens,
+                    completion_tokens=stream_result.usage.completion_tokens,
                 )
             )
         self.__perf_timer.add_milestone(period, "assistant_response_received")
@@ -162,6 +162,25 @@ async def _run_iteration(self) -> bool:
         logger.debug(f"Message from context: {self.__messages_context.messages}")
         return True
 
+    async def accumulate_stream(
+        self, chat_completion_stream: AsyncStream[ChatCompletionChunk]
+    ) -> ChatStreamAccumulator:
+        try:
+            stream_accumulator = await self.__stream_handler_provider.get().process_stream(
+                chunks=chat_completion_stream,
+                config=ChatStreamConfig(
+                    destination=self.__choice,
+                    stream_content=True,
+                    propagate_stages=self.__propagate_orchestrator_stages,
+                ),
+            )
+        except ChatStreamHandlerError:
+            logger.exception("Orchestrator stream handling failed.")
+            raise
+        if stream_accumulator is None:
+            raise RuntimeError("Assistant invocation returned no result.")
+        return stream_accumulator
+
     def _build_tool_execution_history(self) -> list[dict[str, object]]:
         """Build tool execution history by extracting ASSISTANT and TOOL messages.
 

src/quickapp/agent/tool_executor.py

@@ -4,9 +4,9 @@
 
 from injector import inject
 
-from quickapp.agent._models import AccumulatedToolCall
 from quickapp.common import CompletionResult, StagedBaseTool
 from quickapp.common.abstract.completion_result_enricher import CompletionResultEnricher
+from quickapp.common.chat_completion_stream.tool_call import AccumulatedToolCall
 from quickapp.common.perf_timer.perf_timer import PerformanceTimer
 from quickapp.common.utils import sanitize_toolname