aidial-sdk 0.48.0.dev3__tar.gz → 0.48.0.dev5__tar.gz

{aidial_sdk-0.48.0.dev3 → aidial_sdk-0.48.0.dev5}/PKG-INFO

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

aidial_sdk-0.48.0.dev5/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.dev3 → aidial_sdk-0.48.0.dev5}/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.dev3 → aidial_sdk-0.48.0.dev5}/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.dev5/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.dev5/aidial_sdk/translators/messages_to_chat_completions/chat_completions/_reasoning_content.py

@@ -0,0 +1,95 @@
+import re
+from dataclasses import dataclass
+from typing import Literal
+
+from openai.types.chat import ChatCompletionMessage
+from openai.types.chat.chat_completion_chunk import ChoiceDelta
+
+from aidial_sdk.translators.messages_to_chat_completions.chat_completions.dial_extensions import (
+    CustomContent,
+    parse_extras,
+)
+
+_Source = Literal["azure", "claude", "dial"]
+
+_REASONING_STAGE = re.compile(r"think|thought|reason", re.IGNORECASE)
+
+
+@dataclass
+class ReasoningContent:
+    content: str | None
+    signature: str | None
+
+
+class ReasoningContentProvider:
+    """
+    Reads the reasoning of a message, or of a stream of message chunks, from
+    `reasoning_content`, the Claude state and the DIAL reasoning stages. They
+    carry the same text, chunked differently: only the text none of them has
+    delivered before is returned.
+    """
+
+    def __init__(self) -> None:
+        # Every source is a prefix of the same text, so offsets suffice.
+        self._returned = 0
+        self._received: dict[_Source, int] = {
+            "azure": 0,
+            "claude": 0,
+            "dial": 0,
+        }
+        self._reasoning_stages: set[int | None] = set()
+
+    def get_reasoning_content(
+        self, message: ChatCompletionMessage | ChoiceDelta
+    ) -> ReasoningContent:
+        custom_content = parse_extras(message.model_extra).custom_content
+        thinking, signature = _get_claude_reasoning_content(custom_content)
+        content = (
+            self._unseen("azure", _get_azure_reasoning_content(message))
+            + self._unseen("claude", thinking)
+            + self._unseen("dial", self._dial_reasoning_content(custom_content))
+        )
+        return ReasoningContent(
+            content=content or None,
+            signature=signature or None,
+        )
+
+    def _unseen(self, source: _Source, text: str | None) -> str:
+        text = text or ""
+        start = self._received[source]
+        self._received[source] += len(text)
+        unseen = text[self._returned - start :]
+        self._returned = max(self._returned, self._received[source])
+        return unseen
+
+    def _dial_reasoning_content(
+        self, custom_content: CustomContent | None
+    ) -> str:
+        text = ""
+        for stage in (custom_content.stages if custom_content else None) or []:
+            # In a stream, only the first chunk of a stage carries its name.
+            if _is_reasoning_stage(stage.name):
+                self._reasoning_stages.add(stage.index)
+            if stage.index in self._reasoning_stages:
+                text += stage.content or ""
+        return text
+
+
+def _get_azure_reasoning_content(
+    message: ChatCompletionMessage | ChoiceDelta,
+) -> str | None:
+    return (message.model_extra or {}).get("reasoning_content")
+
+
+def _get_claude_reasoning_content(
+    custom_content: CustomContent | None,
+) -> tuple[str | None, str | None]:
+    state = custom_content.state if custom_content else None
+    for block in (state.claude_message_content if state else None) or []:
+        if block.type == "thinking":
+            return block.thinking, block.signature
+    return None, None
+
+
+def _is_reasoning_stage(name: str | None) -> bool:
+    return name is not None and bool(_REASONING_STAGE.search(name))