Framework to create applications and model adapters for AI DIAL
from aidial_sdk.application import DIALApp
from aidial_sdk.exceptions import HTTPException
from aidial_sdk.utils.logging import logger

aidial_sdk/__init__.py
epam/ai-dial-sdk
import logging.config
from json import JSONDecodeError
from typing import Dict, Optional

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse

from aidial_sdk.chat_completion.base import ChatCompletion
from aidial_sdk.chat_completion.request import Request as ChatCompletionRequest
from aidial_sdk.chat_completion.response import (
    Response as ChatCompletionResponse,
)
from aidial_sdk.header_propagator import HeaderPropagator
from aidial_sdk.pydantic_v1 import ValidationError
from aidial_sdk.utils.errors import json_error
from aidial_sdk.utils.log_config import LogConfig
from aidial_sdk.utils.logging import log_debug, set_log_deployment
from aidial_sdk.utils.streaming import merge_chunks

logging.config.dictConfig(LogConfig().dict())


class DIALApp(FastAPI):
    chat_completion_impls: Dict[str, ChatCompletion] = {}

    def __init__(
        self,
        dial_url: Optional[str] = None,
        propagation_auth_headers: bool = False,
        **fast_api_kwargs,
    ):
        super().__init__(**fast_api_kwargs)

        if propagation_auth_headers:
            if not dial_url:
                raise ValueError(
                    "dial_url is required if propagation auth headers is enabled"
                )

            HeaderPropagator(self, dial_url).enable()

        self.add_api_route(
            "/openai/deployments/{deployment_id}/chat/completions",
            self._chat_completion,
            methods=["POST"],
        )

    def add_chat_completion(
        self, deployment_name: str, impl: ChatCompletion
    ) -> None:
        self.chat_completion_impls[deployment_name] = impl

    async def _chat_completion(
        self, deployment_id: str, original_request: Request
    ) -> Response:
        set_log_deployment(deployment_id)
        impl = self.chat_completion_impls.get(deployment_id, None)

        if not impl:
            return JSONResponse(
                status_code=404,
                content=json_error(
                    message="The API deployment for this resource does not exist.",
                    code="deployment_not_found",
                ),
            )

        try:
            body = await original_request.json()
        except JSONDecodeError as e:
            return JSONResponse(
                status_code=400,
                content=json_error(
                    message=f"Your request contained invalid JSON: {str(e)}",
                    type="invalid_request_error",
                ),
            )

        headers = original_request.headers
        try:
            request = ChatCompletionRequest(
                **body,
                api_key=headers["Api-Key"],
                jwt=headers.get("Authorization"),
                deployment_id=deployment_id,
                api_version=original_request.query_params.get("api-version"),
                headers=headers,
            )
        except ValidationError as e:
            error = e.errors()[0]
            path = ".".join(map(str, e.errors()[0]["loc"]))
            return JSONResponse(
                status_code=400,
                content=json_error(
                    message=f"Your request contained invalid structure on path {path}. {error['msg']}",
                    type="invalid_request_error",
                ),
            )

        log_debug(f"request: {body}")

        response = ChatCompletionResponse(request)
        first_chunk = await response._generator(impl.chat_completion, request)

        if request.stream:
            return StreamingResponse(
                response._generate_stream(first_chunk),
                media_type="text/event-stream",
            )
        else:
            response_body = await merge_chunks(
                response._generate_stream(first_chunk)
            )

            log_debug(f"response: {response_body}")
            return JSONResponse(content=response_body)

aidial_sdk/application.py
epam/ai-dial-sdk
from aidial_sdk.chat_completion.base import ChatCompletion
from aidial_sdk.chat_completion.choice import Choice
from aidial_sdk.chat_completion.enums import FinishReason, Status
from aidial_sdk.chat_completion.request import (
    Addon,
    Attachment,
    CustomContent,
    Function,
    FunctionCall,
    Message,
    Request,
    Role,
)
from aidial_sdk.chat_completion.response import Response
from aidial_sdk.chat_completion.stage import Stage

aidial_sdk/chat_completion/__init__.py
epam/ai-dial-sdk
from abc import ABC, abstractmethod

from aidial_sdk.chat_completion.request import Request
from aidial_sdk.chat_completion.response import Response


class ChatCompletion(ABC):
    @abstractmethod
    async def chat_completion(
        self, request: Request, response: Response
    ) -> None:
        """Implement chat completion logic"""

aidial_sdk/chat_completion/base.py
epam/ai-dial-sdk
from asyncio import Queue
from types import TracebackType
from typing import Any, Optional, Type

from aidial_sdk.chat_completion.chunks import (
    AttachmentChunk,
    ContentChunk,
    EndChoiceChunk,
    StartChoiceChunk,
    StateChunk,
)
from aidial_sdk.chat_completion.enums import FinishReason
from aidial_sdk.chat_completion.stage import Stage
from aidial_sdk.pydantic_v1 import ValidationError
from aidial_sdk.utils.errors import runtime_error


class Choice:
    _queue: Queue
    _index: int
    _last_attachment_index: int
    _last_stage_index: int
    _opened: bool
    _closed: bool
    _state_submitted: bool

    def __init__(self, queue: Queue, choice_index: int):
        self._queue = queue
        self._index = choice_index
        self._last_attachment_index = 0
        self._last_stage_index = 0
        self._opened = False
        self._closed = False
        self._state_submitted = False

    def __enter__(self):
        self.open()
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        self.close()
        return False

    def append_content(self, content: str) -> None:
        if not self._opened:
            runtime_error("Trying to append content to an unopened choice")
        if self._closed:
            runtime_error("Trying to append content to a closed choice")

        self._queue.put_nowait(ContentChunk(content, self._index))

    def add_attachment(
        self,
        type: Optional[str] = None,
        title: Optional[str] = None,
        data: Optional[str] = None,
        url: Optional[str] = None,
        reference_url: Optional[str] = None,
        reference_type: Optional[str] = None,
    ) -> None:
        if not self._opened:
            runtime_error("Trying to add attachment to an unopened choice")
        if self._closed:
            runtime_error("Trying to add attachment to a closed choice")

        attachment_chunk = None
        try:
            attachment_chunk = AttachmentChunk(
                choice_index=self._index,
                attachment_index=self._last_attachment_index,
                type=type,
                title=title,
                data=data,
                url=url,
                reference_url=reference_url,
                reference_type=reference_type,
            )
        except ValidationError as e:
            runtime_error(e.errors()[0]["msg"])

        self._queue.put_nowait(attachment_chunk)
        self._last_attachment_index += 1

    def set_state(self, state: Any) -> None:
        if self._state_submitted:
            runtime_error('Trying to set "state" twice')

        if not self._opened:
            runtime_error("Trying to append state to an unopened choice")
        if self._closed:
            runtime_error("Trying to append state to a closed choice")

        self._state_submitted = True
        self._queue.put_nowait(StateChunk(self._index, state))

    def create_stage(self, name: Optional[str] = None) -> Stage:
        if not self._opened:
            runtime_error("Trying to create stage to an unopened choice")
        if self._closed:
            runtime_error("Trying to create stage to a closed choice")

        stage = Stage(self._queue, self._index, self._last_stage_index, name)
        self._last_stage_index += 1

        return stage

    def open(self):
        if self._opened:
            runtime_error("The choice is already open")

        self._opened = True
        self._queue.put_nowait(StartChoiceChunk(choice_index=self._index))

    def close(self, finish_reason: FinishReason = FinishReason.STOP) -> None:
        if not self._opened:
            runtime_error("Trying to close an unopened choice")
        if self._closed:
            runtime_error("The choice is already closed")

        self._closed = True
        self._queue.put_nowait(EndChoiceChunk(finish_reason, self._index))

aidial_sdk/chat_completion/choice.py
epam/ai-dial-sdk
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional

from aidial_sdk.chat_completion.enums import FinishReason, Status
from aidial_sdk.pydantic_v1 import BaseModel, root_validator


class BaseChunk(ABC):
    @abstractmethod
    def to_dict(self) -> Dict[str, Any]:
        pass


class StartChoiceChunk(BaseChunk):
    choice_index: int

    def __init__(self, choice_index: int):
        self.choice_index = choice_index

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {"role": "assistant"},
                }
            ],
            "usage": None,
        }


class EndChoiceChunk(BaseChunk):
    finish_reason: FinishReason
    choice_index: int

    def __init__(self, finish_reason: FinishReason, choice_index: int):
        self.finish_reason = finish_reason
        self.choice_index = choice_index

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": self.finish_reason.value,
                    "delta": {},
                }
            ],
            "usage": None,
        }


class ContentChunk(BaseChunk):
    content: str
    choice_index: int

    def __init__(self, content: str, choice_index: int):
        self.content = content
        self.choice_index = choice_index

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {"content": self.content},
                }
            ],
            "usage": None,
        }


class StartStageChunk(BaseChunk):
    choice_index: int
    stage_index: int
    name: Optional[str]

    def __init__(
        self, choice_index: int, stage_index: int, name: Optional[str]
    ):
        self.choice_index = choice_index
        self.stage_index = stage_index
        self.name = name

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "stages": [
                                {
                                    "index": self.stage_index,
                                    "name": self.name,
                                    "status": None,
                                }
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class FinishStageChunk(BaseChunk):
    choice_index: int
    stage_index: int
    status: Status

    def __init__(self, choice_index: int, stage_index: int, status: Status):
        self.choice_index = choice_index
        self.stage_index = stage_index
        self.status = status

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "stages": [
                                {
                                    "index": self.stage_index,
                                    "status": self.status.value,
                                }
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class ContentStageChunk(BaseChunk):
    choice_index: int
    stage_index: int
    content: str

    def __init__(self, choice_index: int, stage_index: int, content: str):
        self.choice_index = choice_index
        self.stage_index = stage_index
        self.content = content

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "stages": [
                                {
                                    "index": self.stage_index,
                                    "content": self.content,
                                    "status": None,
                                }
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class NameStageChunk(BaseChunk):
    choice_index: int
    stage_index: int
    name: str

    def __init__(self, choice_index: int, stage_index: int, name: str):
        self.choice_index = choice_index
        self.stage_index = stage_index
        self.name = name

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "stages": [
                                {
                                    "index": self.stage_index,
                                    "name": self.name,
                                    "status": None,
                                }
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class Attachment(BaseModel):
    choice_index: int
    attachment_index: int

    type: Optional[str]
    title: Optional[str]
    data: Optional[str]
    url: Optional[str]
    reference_url: Optional[str]
    reference_type: Optional[str]

    @root_validator
    def check_data_or_url(cls, values):
        data, url = values.get("data"), values.get("url")

        if not data and not url:
            raise ValueError("Trying to add attachment without data and url")
        if data and url:
            raise ValueError("Trying to add attachment with data and url")

        return values

    def attachment_dict(self, index: int):
        attachment: Dict[str, Any] = {"index": index}

        if self.type:
            attachment["type"] = self.type
        if self.title:
            attachment["title"] = self.title
        if self.data:
            attachment["data"] = self.data
        if self.url:
            attachment["url"] = self.url
        if self.reference_url:
            attachment["reference_url"] = self.reference_url
        if self.reference_type:
            attachment["reference_type"] = self.reference_type

        return attachment


class AttachmentChunk(Attachment, BaseChunk):
    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "attachments": [
                                self.attachment_dict(self.attachment_index)
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class AttachmentStageChunk(Attachment, BaseChunk):
    stage_index: int

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {
                        "custom_content": {
                            "stages": [
                                {
                                    "index": self.stage_index,
                                    "attachments": [
                                        self.attachment_dict(
                                            self.attachment_index
                                        )
                                    ],
                                    "status": None,
                                }
                            ]
                        }
                    },
                }
            ],
            "usage": None,
        }


class StateChunk(BaseChunk):
    choice_index: int
    state: Any

    def __init__(self, choice_index: int, state: Any):
        self.state = state
        self.choice_index = choice_index

    def to_dict(self):
        return {
            "choices": [
                {
                    "index": self.choice_index,
                    "finish_reason": None,
                    "delta": {"custom_content": {"state": self.state}},
                }
            ],
            "usage": None,
        }


class UsageChunk(BaseChunk):
    prompt_tokens: int
    completion_tokens: int

    def __init__(self, prompt_tokens: int, completion_tokens: int):
        self.prompt_tokens = prompt_tokens
        self.completion_tokens = completion_tokens

    def to_dict(self):
        return {
            "usage": {
                "prompt_tokens": self.prompt_tokens,
                "completion_tokens": self.completion_tokens,
                "total_tokens": self.prompt_tokens + self.completion_tokens,
            }
        }


class UsagePerModelChunk(BaseChunk):
    index: int
    model: str
    prompt_tokens: int
    completion_tokens: int

    def __init__(
        self,
        index: int,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
    ):
        self.index = index
        self.model = model
        self.prompt_tokens = prompt_tokens
        self.completion_tokens = completion_tokens

    def to_dict(self):
        return {
            "statistics": {
                "usage_per_model": [
                    {
                        "index": self.index,
                        "model": self.model,
                        "prompt_tokens": self.prompt_tokens,
                        "completion_tokens": self.completion_tokens,
                        "total_tokens": self.prompt_tokens
                        + self.completion_tokens,
                    }
                ]
            }
        }


class EndChunk:
    exc: Optional[Exception]

    def __init__(self, exc: Optional[Exception] = None):
        self.exc = exc

aidial_sdk/chat_completion/chunks.py
epam/ai-dial-sdk
from enum import Enum


class FinishReason(Enum):
    STOP = "stop"
    LENGTH = "length"
    FUNCTION_CALL = "function_call"
    CONTENT_FILTER = "content_filter"


class Status(Enum):
    COMPLETED = "completed"
    FAILED = "failed"

aidial_sdk/chat_completion/enums.py
epam/ai-dial-sdk
from enum import Enum
from typing import Any, List, Mapping, Optional, Union

from aidial_sdk.pydantic_v1 import (
    BaseModel,
    ConstrainedFloat,
    ConstrainedInt,
    ConstrainedList,
    PositiveInt,
    StrictStr,
)


class ExtraForbidModel(BaseModel):
    class Config:
        extra = "forbid"


class Attachment(ExtraForbidModel):
    type: Optional[StrictStr] = "text/markdown"
    title: Optional[StrictStr] = None
    data: Optional[StrictStr] = None
    url: Optional[StrictStr] = None
    reference_type: Optional[StrictStr] = None
    reference_url: Optional[StrictStr] = None


class CustomContent(ExtraForbidModel):
    attachments: Optional[List[Attachment]] = None
    state: Optional[Any] = None


class Role(Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"


class FunctionCall(ExtraForbidModel):
    name: str
    arguments: str


class Message(ExtraForbidModel):
    role: Role
    content: Optional[StrictStr] = None
    custom_content: Optional[CustomContent] = None
    name: Optional[StrictStr] = None
    function_call: Optional[FunctionCall] = None


class Addon(ExtraForbidModel):
    name: Optional[StrictStr] = None
    url: Optional[StrictStr] = None


class Function(ExtraForbidModel):
    name: StrictStr
    description: StrictStr
    parameters: StrictStr


class Temperature(ConstrainedFloat):
    ge = 0
    le = 2


class TopP(ConstrainedFloat):
    ge = 0
    le = 1


class N(ConstrainedInt):
    ge = 1
    le = 128


class Stop(ConstrainedList):
    max_items: int = 4
    __args__ = tuple([StrictStr])


class Penalty(ConstrainedFloat):
    ge = -2
    le = 2


class Request(ExtraForbidModel):
    model: Optional[StrictStr] = None
    messages: List[Message]
    functions: Optional[List[Function]] = None
    function_call: Optional[
        Union[StrictStr, Mapping[StrictStr, StrictStr]]
    ] = None
    addons: Optional[List[Addon]] = None
    stream: bool = False
    temperature: Optional[Temperature] = None
    top_p: Optional[TopP] = None
    n: Optional[N] = None
    stop: Optional[Union[StrictStr, Stop]] = None
    max_tokens: Optional[PositiveInt] = None
    presence_penalty: Optional[Penalty] = None
    frequency_penalty: Optional[Penalty] = None
    logit_bias: Optional[Mapping[int, float]] = None
    user: Optional[StrictStr] = None

    api_key: StrictStr
    jwt: Optional[StrictStr] = None
    deployment_id: StrictStr
    api_version: Optional[StrictStr] = None
    headers: Mapping[StrictStr, StrictStr]

aidial_sdk/chat_completion/request.py
epam/ai-dial-sdk
import asyncio
from time import time
from typing import Any, AsyncGenerator, Callable, Coroutine, Dict, Optional
from uuid import uuid4

from fastapi import HTTPException

from aidial_sdk.chat_completion.choice import Choice
from aidial_sdk.chat_completion.chunks import (
    BaseChunk,
    EndChoiceChunk,
    EndChunk,
    UsageChunk,
    UsagePerModelChunk,
)
from aidial_sdk.chat_completion.request import Request
from aidial_sdk.exceptions import HTTPException as DialHttpException
from aidial_sdk.utils.errors import json_error, runtime_error
from aidial_sdk.utils.logging import log_error, log_exception
from aidial_sdk.utils.merge_chunks import merge
from aidial_sdk.utils.streaming import DONE_MARKER, format_chunk


class Response:
    request: Request

    _queue: asyncio.Queue
    _last_choice_index: int
    _last_usage_per_model_index: int
    _generation_started: bool
    _usage_generated: bool
    _response_id: str
    _model: Optional[str]
    _created: int

    def __init__(self, request: Request):
        self._queue = asyncio.Queue()
        self._last_choice_index = 0
        self._last_usage_per_model_index = 0
        self._generation_started = False
        self._usage_generated = False

        self.request = request
        self._response_id = str(uuid4())
        self._model = None
        self._created = int(time())

    def _add_default_fields(self, target: Dict[str, Any]) -> None:
        target["id"] = self._response_id
        if self._model:
            target["model"] = self._model
        target["created"] = self._created
        target["object"] = (
            "chat.completion.chunk"
            if self.request.stream
            else "chat.completion"
        )

    async def _generate_stream(
        self, first_chunk: BaseChunk
    ) -> AsyncGenerator[Any, None]:
        chunk = first_chunk.to_dict()

        # NOTE: add default fields only to the first chunk in a non-streaming mode and to all chunks in a streaming mode
        self._add_default_fields(chunk)

        if self.request.stream:
            formatted_chunk = format_chunk(chunk)
            yield formatted_chunk
        else:
            yield chunk

        self._queue.task_done()

        user_task_finished = False

        last_end_choice_chunk = None
        usage_chunk = {}
        while True:
            get_task = asyncio.create_task(self._queue.get())
            done = (
                await asyncio.wait(
                    [get_task, self.user_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
            )[0]
            if self.user_task in done:
                if not user_task_finished:
                    end_chunk_generated = False
                    try:
                        self.user_task.result()
                    except DialHttpException as e:
                        if self.request.stream:
                            self._queue.put_nowait(EndChunk(e))
                            end_chunk_generated = True
                        else:
                            raise HTTPException(
                                status_code=e.status_code,
                                detail=json_error(
                                    message=e.message,
                                    type=e.type,
                                    param=e.param,
                                    code=e.code,
                                ),
                            )
                    except Exception as e:
                        log_exception(e)

                        if self.request.stream:
                            self._queue.put_nowait(EndChunk(e))
                            end_chunk_generated = True
                        else:
                            raise HTTPException(
                                status_code=500,
                                detail=json_error(
                                    message="Error during processing the request",
                                    type="runtime_error",
                                ),
                            )

                    if not end_chunk_generated:
                        self._queue.put_nowait(EndChunk())
                    user_task_finished = True
            item = get_task.result() if get_task in done else await get_task

            if isinstance(item, EndChoiceChunk):
                if item.choice_index == (self.request.n or 1) - 1:
                    last_end_choice_chunk = item.to_dict()
                    self._queue.task_done()
                    continue

            if isinstance(
                item,
                (UsageChunk, UsagePerModelChunk),
            ):
                usage_chunk = merge(usage_chunk, item.to_dict())
            elif isinstance(item, BaseChunk):
                chunk = item.to_dict()

                if self.request.stream:
                    self._add_default_fields(chunk)
                    formatted_chunk = format_chunk(chunk)
                    yield formatted_chunk
                else:
                    yield chunk
            elif isinstance(item, EndChunk):
                if last_end_choice_chunk:
                    chunk = merge(last_end_choice_chunk, usage_chunk)

                    if self.request.stream:
                        self._add_default_fields(chunk)
                        formatted_chunk = format_chunk(chunk)
                        yield formatted_chunk
                    else:
                        yield chunk

                if item.exc:
                    if isinstance(item.exc, DialHttpException):
                        formatted_chunk = format_chunk(
                            json_error(
                                message=item.exc.message,
                                type=item.exc.type,
                                param=item.exc.param,
                                code=item.exc.code,
                            )
                        )
                    else:
                        formatted_chunk = format_chunk(
                            json_error(
                                message="Error during processing the request",
                                type="runtime_error",
                            )
                        )
                    yield formatted_chunk
                else:
                    if self._last_choice_index != (self.request.n or 1):
                        log_error("Not all choices were generated")

                        error = json_error(
                            message="Error during processing the request",
                            type="runtime_error",
                        )

                        if self.request.stream:
                            formatted_chunk = format_chunk(error)
                            yield formatted_chunk
                        else:
                            raise HTTPException(
                                status_code=500,
                                detail=error,
                            )

                if self.request.stream:
                    yield format_chunk(DONE_MARKER)

                self._queue.task_done()

                return

            self._queue.task_done()

    async def _generator(
        self,
        producer: Callable[[Any, Any], Coroutine[Any, Any, Any]],
        request: Request,
    ) -> BaseChunk:
        self.user_task = asyncio.create_task(producer(request, self))

        get_task = asyncio.create_task(self._queue.get())
        done = (
            await asyncio.wait(
                [get_task, self.user_task], return_when=asyncio.FIRST_COMPLETED
            )
        )[0]
        if self.user_task in done:
            try:
                self.user_task.result()
            except DialHttpException as e:
                raise HTTPException(
                    status_code=e.status_code,
                    detail=json_error(
                        message=e.message,
                        type=e.type,
                        param=e.param,
                        code=e.code,
                    ),
                )
            except Exception as e:
                log_exception(e)
                raise HTTPException(
                    status_code=500,
                    detail=json_error(
                        message="Error during processing the request",
                        type="runtime_error",
                    ),
                )

        return get_task.result() if get_task in done else await get_task

    def create_choice(self) -> Choice:
        self._generation_started = True

        if self._last_choice_index >= (self.request.n or 1):
            runtime_error("Trying to generate more chunks than requested")

        choice = Choice(self._queue, self._last_choice_index)
        self._last_choice_index += 1

        return choice

    def create_single_choice(self) -> Choice:
        if self._last_choice_index > 0:
            runtime_error("Trying to generate a single choice after choice")
        if (self.request.n or 1) > 1:
            raise DialHttpException(
                status_code=422,
                message=f"{self.request.deployment_id} deployment doesn't support n > 1",
                type="invalid_request_error",
            )

        return self.create_choice()

    def add_usage_per_model(
        self, model: str, prompt_tokens: int = 0, completion_tokens: int = 0
    ):
        self._generation_started = True

        if self._last_choice_index != (self.request.n or 1):
            runtime_error(
                'Trying to set "usage_per_model" before generating all choices',
            )

        self._queue.put_nowait(
            UsagePerModelChunk(
                self._last_usage_per_model_index,
                model,
                prompt_tokens,
                completion_tokens,
            )
        )
        self._last_usage_per_model_index += 1

    def set_usage(self, prompt_tokens: int = 0, completion_tokens: int = 0):
        self._generation_started = True

        if self._usage_generated:
            runtime_error('Trying to set "usage" twice')
        if self._last_choice_index != (self.request.n or 1):
            runtime_error(
                'Trying to set "usage" before generating all choices',
            )

        self._usage_generated = True
        self._queue.put_nowait(UsageChunk(prompt_tokens, completion_tokens))

    async def aflush(self):
        await self._queue.join()

    def set_created(self, created: int):
        if self._generation_started:
            runtime_error('Trying to set "created" after start of generation')

        self._created = created

    def set_model(self, model: str):
        if self._generation_started:
            runtime_error('Trying to set "model" after start of generation')

        self._model = model

    def set_response_id(self, response_id: str):
        if self._generation_started:
            runtime_error(
                'Trying to set "response_id" after start of generation',
            )

        self._response_id = response_id

aidial_sdk/chat_completion/response.py
epam/ai-dial-sdk
from asyncio import Queue
from types import TracebackType
from typing import Optional, Type

from aidial_sdk.chat_completion.chunks import (
    AttachmentStageChunk,
    ContentStageChunk,
    FinishStageChunk,
    NameStageChunk,
    StartStageChunk,
)
from aidial_sdk.chat_completion.enums import Status
from aidial_sdk.pydantic_v1 import ValidationError
from aidial_sdk.utils.errors import runtime_error


class Stage:
    _queue: Queue
    _choice_index: int
    _stage_index: int
    _name: Optional[str]
    _last_attachment_index: int
    _closed: bool
    _opened: bool

    def __init__(
        self,
        queue: Queue,
        choice_index: int,
        stage_index: int,
        name: Optional[str] = None,
    ):
        self._queue = 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

    def __enter__(self):
        self.open()
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        if not exc:
            self.close(Status.COMPLETED)
        else:
            self.close(Status.FAILED)

        return False

    def append_content(self, content: str):
        if not self._opened:
            runtime_error("Trying to append content to an unopened stage")
        if self._closed:
            runtime_error("Trying to append content to a closed stage")

        self._queue.put_nowait(
            ContentStageChunk(self._choice_index, self._stage_index, content)
        )

    def append_name(self, name: str):
        if not self._opened:
            runtime_error("Trying to append name to an unopened stage")
        if self._closed:
            runtime_error("Trying to append name to a closed stage")

        self._queue.put_nowait(
            NameStageChunk(self._choice_index, self._stage_index, name)
        )

    def add_attachment(
        self,
        type: Optional[str] = None,
        title: Optional[str] = None,
        data: Optional[str] = None,
        url: Optional[str] = None,
        reference_url: Optional[str] = None,
        reference_type: Optional[str] = None,
    ) -> None:
        if not self._opened:
            runtime_error("Trying to add attachment to an unopened stage")
        if self._closed:
            runtime_error("Trying to add attachment to a closed stage")

        attachment_stage_chunk = None
        try:
            attachment_stage_chunk = AttachmentStageChunk(
                choice_index=self._choice_index,
                stage_index=self._stage_index,
                attachment_index=self._last_attachment_index,
                type=type,
                title=title,
                data=data,
                url=url,
                reference_url=reference_url,
                reference_type=reference_type,
            )
        except ValidationError as e:
            runtime_error(e.errors()[0]["msg"])

        self._queue.put_nowait(attachment_stage_chunk)
        self._last_attachment_index += 1

    def open(self):
        if self._opened:
            runtime_error("The stage is already open")

        self._opened = True
        self._queue.put_nowait(
            StartStageChunk(self._choice_index, self._stage_index, self._name)
        )

    def close(self, status: Status = Status.COMPLETED):
        if not self._opened:
            runtime_error("Trying to close an unopened stage")
        if self._closed:
            runtime_error("The stage is already closed")

        self._closed = True
        self._queue.put_nowait(
            FinishStageChunk(self._choice_index, self._stage_index, status)
        )

aidial_sdk/chat_completion/stage.py
epam/ai-dial-sdk
from typing import Optional


class HTTPException(Exception):
    def __init__(
        self,
        message: str,
        status_code: int = 500,
        type: str = "runtime_error",
        param: Optional[str] = None,
        code: Optional[str] = None,
    ) -> None:
        self.message = message
        self.status_code = status_code
        self.type = type
        self.param = param
        self.code = code

    def __repr__(self):
        return "%s(message=%r, status_code=%r, type=%r, param=%r, code=%r)" % (
            self.__class__.__name__,
            self.message,
            self.status_code,
            self.type,
            self.param,
            self.code,
        )

aidial_sdk/exceptions.py
epam/ai-dial-sdk
import functools
import types
from contextvars import ContextVar
from typing import Optional

import aiohttp
import wrapt
from fastapi import FastAPI
from requests import PreparedRequest
from requests.sessions import Session
from starlette.middleware.exceptions import ExceptionMiddleware


class FastAPIMiddleware:
    def __init__(
        self,
        app: ExceptionMiddleware,
        api_key: ContextVar[Optional[str]],
        authorization: ContextVar[Optional[str]],
    ) -> None:
        self.app = app
        self.api_key = api_key
        self.authorization = authorization

    async def __call__(self, scope, receive, send) -> None:
        for header in scope["headers"]:
            if header[0] == b"api-key":
                self.api_key.set(header[1].decode("utf-8"))
            elif header[0] == b"authorization":
                self.authorization.set(header[1].decode("utf-8"))

        await self.app(scope, receive, send)


class HeaderPropagator:
    _app: FastAPI
    _dial_url: str
    _api_key: ContextVar[Optional[str]]
    _authorization: ContextVar[Optional[str]]
    _enabled: bool

    def __init__(self, app: FastAPI, dial_url: str):
        self._app = app
        self._dial_url = dial_url

        self._api_key: ContextVar[Optional[str]] = ContextVar(
            "api_key", default=None
        )
        self._authorization: ContextVar[Optional[str]] = ContextVar(
            "authorization", default=None
        )

        self._enabled = False

    def enable(self):
        if self._enabled:
            return

        self._instrument_fast_api(self._app)
        self._instrument_aiohttp()
        self._instrument_requests()
        self._enabled = True

    def _instrument_fast_api(self, app: FastAPI):
        app.add_middleware(
            FastAPIMiddleware,
            api_key=self._api_key,
            authorization=self._authorization,
        )

    def _instrument_aiohttp(self):
        def instrumented_init(wrapped, instance, args, kwargs):
            trace_config = aiohttp.TraceConfig()
            trace_config.on_request_start.append(self._on_aiohttp_request_start)

            trace_configs = list(kwargs.get("trace_configs") or [])
            trace_configs.append(trace_config)

            kwargs["trace_configs"] = trace_configs
            return wrapped(*args, **kwargs)

        wrapt.wrap_function_wrapper(
            aiohttp.ClientSession, "__init__", instrumented_init
        )

    async def _on_aiohttp_request_start(
        self,
        unused_session: aiohttp.ClientSession,
        trace_config_ctx: types.SimpleNamespace,
        params: aiohttp.TraceRequestStartParams,
    ):
        if not str(params.url).startswith(self._dial_url):
            return

        api_key_val = self._api_key.get()
        authorization_val = self._authorization.get()

        if api_key_val:
            params.headers["api-key"] = api_key_val
        if authorization_val:
            params.headers["authorization"] = authorization_val

    def _instrument_requests(self):
        wrapped_send = Session.send

        @functools.wraps(wrapped_send)
        def instrumented_send(self, request: PreparedRequest, **kwargs):
            if request.url and request.url.startswith(self._dial_url):
                api_key_val = self._dial_api_key.get()
                authorization_val = self._dial_authorization.get()

                if api_key_val:
                    request.headers["api-key"] = api_key_val
                if authorization_val:
                    request.headers["authorization"] = authorization_val

            return wrapped_send(self, request, **kwargs)

        Session._dial_url = self._dial_url  # type: ignore
        Session._dial_api_key = self._api_key  # type: ignore
        Session._dial_authorization = self._authorization  # type: ignore
        Session.send = instrumented_send

aidial_sdk/header_propagator.py
epam/ai-dial-sdk
try:
    from pydantic.v1 import *  # type: ignore
except ImportError:
    from pydantic import *  # type: ignore

aidial_sdk/pydantic_v1/__init__.py
epam/ai-dial-sdk
from typing import Optional

from aidial_sdk.exceptions import HTTPException
from aidial_sdk.utils.logging import log_error


def runtime_error(reason: str):
    log_error(reason)

    raise HTTPException(
        status_code=500,
        message="Error during processing the request",
        type="runtime_error",
    )


def json_error(
    message: Optional[str] = None,
    type: Optional[str] = None,
    param: Optional[str] = None,
    code: Optional[str] = None,
):
    return {
        "error": {
            "message": message,
            "type": type,
            "param": param,
            "code": code,
        }
    }

aidial_sdk/utils/errors.py
epam/ai-dial-sdk
from aidial_sdk.pydantic_v1 import BaseModel


class LogConfig(BaseModel):
    """Logging configuration to be set for the server"""

    version = 1
    disable_existing_loggers = False
    formatters = {
        "default": {
            "()": "uvicorn.logging.DefaultFormatter",
            "fmt": "%(levelprefix)s | %(asctime)s | %(name)s | %(process)d | %(message)s",
            "datefmt": "%Y-%m-%d %H:%M:%S",
            "use_colors": True,
        },
    }
    handlers = {
        "default": {
            "formatter": "default",
            "class": "logging.StreamHandler",
            "stream": "ext://sys.stderr",
        },
    }
    loggers = {
        "aidial_sdk": {"handlers": ["default"]},
        "uvicorn": {
            "handlers": ["default"],
            "propagate": False,
        },
    }

aidial_sdk/utils/log_config.py
epam/ai-dial-sdk
import logging
from contextvars import ContextVar

logger = logging.getLogger("aidial_sdk")
deployment_id = ContextVar("deployment_id", default=None)


def set_log_deployment(new_deployment_id):
    deployment_id.set(new_deployment_id)


def log_info(message: str, *args, **kwargs):
    logger.info(f"[{deployment_id.get()}] {message}", *args, **kwargs)


def log_debug(message: str, *args, **kwargs):
    logger.debug(f"[{deployment_id.get()}] {message}", *args, **kwargs)


def log_warning(message: str, *args, **kwargs):
    logger.warning(f"[{deployment_id.get()}] {message}", *args, **kwargs)


def log_error(message: str, *args, **kwargs):
    logger.error(f"[{deployment_id.get()}] {message}", *args, **kwargs)


def log_exception(message: Exception, *args, **kwargs):
    logger.exception(message, *args, **kwargs)

aidial_sdk/utils/logging.py
epam/ai-dial-sdk
from typing import Any, List, TypeVar, Union, cast

T = TypeVar("T")

Path = List[Union[int, str]]


LIST_OF_DICTS_ERROR_MESSAGE = (
    "Lists could be merged only if their elements are dictionaries"
)

INDEX_ERROR_MESSAGE = "A list element must have 'index' field to identify position of the element in the list"

INCONSISTENT_INDEXED_LIST_ERROR_MESSAGE = (
    "All elements of a list must be either indexed or not indexed"
)

CANNOT_MERGE_NON_INDEXED_LISTS_ERROR_MESSAGE = (
    "Cannot merge two non-indexed non-empty lists"
)

CANNOT_MERGE_NON_INDEXED_AND_INDEXED_LISTS_ERROR_MESSAGE = (
    "Cannot merge a non-indexed list with an indexed list"
)


def show_json_path(path: Path) -> str:
    ret = "$"
    for elem in path:
        if isinstance(elem, int):
            ret += f"[{elem}]"
        else:
            ret += f".{elem}"
    return ret


def merge_str(target: str, source: str, path: Path) -> str:
    return target + source


def merge_int(target: int, source: int, path: Path) -> int:
    return source


def merge_float(target: float, source: float, path: Path) -> float:
    return source


def merge_bool(target: bool, source: bool, path: Path) -> bool:
    return source


def merge_dicts(target: dict, source: dict, path: Path) -> dict:
    for key, value in source.items():
        path.append(key)
        target[key] = merge_recursive(target.get(key), value, path)
        path.pop()

    return target


def is_indexed_list(xs: list) -> bool:
    if len(xs) == 0:
        return False

    all_indexed = True
    any_indexed = False
    for elem in xs:
        if isinstance(elem, dict) and "index" in elem:
            any_indexed = True
        else:
            all_indexed = False

    if any_indexed and not all_indexed:
        raise AssertionError(INCONSISTENT_INDEXED_LIST_ERROR_MESSAGE)

    return all_indexed


def merge_indexed_lists(target: list, source: list, path: Path) -> list:
    for elem in source:
        assert isinstance(elem, dict), LIST_OF_DICTS_ERROR_MESSAGE

        index = elem.get("index")
        assert isinstance(index, int), INDEX_ERROR_MESSAGE

        path.append(index)

        if index < len(target):
            target[index] = merge_recursive(target[index], elem, path)
        else:
            target.extend([{"index": idx} for idx in range(len(target), index)])
            target.append(elem)

        path.pop()

    return target


def merge_lists(target: list, source: list, path: Path) -> list:
    is_target_indexed = is_indexed_list(target)
    is_source_indexed = is_indexed_list(source)

    if len(source) == 0:
        return target

    if len(target) == 0:
        if is_source_indexed:
            return merge_indexed_lists(target, source, path)
        else:
            return source

    if not is_target_indexed and not is_source_indexed:
        raise AssertionError(CANNOT_MERGE_NON_INDEXED_LISTS_ERROR_MESSAGE)

    assert (
        is_target_indexed and is_source_indexed
    ), CANNOT_MERGE_NON_INDEXED_AND_INDEXED_LISTS_ERROR_MESSAGE

    return merge_indexed_lists(target, source, path)


def merge_recursive(target: T, source: Any, path: Path) -> T:
    if source is None:
        return target

    if target is None:
        if isinstance(source, dict):
            target = cast(T, {})
        elif isinstance(source, list):
            target = cast(T, [])
        else:
            return source

    if isinstance(target, list) and isinstance(source, list):
        return merge_lists(target, source, path)
    elif isinstance(target, dict) and isinstance(source, dict):
        return merge_dicts(target, source, path)
    elif isinstance(target, int) and isinstance(source, int):
        return merge_int(target, source, path)
    elif isinstance(target, float) and isinstance(source, float):
        return merge_float(target, source, path)
    elif isinstance(target, bool) and isinstance(source, bool):
        return merge_bool(target, source, path)
    elif isinstance(target, str) and isinstance(source, str):
        return merge_str(target, source, path)

    raise TypeError(
        f"Cannot merge '{type(target).__name__}' with incoming '{type(source).__name__}' at path {show_json_path(path)}"
    )


def merge(*chunks: T) -> T:
    assert len(chunks) > 0, "At least one chunk must be provided"
    ret: T = chunks[0]
    for chunk in chunks[1:]:
        ret = merge_recursive(ret, chunk, path=[])
    return ret


def cleanup_indices(chunk: T) -> T:
    if isinstance(chunk, list):
        ret = []
        for elem in chunk:
            if isinstance(elem, dict) and "index" in elem:
                elem = elem.copy()
                del elem["index"]
            ret.append(cleanup_indices(elem))
        return cast(T, ret)

    if isinstance(chunk, dict):
        return cast(
            T, {key: cleanup_indices(value) for key, value in chunk.items()}
        )

    return chunk

aidial_sdk/utils/merge_chunks.py
epam/ai-dial-sdk
import json
from typing import Any, AsyncGenerator, Dict

from aidial_sdk.utils.logging import log_debug
from aidial_sdk.utils.merge_chunks import merge

DONE_MARKER = "[DONE]"


async def merge_chunks(
    chunk_stream: AsyncGenerator[Any, None]
) -> Dict[str, Any]:
    response: Dict[str, Any] = {}
    async for chunk in chunk_stream:
        response = merge(response, chunk)

    for choice in response["choices"]:
        choice["message"] = choice["delta"]
        del choice["delta"]

    return response


def format_chunk(data: Any):
    data = "data: " + (
        json.dumps(data, separators=(",", ":"))
        if isinstance(data, dict)
        else data
    )
    log_debug(data)
    return f"{data}\n\n"

aidial_sdk/utils/streaming.py
epam/ai-dial-sdk
import uvicorn

from aidial_sdk import DIALApp
from aidial_sdk.chat_completion import ChatCompletion, Request, Response


# ChatCompletion is an abstract class for applications and model adapters
class EchoApplication(ChatCompletion):
    async def chat_completion(
        self, request: Request, response: Response
    ) -> None:
        # Get last message (the newest) from the history
        last_user_message = request.messages[-1]

        # Generate response with a single choice
        with response.create_single_choice() as choice:
            # Fill the content of the response with the last user's content
            choice.append_content(last_user_message.content or "")


# DIALApp extends FastAPI to provide an user-friendly interface for routing requests to your applications
app = DIALApp()
app.add_chat_completion("echo", EchoApplication())

# Run builded app
if __name__ == "__main__":
    uvicorn.run(app, port=5000)

examples/echo.py
epam/ai-dial-sdk
import nox

nox.options.reuse_existing_virtualenvs = True

SRC = "."


def format_with_args(session: nox.Session, *args):
    session.run("autoflake", *args)
    session.run("isort", *args)
    session.run("black", *args)


@nox.session
def lint(session: nox.Session):
    """Runs linters and fixers"""
    try:
        session.run("poetry", "install", external=True)
        session.run("poetry", "check", "--lock", external=True)
        session.run("pyright", SRC)
        session.run("flake8", SRC)
        format_with_args(session, SRC, "--check")
    except Exception:
        session.error(
            "linting has failed. Run 'make format' to fix formatting and fix other errors manually"
        )


@nox.session
def format(session: nox.Session):
    """Runs linters and fixers"""
    session.run("poetry", "install", external=True)
    format_with_args(session, SRC)


@nox.session(python=["3.8", "3.9", "3.10", "3.11"])
def test(session: nox.Session) -> None:
    """Runs tests"""
    session.run("poetry", "install", external=True)
    session.run("pytest")

noxfile.py
epam/ai-dial-sdk
from aidial_sdk.chat_completion import ChatCompletion, Request, Response


class SingleChoiceApplication(ChatCompletion):
    async def chat_completion(
        self, request: Request, response: Response
    ) -> None:
        response.set_response_id("test_id")
        response.set_created(0)

        with response.create_single_choice() as choice:
            choice.append_content("Test response content")

tests/applications/single_choice_application.py
epam/ai-dial-sdk
import re
from dataclasses import dataclass
from typing import Any, List, Union

import pytest

from aidial_sdk.utils.merge_chunks import (
    CANNOT_MERGE_NON_INDEXED_AND_INDEXED_LISTS_ERROR_MESSAGE,
    CANNOT_MERGE_NON_INDEXED_LISTS_ERROR_MESSAGE,
    INCONSISTENT_INDEXED_LIST_ERROR_MESSAGE,
    cleanup_indices,
    merge,
)


@dataclass
class Test:
    __test__ = False  # Hide from pytest test discovery

    chunks: List[Any]
    expected: Union[Any, Exception]
    desc: str


test_cases: List[Test] = [
    Test(chunks=[1, 2], expected=2, desc="Merge ints"),
    Test(chunks=[1.0, 2.0], expected=2.0, desc="Merge floats"),
    Test(chunks=["foo", "bar"], expected="foobar", desc="Merge strings"),
    Test(chunks=[True, False], expected=False, desc="Merge bools"),
    Test(chunks=[{}], expected={}, desc="Merge empty dicts"),
    Test(chunks=[1, None], expected=1, desc="Merge with None right"),
    Test(chunks=[None, 1], expected=1, desc="Merge with None left"),
    Test(chunks=[None, []], expected=[], desc="Merge with None left list"),
    Test(chunks=[[], None], expected=[], desc="Merge with None right list"),
    Test(chunks=[None, {}], expected={}, desc="Merge with None left dict"),
    Test(chunks=[{}, None], expected={}, desc="Merge with None right dict"),
    Test(
        chunks=[{"a": {"b": "foo"}}, {"a": {"b": 1}}],
        expected=TypeError(
            "Cannot merge 'str' with incoming 'int' at path $.a.b"
        ),
        desc="str+int type-error",
    ),
    Test(
        chunks=[{"a": {"b": 1}}, {"a": {"b": "foo"}}],
        expected=TypeError(
            "Cannot merge 'int' with incoming 'str' at path $.a.b"
        ),
        desc="int+str type-error",
    ),
    Test(
        chunks=[{}, {"a": {"b": 1}}],
        expected={"a": {"b": 1}},
        desc="Merge to empty dict",
    ),
    Test(
        chunks=[{}, {"a": []}],
        expected={"a": []},
        desc="Merge to empty dict with empty list",
    ),
    Test(
        chunks=[{}, {"a": [1]}],
        expected={"a": [1]},
        desc="Merge to empty dict with non-empty non-indexed list",
    ),
    Test(
        chunks=[{}, {"a": [{"index": 0}, {"value": 1}]}],
        expected=AssertionError(INCONSISTENT_INDEXED_LIST_ERROR_MESSAGE),
        desc="Inconsistent list indexing",
    ),
    Test(
        chunks=[{"a": [2]}, {"a": [{"index": 0}]}],
        expected=AssertionError(
            CANNOT_MERGE_NON_INDEXED_AND_INDEXED_LISTS_ERROR_MESSAGE
        ),
        desc="Merge non-indexed and indexed lists",
    ),
    Test(
        chunks=[{"a": [{"index": 0}]}, {"a": [2]}],
        expected=AssertionError(
            CANNOT_MERGE_NON_INDEXED_AND_INDEXED_LISTS_ERROR_MESSAGE
        ),
        desc="Merge indexed and non-indexed lists",
    ),
    Test(
        chunks=[{"a": [1]}, {"a": [2]}],
        expected=AssertionError(CANNOT_MERGE_NON_INDEXED_LISTS_ERROR_MESSAGE),
        desc="Merge lists of non-dicts",
    ),
    Test(
        chunks=[{"a": [{"b": 1}]}, {"a": [{"b": 2}]}],
        expected=AssertionError(CANNOT_MERGE_NON_INDEXED_LISTS_ERROR_MESSAGE),
        desc="Merge lists of non-indexed dicts",
    ),
    Test(
        chunks=[{"a": 1, "b": 2}, {"c": 3, "d": 4}],
        expected={"a": 1, "b": 2, "c": 3, "d": 4},
        desc="Merge dicts with non-overlapping keys",
    ),
    Test(
        chunks=[{"a": 1, "b": 2}, {"c": 3, "b": 4}],
        expected={"a": 1, "b": 4, "c": 3},
        desc="Merge dicts with overlapping keys",
    ),
    Test(
        chunks=[
            {"a": [{"index": 0, "value": 1}]},
            {"a": [{"index": 0, "value": 2}]},
        ],
        expected={"a": [{"value": 2}]},
        desc="Merge lists with overlapping indices",
    ),
    Test(
        chunks=[
            {"a": [{"index": 0, "value": 0}]},
            {"a": [{"index": 1, "value": 1}]},
        ],
        expected={"a": [{"value": 0}, {"value": 1}]},
        desc="Merge lists with non-overlapping indices",
    ),
    Test(
        chunks=[
            {"a": []},
            {"a": [{"index": 1, "value": 1}]},
            {"a": [{"index": 0, "value": 0}]},
        ],
        expected={"a": [{"value": 0}, {"value": 1}]},
        desc="Merge lists out-of-order",
    ),
    Test(
        chunks=[
            {},
            {"a": [{"index": 5, "value": 5}]},
            {"a": [{"index": 4, "value": 4}]},
            {"a": [{"index": 2, "value": 2}]},
            {"a": [{"index": 1, "value": 1}]},
        ],
        expected={
            "a": [
                {},
                {"value": 1},
                {"value": 2},
                {},
                {"value": 4},
                {"value": 5},
            ]
        },
        desc="Merge lists out-of-order (no starting point)",
    ),
    Test(
        chunks=[
            {"a": [{"index": 0, "value": 0}]},
            {"a": [{"index": 2, "value": 2}]},
        ],
        expected={"a": [{"value": 0}, {}, {"value": 2}]},
        desc="Merge lists with a forward gap",
    ),
    Test(
        chunks=[{"a": "Hello "}, {"a": "world!"}],
        expected={"a": "Hello world!"},
        desc="Merge nested strings",
    ),
    Test(
        chunks=[
            {"usage": {"prompt_tokens": 1}},
            {"usage": {"prompt_tokens": 2}},
        ],
        expected={"usage": {"prompt_tokens": 2}},
        desc="Merge top-level usage",
    ),
    Test(
        chunks=[
            {"a": {"usage": {"prompt_tokens": 1}}},
            {"a": {"usage": {"prompt_tokens": 2}}},
        ],
        expected={"a": {"usage": {"prompt_tokens": 2}}},
        desc="Merge nested usage",
    ),
]


@pytest.mark.parametrize("test", test_cases, ids=lambda t: t.desc)
def test_merge_chunks(test: Test):
    if isinstance(test.expected, Exception):
        with pytest.raises(
            type(test.expected), match=re.escape(str(test.expected))
        ):
            cleanup_indices(merge(*test.chunks))
    else:
        assert cleanup_indices(merge(*test.chunks)) == test.expected

tests/test_merge_chunks.py
epam/ai-dial-sdk
import json

from starlette.testclient import TestClient

from aidial_sdk import DIALApp
from tests.applications.single_choice_application import SingleChoiceApplication


def test_single_choice():
    dial_app = DIALApp()
    dial_app.add_chat_completion("test_app", SingleChoiceApplication())

    test_app = TestClient(dial_app)

    response = test_app.post(
        "/openai/deployments/test_app/chat/completions",
        json={
            "messages": [{"role": "user", "content": "Test content"}],
            "stream": False,
        },
        headers={"Api-Key": "TEST_API_KEY"},
    )

    assert response.status_code == 200 and response.json() == {
        "choices": [
            {
                "index": 0,
                "finish_reason": "stop",
                "message": {
                    "role": "assistant",
                    "content": "Test response content",
                },
            }
        ],
        "usage": None,
        "id": "test_id",
        "created": 0,
        "object": "chat.completion",
    }


def test_single_choice_streaming():
    dial_app = DIALApp()
    dial_app.add_chat_completion("test_app", SingleChoiceApplication())

    test_app = TestClient(dial_app)

    response = test_app.post(
        "/openai/deployments/test_app/chat/completions",
        json={
            "messages": [{"role": "user", "content": "Test content"}],
            "stream": True,
        },
        headers={"Api-Key": "TEST_API_KEY"},
    )

    for index, value in enumerate(response.iter_lines()):
        if index % 2:
            assert value == ""
            continue

        assert value.startswith("data: ")
        data = value[6:]

        if index == 0:
            assert json.loads(data) == {
                "choices": [
                    {
                        "index": 0,
                        "finish_reason": None,
                        "delta": {"role": "assistant"},
                    }
                ],
                "usage": None,
                "id": "test_id",
                "created": 0,
                "object": "chat.completion.chunk",
            }
        elif index == 2:
            assert json.loads(data) == {
                "choices": [
                    {
                        "index": 0,
                        "finish_reason": None,
                        "delta": {"content": "Test response content"},
                    }
                ],
                "usage": None,
                "id": "test_id",
                "created": 0,
                "object": "chat.completion.chunk",
            }
        elif index == 4:
            assert json.loads(data) == {
                "choices": [{"index": 0, "finish_reason": "stop", "delta": {}}],
                "usage": None,
                "id": "test_id",
                "created": 0,
                "object": "chat.completion.chunk",
            }
        elif index == 6:
            assert data == "[DONE]"

tests/test_single_choice.py
epam/ai-dial-sdk