aidial-sdk 0.48.0.dev2__tar.gz → 0.48.0.dev4__tar.gz

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/PKG-INFO

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

aidial_sdk-0.48.0.dev4/aidial_sdk/_errors.py

@@ -0,0 +1,81 @@
+from collections.abc import Awaitable, Callable
+
+from fastapi import HTTPException, Request
+from fastapi.responses import JSONResponse, Response
+from fastapi.routing import APIRoute
+
+from aidial_sdk._pydantic import ValidationError
+from aidial_sdk.exceptions import (
+    ExceptionType,
+    InternalServerError,
+    InvalidRequestError,
+)
+from aidial_sdk.exceptions import HTTPException as DIALException
+from aidial_sdk.utils.logging import logger as log
+
+INTERNAL_ERROR_MESSAGE = "Internal server error"
+
+Handler = Callable[[Request], Awaitable[Response]]
+
+
+def _get_exception_type(request: Request) -> ExceptionType:
+    route = request.scope.get("route")
+    if isinstance(route, APIRoute) and route.path.startswith("/anthropic/"):
+        return "anthropic"
+    return "openai"
+
+
+def invalid_request_error(e: ValidationError) -> InvalidRequestError:
+    error = e.errors()[0]
+    path = ".".join(map(str, error["loc"]))
+    return InvalidRequestError(
+        f"Your request contained invalid structure on path {path}. {error['msg']}"
+    )
+
+
+def to_dial_exception(e: Exception) -> DIALException:
+    match e:
+        case DIALException():
+            return e
+        case _:
+            log.exception("Unexpected error in the request handler")
+            return InternalServerError(INTERNAL_ERROR_MESSAGE)
+
+
+def _error_response(
+    request: Request, e: Exception, api: ExceptionType | None
+) -> Response:
+    if isinstance(e, HTTPException):
+        return JSONResponse(
+            status_code=e.status_code,
+            content=e.detail,
+            headers=e.headers,
+        )
+
+    api = api or _get_exception_type(request)
+    return to_dial_exception(e).to_fastapi_response(api)
+
+
+def error_handler(api: ExceptionType | None) -> Callable[[Handler], Handler]:
+    def _decorator(handler: Handler) -> Handler:
+        async def _handler(request: Request) -> Response:
+            try:
+                return await handler(request)
+            except Exception as e:
+                return _error_response(request, e, api)
+
+        return _handler
+
+    return _decorator
+
+
+def dial_exception_handler(request: Request, exc: Exception) -> Response:
+    return _error_response(request, exc, None)
+
+
+def unhandled_exception_handler(request: Request, exc: Exception) -> Response:
+    # ServerErrorMiddleware re-raises the exception for the server to log it
+    error = InternalServerError(
+        f"Unhandled SDK exception: {INTERNAL_ERROR_MESSAGE}"
+    )
+    return _error_response(request, error, None)

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/aidial_sdk/application.py

@@ -10,9 +10,11 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse
 from starlette.types import Receive, Scope, Send
 
 from aidial_sdk._errors import (
+    Handler,
     dial_exception_handler,
-    fastapi_exception_handler,
-    pydantic_validation_exception_handler,
+    error_handler,
+    invalid_request_error,
+    unhandled_exception_handler,
 )
 from aidial_sdk._pydantic import ValidationError
 from aidial_sdk._pydantic._compat import BaseModel
@@ -41,11 +43,11 @@ from aidial_sdk.utils._interpolation import (
     interpolate_deployment_id,
     match_deployment_id,
 )
+from aidial_sdk.utils._logging_middleware import LoggingMiddleware
 from aidial_sdk.utils._reflection import get_method_implementation
 from aidial_sdk.utils.env import env_float, env_var_list
 from aidial_sdk.utils.log_config import configure_sdk_logger
 from aidial_sdk.utils.logging import (
-    log_debug,
     reset_log_context,
     set_log_deployment,
 )
@@ -60,8 +62,6 @@ configure_sdk_logger()
 
 RequestType = TypeVar("RequestType", bound=FromRequestMixin)
 
-Handler = Callable[[Request], Coroutine[Any, Any, Response]]
-
 
 class PathFilter(logging.Filter):
     path: str
@@ -145,14 +145,12 @@ class DIALApp(FastAPI):
 
         self.add_middleware(DisconnectMiddleware)
         self.add_middleware(NonForwardableHeadersMiddleware)
+        self.add_middleware(LoggingMiddleware)
 
-        self.add_exception_handler(
-            ValidationError, pydantic_validation_exception_handler
-        )
-
-        self.add_exception_handler(HTTPException, fastapi_exception_handler)
-
+        # The routes DIALApp adds have error_handler, these are for custom ones
+        self.add_exception_handler(HTTPException, dial_exception_handler)
         self.add_exception_handler(DIALException, dial_exception_handler)
+        self.add_exception_handler(Exception, unhandled_exception_handler)
 
     def configure_telemetry(self, config: TelemetryConfig | None):
         if config is None or config.is_noop():
@@ -184,7 +182,7 @@ class DIALApp(FastAPI):
     ) -> "DIALApp":
         self.add_api_route(
             f"/openai/deployments/{deployment_name}/embeddings",
-            self._embeddings(deployment_name, impl),
+            error_handler("openai")(self._embeddings(deployment_name, impl)),
             methods=["POST"],
         )
 
@@ -193,6 +191,7 @@ class DIALApp(FastAPI):
             "/openai/v1/embeddings",
             self._embeddings(None, impl),
             methods=["POST"],
+            handler_decorator=error_handler("openai"),
         )
 
         return self
@@ -209,7 +208,7 @@ class DIALApp(FastAPI):
         ):
             self.add_api_route(
                 f"/openai/deployments/{deployment_name}/{endpoint}",
-                handler,
+                error_handler("openai")(handler),
                 methods=methods,
             )
 
@@ -221,6 +220,7 @@ class DIALApp(FastAPI):
                 f"/openai/v1/{endpoint}",
                 handler,
                 methods=methods,
+                handler_decorator=error_handler("openai"),
             )
 
         return self
@@ -281,12 +281,13 @@ class DIALApp(FastAPI):
         handler: Handler,
         *,
         methods: list[str],
+        handler_decorator: Callable[[Handler], Handler],
     ) -> None:
         if endpoint not in self._generic_handlers:
             self._generic_handlers[endpoint] = {}
             self.add_api_route(
                 endpoint,
-                self._generic_dispatcher(endpoint),
+                handler_decorator(self._generic_dispatcher(endpoint)),
                 methods=methods,
             )
         self._generic_handlers[endpoint][deployment_name] = handler
@@ -337,7 +338,6 @@ class DIALApp(FastAPI):
             request = await self._parse_request(
                 request_type, original_request, deployment_id
             )
-            log_debug(f"request[{endpoint}]: {request}")
 
             response = await endpoint_impl(request)
 
@@ -350,8 +350,6 @@ class DIALApp(FastAPI):
                     f"Unexpected response type from {endpoint}: {type(response)}"
                 )
 
-            log_debug(f"response[{endpoint}]: {response_json}")
-
             return JSONResponse(content=response_json)
 
         return _handler
@@ -379,11 +377,14 @@ class DIALApp(FastAPI):
             )
             set_log_deployment(deployment_id)
 
-        ret = await request.from_request(
-            original_request, deployment_id, self._dial_url
-        )
-        if not self._allow_extra_request_fields:
-            model_validate_extra_fields(ret)
+        try:
+            ret = await request.from_request(
+                original_request, deployment_id, self._dial_url
+            )
+            if not self._allow_extra_request_fields:
+                model_validate_extra_fields(ret)
+        except ValidationError as e:
+            raise invalid_request_error(e) from e
         return ret
 
     def _chat_completion(
@@ -392,7 +393,7 @@ class DIALApp(FastAPI):
         impl: ChatCompletion,
         *,
         heartbeat_interval: float | None,
-    ):
+    ) -> Handler:
         heartbeat_interval = heartbeat_interval or env_float(
             "DIAL_SDK_SSE_HEARTBEAT_INTERVAL"
         )
@@ -403,7 +404,6 @@ class DIALApp(FastAPI):
             )
 
             response = ChatCompletionResponse(request)
-
             stream = response._generate_stream(impl.chat_completion)
 
             if request.stream:
@@ -411,7 +411,6 @@ class DIALApp(FastAPI):
                     stream = add_heartbeat(
                         stream,
                         heartbeat_interval=heartbeat_interval,
-                        heartbeat_callback=lambda: log_debug("heartbeat"),
                         heartbeat_object=": heartbeat\n\n",
                     )
 
@@ -421,7 +420,6 @@ class DIALApp(FastAPI):
                 )
             else:
                 response_json = await to_block_response(stream)
-                log_debug(f"response: {response_json}")
                 resp = JSONResponse(content=response_json)
 
             for key, value in response.headers:

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/aidial_sdk/chat_completion/response.py

@@ -70,7 +70,9 @@ class Response:
                 if self.request.stream
                 else "chat.completion"
             ),
-            model=get_client_deployment_id(request.headers),
+            model=get_client_deployment_id(
+                request.headers, request.deployment_id
+            ),
         )
 
         self._headers = []

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/aidial_sdk/deployment/from_request_mixin.py

@@ -44,14 +44,24 @@ def get_upstream_deployment_id(
     return (
         headers.get(DIAL_OVERRIDE_NAME)
         or deployment_id
-        or get_client_deployment_id(headers)
+        or _get_dial_deployment_id(headers)
     )
 
 
-def get_client_deployment_id(headers: Mapping[str, str]) -> str:
+def get_client_deployment_id(
+    headers: Mapping[str, str], deployment_id: str | None
+) -> str:
     """The client-facing deployment id that a DIAL application is expected
     to return as a model parameter in the response."""
 
+    # NOTE: Once https://github.com/epam/ai-dial-core/issues/2243 is fixed,
+    # this could be replaced with just _get_dial_deployment_id.
+    return headers.get(DIAL_DEPLOYMENT_ID) or get_upstream_deployment_id(
+        headers, deployment_id
+    )
+
+
+def _get_dial_deployment_id(headers: Mapping[str, str]) -> str:
     if (deployment_id := headers.get(DIAL_DEPLOYMENT_ID)) is None:
         raise InternalServerError(
             f"The request headers are missing {DIAL_DEPLOYMENT_ID} header."

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/aidial_sdk/exceptions.py

@@ -1,12 +1,38 @@
 import functools
 import warnings
+from enum import Enum
 from http import HTTPStatus
+from typing import Literal, assert_never
 
 from fastapi import HTTPException as FastAPIException
 from fastapi.responses import JSONResponse
 
 from aidial_sdk.utils.json import remove_nones
 
+ExceptionType = Literal["openai", "anthropic"]
+
+
+class _AnthropicErrorType(Enum):
+    INVALID_REQUEST = "invalid_request_error", 400, 422
+    AUTHENTICATION = "authentication_error", 401
+    PERMISSION = "permission_error", 403
+    NOT_FOUND = "not_found_error", 404
+    REQUEST_TOO_LARGE = "request_too_large", 413
+    RATE_LIMIT = "rate_limit_error", 429
+    API = "api_error", 500, 502
+    OVERLOADED = "overloaded_error", 529, 503
+
+    def __init__(self, code: str, *status_codes: int) -> None:
+        self.code = code
+        self.status_codes = status_codes
+
+    @classmethod
+    def from_status_code(cls, status_code: int) -> "_AnthropicErrorType":
+        for error_type in cls:
+            if status_code in error_type.status_codes:
+                return error_type
+        return cls.API
+
 
 class HTTPException(Exception):
     def __init__(
@@ -47,31 +73,47 @@ class HTTPException(Exception):
         args = ", ".join([f"{k}={v!r}" for k, v in error.items()])
         return f"{self.__class__.__name__}({args})"
 
-    def json_error(self) -> dict:
-        return {
-            "error": remove_nones(
-                {
-                    "message": self.message,
-                    "type": self.type,
-                    "param": self.param,
-                    "code": self.code,
-                    "display_message": self.display_message,
-                    **self.extra_fields,
+    def json_error(self, api: ExceptionType = "openai") -> dict:
+        match api:
+            case "openai":
+                return {
+                    "error": remove_nones(
+                        {
+                            "message": self.message,
+                            "type": self.type,
+                            "param": self.param,
+                            "code": self.code,
+                            "display_message": self.display_message,
+                            **self.extra_fields,
+                        }
+                    )
                 }
-            )
-        }
+            case "anthropic":
+                error_type = _AnthropicErrorType.from_status_code(
+                    self.status_code
+                )
+                return {
+                    "type": "error",
+                    "error": {"type": error_type.code, "message": self.message},
+                }
+            case _:
+                assert_never(api)
 
-    def to_fastapi_response(self) -> JSONResponse:
+    def to_fastapi_response(
+        self, api: ExceptionType = "openai"
+    ) -> JSONResponse:
         return JSONResponse(
             status_code=self.status_code,
-            content=self.json_error(),
+            content=self.json_error(api),
             headers=self.headers,
         )
 
-    def to_fastapi_exception(self) -> FastAPIException:
+    def to_fastapi_exception(
+        self, api: ExceptionType = "openai"
+    ) -> FastAPIException:
         return FastAPIException(
             status_code=self.status_code,
-            detail=self.json_error(),
+            detail=self.json_error(api),
             headers=self.headers,
         )
 

aidial_sdk-0.48.0.dev4/aidial_sdk/translators/messages_to_chat_completions/backend/http.py

@@ -0,0 +1,161 @@
+from collections.abc import AsyncGenerator, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from urllib.parse import quote
+
+import httpx
+import openai
+from fastapi import Request
+from openai import AsyncOpenAI, AsyncStream, Omit, omit
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+
+from aidial_sdk.exceptions import HTTPException as DIALException
+from aidial_sdk.translators.messages_to_chat_completions.settings import (
+    get_dial_api_version,
+)
+from aidial_sdk.utils.logging import logger as log
+
+# The headers of the inbound request that DIAL Core reads when the translator
+# calls it back. Any other header is not forwarded.
+_FORWARDED_HEADERS = {
+    # Core's per-request key
+    "api-key",
+    # Prompt caching policy of the request
+    "x-dial-cache-policy",
+    # Pins the request to the same upstream
+    "x-upstream-id",
+}
+
+
+def _client(
+    base_url: str,
+    http_client: httpx.AsyncClient,
+    default_query: dict[str, str] | None = None,
+) -> AsyncOpenAI:
+    return AsyncOpenAI(
+        # Required by the SDK; Core authenticates the forwarded `api-key` header.
+        api_key="-",
+        base_url=base_url,
+        default_query=default_query,
+        max_retries=0,
+        http_client=http_client,
+    )
+
+
+def core_v1_client(
+    dial_url: str, http_client: httpx.AsyncClient
+) -> AsyncOpenAI:
+    return _client(f"{dial_url}/openai/v1", http_client)
+
+
+def core_legacy_client(
+    dial_url: str, deployment: str, http_client: httpx.AsyncClient
+) -> AsyncOpenAI:
+    return _client(
+        f"{dial_url}/openai/deployments/{quote(deployment, safe='')}",
+        http_client,
+        {"api-version": get_dial_api_version()},
+    )
+
+
+def core_headers(request: Request) -> dict[str, str | Omit]:
+    # Without the override the SDK sends `Authorization: Bearer -`
+    headers: dict[str, str | Omit] = {"Authorization": omit}
+    headers.update(
+        {
+            name: value
+            for name, value in request.headers.items()
+            if name.lower() in _FORWARDED_HEADERS
+        }
+    )
+    return headers
+
+
+# An error event in a stream has no status code, only a type or a code
+_STREAM_ERROR_STATUS_CODES: dict[str | None, int] = {
+    "invalid_request_error": 400,
+    "rate_limit_exceeded": 429,
+}
+
+
+def _extract_error_message(e: openai.APIStatusError) -> str:
+    body = e.body
+    if isinstance(body, dict):
+        error = body.get("error")
+        if isinstance(error, dict) and isinstance(error.get("message"), str):
+            return error["message"]
+        if isinstance(error, str) and error:
+            return error
+        if isinstance(body.get("message"), str):
+            return body["message"]
+
+    if isinstance(body, str) and body.strip():
+        return body.strip()
+
+    return e.message
+
+
+def _from_core_error(e: openai.APIError) -> DIALException:
+    match e:
+        case openai.APIStatusError():
+            return DIALException(
+                _extract_error_message(e), status_code=e.status_code
+            )
+        case openai.APIConnectionError():
+            log.error(f"Failed to reach DIAL Core: {e}")
+            return DIALException("Failed to reach DIAL Core", status_code=502)
+        case _:
+            log.error(f"DIAL Core reported an error in the stream: {e}")
+            status_code = (
+                _STREAM_ERROR_STATUS_CODES.get(e.type)
+                or _STREAM_ERROR_STATUS_CODES.get(e.code)
+                or 500
+            )
+            return DIALException(e.message, status_code=status_code)
+
+
+@contextmanager
+def _core_errors() -> Iterator[None]:
+    try:
+        yield
+    except openai.APIError as e:
+        raise _from_core_error(e) from e
+
+
+async def _chunks(
+    stream: AsyncStream[ChatCompletionChunk],
+) -> AsyncGenerator[ChatCompletionChunk, None]:
+    with _core_errors():
+        async with stream:
+            async for item in stream:
+                yield item
+
+
+@dataclass
+class ChatClientHttp:
+    client: AsyncOpenAI
+    core_headers: dict[str, str | Omit]
+
+    async def create_stream(
+        self, body: dict
+    ) -> AsyncGenerator[ChatCompletionChunk, None]:
+        with _core_errors():
+            stream = await self.client.post(
+                "/chat/completions",
+                cast_to=ChatCompletion,
+                body=body,
+                options={"headers": self.core_headers},
+                stream=True,
+                stream_cls=AsyncStream[ChatCompletionChunk],
+            )
+        return _chunks(stream)
+
+    async def create_block(self, body: dict) -> ChatCompletion:
+        with _core_errors():
+            response = await self.client.post(
+                "/chat/completions",
+                cast_to=ChatCompletion,
+                body=body,
+                options={"headers": self.core_headers},
+            )
+        return response

{aidial_sdk-0.48.0.dev2 → aidial_sdk-0.48.0.dev4}/aidial_sdk/translators/messages_to_chat_completions/chat_completions/app.py

@@ -8,7 +8,9 @@ from anthropic.types.beta.message_create_params import MessageCreateParams
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
+from aidial_sdk._errors import error_handler
 from aidial_sdk.deployment.from_request_mixin import get_client_deployment_id
+from aidial_sdk.exceptions import InvalidRequestError, ResourceNotFoundError
 from aidial_sdk.translators.messages_to_chat_completions.backend.http import (
     ChatClientHttp,
     core_headers,
@@ -24,39 +26,33 @@ from aidial_sdk.translators.messages_to_chat_completions.chat_completions.stream
 from aidial_sdk.translators.messages_to_chat_completions.chat_completions.to_chat_completions import (
     to_chat_completions_request,
 )
-from aidial_sdk.translators.messages_to_chat_completions.common import (
-    not_found,
-    with_debug_logging,
-)
-from aidial_sdk.translators.messages_to_chat_completions.errors import (
-    AnthropicErrorType,
-    AnthropicHTTPError,
-    translator_error_handler,
-)
 from aidial_sdk.translators.messages_to_chat_completions.settings import (
     get_dial_url,
 )
 from aidial_sdk.translators.messages_to_chat_completions.translation_state import (
     TranslationState,
 )
+from aidial_sdk.utils._logging_middleware import LoggingMiddleware
 from aidial_sdk.utils.logging import logger as log
 
 ClientT = TypeVar("ClientT")
 ClientFactory = Callable[[Request], Awaitable[ClientT]] | ClientT
 
 
+async def _not_found(request: Request) -> Response:
+    error = ResourceNotFoundError(f"Unknown endpoint: {request.url.path}")
+    return error.to_fastapi_response("anthropic")
+
+
 async def handle_anthropic_messages(
     request: Request, client: ChatClientHttp
 ) -> Response:
-    deployment = get_client_deployment_id(request.headers)
+    deployment = get_client_deployment_id(request.headers, None)
 
     try:
         req: MessageCreateParams = await request.json()
     except ValueError as e:
-        raise AnthropicHTTPError(
-            AnthropicErrorType.INVALID_REQUEST,
-            f"Request body is not valid JSON: {e}",
-        ) from e
+        raise InvalidRequestError(f"Request body is not valid JSON: {e}") from e
 
     translation_state = TranslationState()
     chat_request = to_chat_completions_request(
@@ -65,7 +61,7 @@ async def handle_anthropic_messages(
 
     body = chat_request.model_dump(mode="json", exclude_none=True)
     if log.isEnabledFor(logging.DEBUG):
-        log.debug(f"Translated request to Core: {json.dumps(body)}")
+        log.debug(f"chat request: {json.dumps(body)}")
 
     if chat_request.stream:
         events = await client.create_stream(body)
@@ -85,6 +81,7 @@ def create_anthropic_api_app(
     get_client: ClientFactory[httpx.AsyncClient],
 ) -> FastAPI:
     app = FastAPI()
+    app.add_middleware(LoggingMiddleware)
 
     async def _chat_client(request: Request, *, legacy: bool) -> ChatClientHttp:
         http_client = (
@@ -92,7 +89,7 @@ def create_anthropic_api_app(
         )
         dial_url = get_dial_url()
         if legacy:
-            deployment_id = get_client_deployment_id(request.headers)
+            deployment_id = get_client_deployment_id(request.headers, None)
             client = core_legacy_client(dial_url, deployment_id, http_client)
         else:
             client = core_v1_client(dial_url, http_client)
@@ -100,16 +97,14 @@ def create_anthropic_api_app(
         return ChatClientHttp(client, core_headers(request))
 
     @app.post("/anthropic/v1/messages")
-    @with_debug_logging
-    @translator_error_handler
+    @error_handler("anthropic")
     async def _handle_v1_messages(request: Request) -> Response:
         return await handle_anthropic_messages(
             request, await _chat_client(request, legacy=False)
         )
 
     @app.post("/legacy/anthropic/v1/messages")
-    @with_debug_logging
-    @translator_error_handler
+    @error_handler("anthropic")
     async def _handle_legacy_messages(request: Request) -> Response:
         return await handle_anthropic_messages(
             request, await _chat_client(request, legacy=True)
@@ -117,7 +112,7 @@ def create_anthropic_api_app(
 
     app.add_api_route(
         "/{full_path:path}",
-        not_found,
+        _not_found,
         methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
     )