The project implements AI DIAL API for language models from Azure OpenAI
import json
import logging.config
import os

import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from openai import ChatCompletion, Embedding, error

from aidial_adapter_openai.openai_override import OpenAIException
from aidial_adapter_openai.utils.exceptions import HTTPException
from aidial_adapter_openai.utils.log_config import LogConfig
from aidial_adapter_openai.utils.parsers import (
    ApiType,
    parse_body,
    parse_upstream,
)
from aidial_adapter_openai.utils.streaming import generate_stream

logging.config.dictConfig(LogConfig().dict())
app = FastAPI()
model_aliases = json.loads(os.getenv("MODEL_ALIASES", "{}"))
azure_api_version = os.getenv("AZURE_API_VERSION", "2023-03-15-preview")


async def handle_exceptions(call):
    try:
        return await call
    except OpenAIException as e:
        return Response(status_code=e.code, headers=e.headers, content=e.body)
    except error.Timeout:
        raise HTTPException("Request timed out", 504, "timeout")
    except error.APIConnectionError:
        raise HTTPException(
            "Error communicating with OpenAI", 502, "connection"
        )


@app.post("/openai/deployments/{deployment_id}/chat/completions")
async def chat_completion(deployment_id: str, request: Request):
    data = await parse_body(request)

    is_stream = data.get("stream", False)
    dial_api_key = request.headers["X-UPSTREAM-KEY"]

    api_base, upstream_deployment = parse_upstream(
        request.headers["X-UPSTREAM-ENDPOINT"], ApiType.CHAT_COMPLETION
    )

    response = await handle_exceptions(
        ChatCompletion().acreate(
            engine=upstream_deployment,
            api_key=dial_api_key,
            api_base=api_base,
            api_type="azure",
            api_version=azure_api_version,
            request_timeout=(10, 600),  # connect timeout and total timeout
            **data
        )
    )

    if is_stream:
        return StreamingResponse(
            generate_stream(
                data["messages"],
                response,
                model_aliases.get(deployment_id, deployment_id),
                deployment_id,
            ),
            media_type="text/event-stream",
        )
    else:
        return response


@app.post("/openai/deployments/{deployment_id}/embeddings")
async def embedding(deployment_id: str, request: Request):
    data = await parse_body(request)

    dial_api_key = request.headers["X-UPSTREAM-KEY"]
    api_base, upstream_deployment = parse_upstream(
        request.headers["X-UPSTREAM-ENDPOINT"], ApiType.EMBEDDING
    )

    return await handle_exceptions(
        Embedding().acreate(
            deployment_id=upstream_deployment,
            api_key=dial_api_key,
            api_base=api_base,
            api_type="azure",
            api_version=azure_api_version,
            request_timeout=(10, 600),  # connect timeout and total timeout
            **data
        )
    )


@app.exception_handler(HTTPException)
def exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "error": {
                "message": exc.message,
                "type": exc.type,
                "param": exc.param,
                "code": exc.code,
            }
        },
    )


@app.get("/health")
def health():
    return {"status": "ok"}


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

aidial_adapter_openai/app.py
epam/ai-dial-adapter-openai
"""
OpenAI SDK translates various HTTP errors received from OpenAI API
into Python exceptions: error.RateLimitError, error.InvalidRequestError,
error.AuthenticationError etc.

We want to retranslate the original HTTP errors to the user.
So the standard error handlers in the openai.api_requestor.APIRequestor class
are rewritten to wrap the original HTTP errors into OpenAIException and raise it.
"""
import json
from json import JSONDecodeError

import wrapt
from openai.api_requestor import APIRequestor
from openai.openai_response import OpenAIResponse


class OpenAIException(Exception):
    def __init__(self, body, code, resp, headers):
        self.body = body
        self.code = code
        self.resp = resp
        self.headers = headers

        super().__init__(resp)


# Overridden to proxy original errors
def handle_error_response_wrapper(wrapped, self, args, kwargs):
    raise OpenAIException(*args)


# Overridden to proxy original errors
def interpret_response_line_wrapper(wrapped, self: APIRequestor, args, kwargs):
    rbody, rcode, rheaders = args
    stream = kwargs.get("stream", False)

    # HTTP 204 response code does not have any content in the body.
    if rcode == 204:
        return OpenAIResponse(None, rheaders)

    if rcode == 503:
        raise self.handle_error_response(  # overrided
            rbody, rcode, None, rheaders, stream_error=False
        )
    try:
        if "text/plain" in rheaders.get("Content-Type", ""):
            data = rbody
        else:
            data = json.loads(rbody)
    except (JSONDecodeError, UnicodeDecodeError):
        raise self.handle_error_response(  # overrided
            rbody,
            rcode,
            None,
            rheaders,
            stream_error=False,
        )
    resp = OpenAIResponse(data, rheaders)
    # In the future, we might add a "status" parameter to errors
    # to better handle the "error while streaming" case.
    stream_error = stream and "error" in resp.data
    if stream_error or not 200 <= rcode < 300:
        raise self.handle_error_response(
            rbody, rcode, resp.data, rheaders, stream_error=stream_error
        )
    return resp


wrapt.wrap_function_wrapper(
    APIRequestor, "handle_error_response", handle_error_response_wrapper
)

wrapt.wrap_function_wrapper(
    APIRequestor, "_interpret_response_line", interpret_response_line_wrapper
)

aidial_adapter_openai/openai_override.py
epam/ai-dial-adapter-openai
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_adapter_openai/utils/exceptions.py
epam/ai-dial-adapter-openai
import logging
import os

from pydantic import BaseModel

LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")


class LogConfig(BaseModel):
    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_adapter_openai": {"handlers": ["default"], "level": LOG_LEVEL},
        "uvicorn": {
            "handlers": ["default"],
            "level": LOG_LEVEL,
            "propagate": False,
        },
    }


logger = logging.getLogger("aidial_adapter_openai")

aidial_adapter_openai/utils/log_config.py
epam/ai-dial-adapter-openai
import re
from enum import Enum
from json import JSONDecodeError
from typing import Any, Mapping

from fastapi import Request

from aidial_adapter_openai.utils.exceptions import HTTPException


class ApiType(Enum):
    CHAT_COMPLETION = r"(.+?)/openai/deployments/(.+?)/chat/completions"
    EMBEDDING = r"(.+?)/openai/deployments/(.+?)/embeddings"

    def pattern(self):
        return self.value


def parse_upstream(
    api_endpoint: str,
    api_type: ApiType,
) -> tuple[str, str]:
    match = re.search(api_type.pattern(), api_endpoint)

    if not match:
        raise HTTPException(
            "Invalid upstream endpoint format", 400, "invalid_request_error"
        )

    return match[1], match[2]


async def parse_body(
    request: Request,
) -> Mapping[str, Any]:
    try:
        data = await request.json()
    except JSONDecodeError as e:
        raise HTTPException(
            "Your request contained invalid JSON: " + str(e),
            400,
            "invalid_request_error",
        )

    if type(data) != dict:
        raise HTTPException(
            str(data) + " is not of type 'object'", 400, "invalid_request_error"
        )

    return data

aidial_adapter_openai/utils/parsers.py
epam/ai-dial-adapter-openai
import json
from time import time
from typing import Any, Mapping
from uuid import uuid4

import tiktoken

from aidial_adapter_openai.openai_override import OpenAIException
from aidial_adapter_openai.utils.log_config import logger
from aidial_adapter_openai.utils.tokens import calculate_prompt_tokens


def chunk_format(data: str | Mapping[str, Any]):
    if type(data) == str:
        return "data: " + data.strip() + "\n\n"
    else:
        return "data: " + json.dumps(data, separators=(",", ":")) + "\n\n"


END_CHUNK = chunk_format("[DONE]")


async def generate_stream(
    messages: list[Any], response, model: str, deployment: str
):
    encoding = tiktoken.encoding_for_model(model)

    prompt_tokens = calculate_prompt_tokens(messages, model, encoding)

    last_chunk = None
    stream_finished = False
    try:
        total_content = ""
        async for chunk in response:
            chunk_dict = chunk.to_dict_recursive()

            if chunk_dict["choices"][0]["finish_reason"] is not None:
                stream_finished = True
                completion_tokens = len(encoding.encode(total_content))
                chunk_dict["usage"] = {
                    "completion_tokens": completion_tokens,
                    "prompt_tokens": prompt_tokens,
                    "total_tokens": prompt_tokens + completion_tokens,
                }
            else:
                total_content += chunk_dict["choices"][0]["delta"].get(
                    "content", ""
                )

            last_chunk = chunk_dict
            yield chunk_format(chunk_dict)
    except OpenAIException as e:
        yield chunk_format(e.body)
        yield END_CHUNK

        return

    if not stream_finished:
        completion_tokens = len(encoding.encode(total_content))

        if last_chunk is not None:
            logger.warning("Don't received chunk with the finish reason")

            last_chunk["usage"] = {
                "completion_tokens": completion_tokens,
                "prompt_tokens": prompt_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            }
            last_chunk["choices"][0]["delta"]["content"] = ""
            last_chunk["choices"][0]["delta"]["finish_reason"] = "length"

            yield chunk_format(last_chunk)
        else:
            logger.warning("Received 0 chunks")

            yield chunk_format(
                {
                    "id": "chatcmpl-" + str(uuid4()),
                    "object": "chat.completion.chunk",
                    "created": str(int(time())),
                    "model": deployment,
                    "choices": [
                        {"index": 0, "finish_reason": "length", "delta": {}}
                    ],
                    "usage": {
                        "completion_tokens": 0,
                        "prompt_tokens": prompt_tokens,
                        "total_tokens": prompt_tokens,
                    },
                }
            )

    yield END_CHUNK

aidial_adapter_openai/utils/streaming.py
epam/ai-dial-adapter-openai
"""
Implemented based on the official recipe: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken
"""
from typing import Any, List

from tiktoken import Encoding


def calculate_prompt_tokens(
    messages: List[Any], model: str, encoding: Encoding
):
    prompt_tokens = 3

    if model == "gpt-3.5-turbo-0301":
        tokens_per_message = 4
        tokens_per_name = -1
    else:
        tokens_per_message = 3
        tokens_per_name = 1

    for message in messages:
        prompt_tokens += tokens_per_message

        for key, value in message.items():
            prompt_tokens += len(encoding.encode(value))
            if key == "name":
                prompt_tokens += tokens_per_name

    return prompt_tokens

aidial_adapter_openai/utils/tokens.py
epam/ai-dial-adapter-openai
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
def test(session: nox.Session):
    """Runs unit tests"""
    session.run("poetry", "install", external=True)
    session.run("pytest", "tests/")

noxfile.py
epam/ai-dial-adapter-openai
import glob
import os
import shutil
from typing import List


def remove_dir(directory_path: str):
    if os.path.exists(directory_path):
        shutil.rmtree(directory_path)
        print("Removed: " + directory_path)


def remove_recursively(pattern: str, exclude_dirs: List[str] = []):
    files = glob.glob(f"**/{pattern}", recursive=True)

    def is_excluded(file: str) -> bool:
        return any(dir in file for dir in exclude_dirs)

    for file in files:
        if not is_excluded(file):
            remove_dir(file)


def main():
    remove_dir(".nox")
    remove_dir("dist")
    remove_recursively("__pycache__", exclude_dirs=[".venv"])
    remove_recursively(".pytest_cache", exclude_dirs=[".venv"])


if __name__ == "__main__":
    main()

scripts/clean.py
epam/ai-dial-adapter-openai
import json

import pytest
from aioresponses import aioresponses
from httpx import AsyncClient

from aidial_adapter_openai.app import app


@pytest.mark.asyncio
async def test_error_during_streaming(aioresponses: aioresponses):
    aioresponses.post(
        "http://localhost:5001/openai/deployments/gpt-4/chat/completions?api-version=2023-03-15-preview",
        status=200,
        body="data: "
        + json.dumps(
            {
                "id": "chatcmpl-test",
                "object": "chat.completion.chunk",
                "created": 1695940483,
                "model": "gpt-4",
                "choices": [
                    {
                        "index": 0,
                        "finish_reason": "stop",
                        "message": {
                            "role": "assistant",
                        },
                    }
                ],
                "usage": None,
            }
        )
        + "\n\n"
        + "data: "
        + json.dumps(
            {
                "error": {
                    "message": "Error test",
                    "type": "runtime_error",
                    "param": None,
                    "code": None,
                }
            }
        )
        + "\n\n"
        + "data: [DONE]\n\n",
        content_type="text/event-stream",
    )
    test_app = AsyncClient(app=app, base_url="http://test.com")

    response = await test_app.post(
        "/openai/deployments/gpt-4/chat/completions",
        json={
            "messages": [{"role": "user", "content": "Test content"}],
            "stream": True,
        },
        headers={
            "X-UPSTREAM-KEY": "TEST_API_KEY",
            "X-UPSTREAM-ENDPOINT": "http://localhost:5001/openai/deployments/gpt-4/chat/completions",
        },
    )

    assert response.status_code == 200

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

        if index == 0:
            assert (
                line
                == 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1695940483,"model":"gpt-4","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant"}}],"usage":{"completion_tokens":0,"prompt_tokens":9,"total_tokens":9}}'
            )
        elif index == 2:
            assert (
                line
                == 'data: {"error": {"message": "Error test", "type": "runtime_error", "param": null, "code": null}}'
            )
        elif index == 4:
            assert line == "data: [DONE]"
        else:
            assert False


@pytest.mark.asyncio
async def test_incorrect_upstream_url(aioresponses: aioresponses):
    aioresponses.post(
        "http://localhost:5001/openai/deployments/gpt-4/chat/completions?api-version=2023-03-15-preview",
        status=200,
        body={},
    )
    test_app = AsyncClient(app=app, base_url="http://test.com")

    response = await test_app.post(
        "/openai/deployments/gpt-4/chat/completions",
        json={"messages": [{"role": "user", "content": "Test content"}]},
        headers={
            "X-UPSTREAM-KEY": "TEST_API_KEY",
            "X-UPSTREAM-ENDPOINT": "http://localhost:5001",  # upstream endpoint should contain the full path
        },
    )

    assert response.status_code == 400
    assert response.json() == {
        "error": {
            "message": "Invalid upstream endpoint format",
            "type": "invalid_request_error",
            "param": None,
            "code": None,
        }
    }


@pytest.mark.asyncio
async def test_asdasd(aioresponses: aioresponses):
    aioresponses.post(
        "http://localhost:5001/openai/deployments/gpt-4/chat/completions?api-version=2023-03-15-preview",
        status=400,
        body="Incorrect format",
    )
    test_app = AsyncClient(app=app, base_url="http://test.com")

    response = await test_app.post(
        "/openai/deployments/gpt-4/chat/completions",
        json={"messages": [{"role": "user", "content": "Test content"}]},
        headers={
            "X-UPSTREAM-KEY": "TEST_API_KEY",
            "X-UPSTREAM-ENDPOINT": "http://localhost:5001/openai/deployments/gpt-4/chat/completions",
        },
    )

    assert response.status_code == 400

    assert response.content == b"Incorrect format"

tests/test_errors.py
epam/ai-dial-adapter-openai
import json

import pytest
from aioresponses import aioresponses
from httpx import AsyncClient

from aidial_adapter_openai.app import app


@pytest.mark.asyncio
async def test_streaming(aioresponses: aioresponses):
    aioresponses.post(
        "http://localhost:5001/openai/deployments/gpt-4/chat/completions?api-version=2023-03-15-preview",
        status=200,
        body="data: "
        + json.dumps(
            {
                "id": "chatcmpl-test",
                "object": "chat.completion.chunk",
                "created": 1695940483,
                "model": "gpt-4",
                "choices": [
                    {
                        "index": 0,
                        "finish_reason": None,
                        "delta": {
                            "role": "assistant",
                        },
                    }
                ],
                "usage": None,
            }
        )
        + "\n\n"
        + "data: "
        + json.dumps(
            {
                "id": "chatcmpl-test",
                "object": "chat.completion.chunk",
                "created": 1695940483,
                "model": "gpt-4",
                "choices": [
                    {
                        "index": 0,
                        "finish_reason": None,
                        "delta": {
                            "content": "Test content",
                        },
                    }
                ],
                "usage": None,
            }
        )
        + "\n\n"
        + "data: "
        + json.dumps(
            {
                "id": "chatcmpl-test",
                "object": "chat.completion.chunk",
                "created": 1696245654,
                "model": "gpt-4",
                "choices": [{"index": 0, "finish_reason": "stop", "delta": {}}],
                "usage": None,
            }
        )
        + "\n\n"
        + "data: [DONE]\n\n",
        content_type="text/event-stream",
    )
    test_app = AsyncClient(app=app, base_url="http://test.com")

    response = await test_app.post(
        "/openai/deployments/gpt-4/chat/completions?api-version=2023-06-15",
        json={
            "messages": [{"role": "user", "content": "Test content"}],
            "stream": True,
        },
        headers={
            "X-UPSTREAM-KEY": "TEST_API_KEY",
            "X-UPSTREAM-ENDPOINT": "http://localhost:5001/openai/deployments/gpt-4/chat/completions",
        },
    )

    assert response.status_code == 200

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

        if index == 0:
            assert (
                line
                == 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1695940483,"model":"gpt-4","choices":[{"index":0,"finish_reason":null,"delta":{"role":"assistant"}}],"usage":null}'
            )
        elif index == 2:
            assert (
                line
                == 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1695940483,"model":"gpt-4","choices":[{"index":0,"finish_reason":null,"delta":{"content":"Test content"}}],"usage":null}'
            )
        elif index == 4:
            assert (
                line
                == 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1696245654,"model":"gpt-4","choices":[{"index":0,"finish_reason":"stop","delta":{}}],"usage":{"completion_tokens":2,"prompt_tokens":9,"total_tokens":11}}'
            )
        elif index == 6:
            assert line == "data: [DONE]"
        else:
            assert False

tests/test_streaming.py
epam/ai-dial-adapter-openai