aidial-sdk 0.44.0.dev4__tar.gz → 0.45.0__tar.gz

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/PKG-INFO

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

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/__init__.py

@@ -1,5 +1,9 @@
 from aidial_sdk.application import DIALApp
 from aidial_sdk.exceptions import HTTPException
+from aidial_sdk.response_headers import (
+    NonForwardableHeadersMiddleware,
+    strip_stale_content_headers,
+)
 from aidial_sdk.utils.log_config import LogConfig, configure_root_logger
 from aidial_sdk.utils.logging import logger
 
@@ -7,6 +11,8 @@ __all__ = [
     "DIALApp",
     "HTTPException",
     "LogConfig",
+    "NonForwardableHeadersMiddleware",
     "configure_root_logger",
     "logger",
+    "strip_stale_content_headers",
 ]

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/application.py

@@ -34,6 +34,7 @@ from aidial_sdk.embeddings.request import Request as EmbeddingsRequest
 from aidial_sdk.exceptions import DeploymentNotFoundError
 from aidial_sdk.exceptions import HTTPException as DIALException
 from aidial_sdk.header_propagator import HeaderPropagator
+from aidial_sdk.response_headers import NonForwardableHeadersMiddleware
 from aidial_sdk.telemetry.types import TelemetryConfig
 from aidial_sdk.utils._disconnect_middleware import DisconnectMiddleware
 from aidial_sdk.utils._reflection import get_method_implementation
@@ -136,6 +137,7 @@ class DIALApp(FastAPI):
             logging.getLogger("uvicorn.access").addFilter(PathFilter(path))
 
         self.add_middleware(DisconnectMiddleware)
+        self.add_middleware(NonForwardableHeadersMiddleware)
 
         self.add_exception_handler(
             ValidationError, pydantic_validation_exception_handler

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/chat_completion/choice.py

@@ -195,12 +195,25 @@ class Choice(ChoiceBase):
         self.send_chunk(FormSchemaChunk(self._index, form_schema))
 
     def create_stage(self, name: str | None = None) -> Stage:
+        return self._create_stage(name, parent=None)
+
+    def _create_stage(
+        self,
+        name: str | None = None,
+        *,
+        parent: Stage | None = None,
+    ) -> Stage:
         if not self._opened:
             raise runtime_error("Trying to create stage to an unopened choice")
         if self._closed:
             raise runtime_error("Trying to create stage to a closed choice")
 
-        stage = Stage(self._queue, self._index, self._last_stage_index, name)
+        stage = Stage(
+            self,
+            self._last_stage_index,
+            name,
+            parent,
+        )
         self._last_stage_index += 1
 
         return stage

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/chat_completion/chunks.py

@@ -211,29 +211,34 @@ class StartStageChunk(BaseChunk):
     choice_index: int
     stage_index: int
     name: str | None
+    parent_stage_index: int | None
 
-    def __init__(self, choice_index: int, stage_index: int, name: str | None):
+    def __init__(
+        self,
+        choice_index: int,
+        stage_index: int,
+        name: str | None,
+        parent_stage_index: int | None = None,
+    ):
         self.choice_index = choice_index
         self.stage_index = stage_index
         self.name = name
+        self.parent_stage_index = parent_stage_index
 
     def to_dict(self):
+        stage: dict[str, Any] = {
+            "index": self.stage_index,
+            "name": self.name,
+            "status": None,
+        }
+        if self.parent_stage_index is not None:
+            stage["parent_stage_index"] = self.parent_stage_index
         return {
             "choices": [
                 {
                     "index": self.choice_index,
                     "finish_reason": None,
-                    "delta": {
-                        "custom_content": {
-                            "stages": [
-                                {
-                                    "index": self.stage_index,
-                                    "name": self.name,
-                                    "status": None,
-                                }
-                            ]
-                        }
-                    },
+                    "delta": {"custom_content": {"stages": [stage]}},
                 }
             ],
             "usage": None,

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/chat_completion/request.py

@@ -46,6 +46,7 @@ class Stage(ExtraAllowModel, IgnoreIndex):
     status: Status
     content: StrictStr | None = None
     attachments: list[Attachment] | None = None
+    parent_stage_index: StrictInt | None = None
 
 
 class CustomContent(ExtraAllowModel):

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/chat_completion/stage.py

@@ -1,5 +1,5 @@
 from types import TracebackType
-from typing import overload
+from typing import TYPE_CHECKING, overload
 
 from aidial_sdk._pydantic import ValidationError
 from aidial_sdk.chat_completion._types import ChunkQueue
@@ -15,31 +15,52 @@ from aidial_sdk.chat_completion.request import Attachment
 from aidial_sdk.utils._attachment import create_attachment
 from aidial_sdk.utils._content_stream import ContentStream
 from aidial_sdk.utils.errors import runtime_error
+from aidial_sdk.utils.logging import log_warning
+
+if TYPE_CHECKING:
+    from aidial_sdk.chat_completion.choice import Choice
 
 
 class Stage:
+    _choice: "Choice"
     _queue: ChunkQueue
     _choice_index: int
     _stage_index: int
     _name: str | None
+    _parent: "Stage | None"
+    _children: list["Stage"]
     _last_attachment_index: int
     _closed: bool
     _opened: bool
 
     def __init__(
         self,
-        queue: ChunkQueue,
-        choice_index: int,
+        choice: "Choice",
         stage_index: int,
         name: str | None = None,
+        parent: "Stage | None" = None,
     ):
-        self._queue = queue
-        self._choice_index = choice_index
+        if parent is not None and parent._choice is not choice:
+            raise runtime_error(
+                "Trying to create a stage whose parent stage belongs to another choice"
+            )
+
+        self._choice = choice
+        self._queue = choice._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
+        self._parent = parent
+        self._children = []
+
+        if parent is not None:
+            parent._children.append(self)
+
+    def create_stage(self, name: str | None = None) -> "Stage":
+        return self._choice._create_stage(name, parent=self)
 
     def __enter__(self):
         self.open()
@@ -119,9 +140,25 @@ class Stage:
         if self._opened:
             raise runtime_error("The stage is already open")
 
+        parent = self._parent
+        if parent is not None:
+            if not parent._opened:
+                raise runtime_error(
+                    "Trying to open a stage whose parent stage is not open"
+                )
+            if parent._closed:
+                raise runtime_error(
+                    "Trying to open a stage whose parent stage is already closed"
+                )
+
         self._opened = True
         self._queue.put_nowait(
-            StartStageChunk(self._choice_index, self._stage_index, self._name)
+            StartStageChunk(
+                self._choice_index,
+                self._stage_index,
+                self._name,
+                parent._stage_index if parent is not None else None,
+            )
         )
 
     def close(self, status: Status = Status.COMPLETED):
@@ -130,6 +167,17 @@ class Stage:
         if self._closed:
             raise runtime_error("The stage is already closed")
 
+        open_children = [
+            child._stage_index
+            for child in self._children
+            if child._opened and not child._closed
+        ]
+        if open_children:
+            log_warning(
+                f"Closing the stage {self._stage_index} which still has "
+                f"open child stages: {open_children}"
+            )
+
         self._closed = True
         self._queue.put_nowait(
             FinishStageChunk(self._choice_index, self._stage_index, status)

aidial_sdk-0.45.0/aidial_sdk/response_headers.py

@@ -0,0 +1,95 @@
+"""Response headers a DIAL component must not relay downstream.
+
+A component sitting in front of an upstream API - a model adapter proxying the
+vendor's response, an interceptor rebuilding it - receives the upstream's
+response headers and tends to pass them on. Some of them describe the
+**upstream hop**, not the response the component itself emits:
+
+* hop-by-hop headers (RFC 9110 7.6.1). ``transfer-encoding`` is the harmful
+  one: the ASGI server computes its own ``content-length`` for the body the
+  component emits, and a message carrying both is rejected by a strict HTTP
+  parser (RFC 9112 6.2) - DIAL Core being one. A relayed ``connection: close``
+  tears down the *downstream* connection once the response is over.
+* ``server`` and ``date``, which the ASGI server supplies itself. Both are
+  singleton fields (RFC 9110 5.6.6), so the upstream's copy duplicates them.
+
+Which headers these are follows from the HTTP spec alone, never from what the
+component does, so ``NonForwardableHeadersMiddleware`` strips them from every
+response the app emits and takes no configuration:
+
+* an app built on ``DIALApp`` has it installed already - nothing to do;
+* an app built on plain ``FastAPI``/``Starlette`` installs it once::
+
+      app = FastAPI()
+      app.add_middleware(NonForwardableHeadersMiddleware)
+
+  It covers the whole app, including sub-apps mounted into it, so a passthrough
+  proxy mounted as one needs no strip of its own.
+
+``content-length`` and ``content-encoding`` are a different matter: they are
+accurate for the upstream body and go stale only once a component reads and
+rebuilds it. The middleware cannot tell a stale one from the correct one the
+ASGI server has just computed, so the rebuilding code drops them itself, on the
+upstream headers it is about to reuse::
+
+      body = await upstream_response.aread()
+      strip_stale_content_headers(upstream_response.headers)
+      return Response(content=body, headers=upstream_response.headers)
+"""
+
+from collections.abc import MutableMapping
+
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+_HOP_BY_HOP_HEADERS = frozenset(
+    {
+        b"connection",
+        b"keep-alive",
+        b"proxy-authenticate",
+        b"proxy-authorization",
+        b"te",
+        b"trailer",
+        b"transfer-encoding",
+        b"upgrade",
+    }
+)
+
+_SERVER_MANAGED_HEADERS = frozenset({b"server", b"date"})
+
+_NON_FORWARDABLE_HEADERS = _HOP_BY_HOP_HEADERS | _SERVER_MANAGED_HEADERS
+
+_STALE_CONTENT_HEADERS = ("content-length", "content-encoding")
+
+
+class NonForwardableHeadersMiddleware:
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(
+        self, scope: Scope, receive: Receive, send: Send
+    ) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        async def send_wrapper(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                message["headers"] = [
+                    (key, value)
+                    for key, value in message["headers"]
+                    if key.lower() not in _NON_FORWARDABLE_HEADERS
+                ]
+            await send(message)
+
+        await self.app(scope, receive, send_wrapper)
+
+
+def strip_stale_content_headers(headers: MutableMapping[str, str]) -> None:
+    """Drop from `headers` the content headers describing the upstream body.
+
+    Call it on the upstream response headers once the body they describe has
+    been read and rebuilt, so that the framing and the content coding of the
+    new body are recomputed rather than inherited.
+    """
+    for header in _STALE_CONTENT_HEADERS:
+        headers.pop(header, None)

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/telemetry/init.py

@@ -37,6 +37,8 @@ from prometheus_client import start_http_server
 from aidial_sdk.telemetry.types import TelemetryConfig
 from aidial_sdk.utils._logging import remove_stream_handlers
 
+_log = logging.getLogger(__name__)
+
 
 def init_telemetry(app: FastAPI | None, config: TelemetryConfig):
     resource = Resource.create(
@@ -142,4 +144,20 @@ def init_telemetry(app: FastAPI | None, config: TelemetryConfig):
 
     if app and (config.tracing is not None or config.metrics is not None):
         # FastAPI instrumentor reports both metrics and traces
+        excl = config.tracing.excluded_asgi_spans if config.tracing else []
+        _instrument_fastapi_app(app, excl)
+
+
+def _instrument_fastapi_app(app: FastAPI, exclude_spans: list[str]) -> None:
+    if not exclude_spans:
+        FastAPIInstrumentor.instrument_app(app)
+        return
+
+    try:
+        FastAPIInstrumentor.instrument_app(app, exclude_spans=exclude_spans)  # type: ignore
+    except TypeError:
+        _log.warning(
+            "OTEL_PYTHON_FASTAPI_EXCLUDE_SPANS is ignored: it requires "
+            "opentelemetry-instrumentation-fastapi>=0.48b0."
+        )
         FastAPIInstrumentor.instrument_app(app)

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/aidial_sdk/telemetry/types.py

@@ -16,6 +16,9 @@ OTEL_EXPORTER_PROMETHEUS_PORT = int(
 OTEL_PYTHON_LOG_CORRELATION = (
     os.getenv("OTEL_PYTHON_LOG_CORRELATION", "false").lower() == "true"
 )
+OTEL_PYTHON_FASTAPI_EXCLUDE_SPANS = env_var_list(
+    "OTEL_PYTHON_FASTAPI_EXCLUDE_SPANS"
+)
 
 
 class LogsConfig(BaseModel):
@@ -30,6 +33,7 @@ class LogsConfig(BaseModel):
 class TracingConfig(BaseModel):
     otlp_export: bool = "otlp" in OTEL_TRACES_EXPORTER
     logging: bool = OTEL_PYTHON_LOG_CORRELATION
+    excluded_asgi_spans: list[str] = OTEL_PYTHON_FASTAPI_EXCLUDE_SPANS
 
 
 class MetricsConfig(BaseModel):

{aidial_sdk-0.44.0.dev4 → aidial_sdk-0.45.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-sdk"
-version = "0.44.0.dev4"
+version = "0.45.0"
 description = "Framework to create applications and model adapters for AI DIAL"
 authors = [{ name = "EPAM RAIL", email = "SpecialEPM-DIALDevTeam@epam.com" }]
 license = "Apache-2.0"