aidial-sdk 0.42.0.dev1__tar.gz → 0.43.0.dev1__tar.gz

{aidial_sdk-0.42.0.dev1 → aidial_sdk-0.43.0.dev1}/PKG-INFO

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

{aidial_sdk-0.42.0.dev1 → aidial_sdk-0.43.0.dev1}/aidial_sdk/application.py

@@ -2,10 +2,12 @@ import logging
 import re
 import warnings
 from collections.abc import Callable, Coroutine, Iterator
+from contextlib import AbstractContextManager, nullcontext
 from typing import Any, Literal, TypeVar
 
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.types import Receive, Scope, Send
 
 from aidial_sdk._errors import (
     dial_exception_handler,
@@ -37,7 +39,11 @@ from aidial_sdk.utils._disconnect_middleware import DisconnectMiddleware
 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, set_log_deployment
+from aidial_sdk.utils.logging import (
+    log_debug,
+    reset_log_context,
+    set_log_deployment,
+)
 from aidial_sdk.utils.pydantic import model_validate_extra_fields
 from aidial_sdk.utils.streaming import (
     add_heartbeat,
@@ -77,6 +83,7 @@ class DIALApp(FastAPI):
     _allow_extra_request_fields: bool
     _dial_url: str | None
     _v1_handlers: dict[str, dict[str, Handler]]
+    _reset_otel_context: Callable[[], AbstractContextManager[None]]
 
     def __init__(
         self,
@@ -103,6 +110,7 @@ class DIALApp(FastAPI):
         self._allow_extra_request_fields = allow_extra_request_fields
         self._dial_url = dial_url
         self._v1_handlers = {}
+        self._reset_otel_context = nullcontext
 
         self.configure_telemetry(telemetry_config)
 
@@ -142,6 +150,7 @@ class DIALApp(FastAPI):
             return
 
         try:
+            from aidial_sdk.telemetry._context import reset_otel_context
             from aidial_sdk.telemetry.init import init_telemetry
 
             init_telemetry(app=self, config=config)
@@ -151,6 +160,16 @@ class DIALApp(FastAPI):
                 "Install the package with the extras: aidial-sdk[telemetry]"
             )
 
+        self._reset_otel_context = reset_otel_context
+
+    async def __call__(
+        self, scope: Scope, receive: Receive, send: Send
+    ) -> None:
+        # Workaround for the bug in asyncio: https://github.com/python/cpython/pull/141158
+        reset_log_context()
+        with self._reset_otel_context():
+            await super().__call__(scope, receive, send)
+
     def add_embeddings(
         self, deployment_name: str, impl: Embeddings
     ) -> "DIALApp":

aidial_sdk-0.43.0.dev1/aidial_sdk/telemetry/_context.py

@@ -0,0 +1,13 @@
+import contextlib
+from collections.abc import Iterator
+
+from opentelemetry.context import Context, attach, detach
+
+
+@contextlib.contextmanager
+def reset_otel_context() -> Iterator[None]:
+    token = attach(Context())
+    try:
+        yield
+    finally:
+        detach(token)

{aidial_sdk-0.42.0.dev1 → aidial_sdk-0.43.0.dev1}/aidial_sdk/utils/logging.py

@@ -12,6 +12,10 @@ def set_log_deployment(new_deployment_id: str):
     deployment_id.set(new_deployment_id)
 
 
+def reset_log_context() -> None:
+    deployment_id.set(None)
+
+
 def log_info(message: str, *args, **kwargs):
     logger.info(f"[{deployment_id.get()}] {message}", *args, **kwargs)
 

{aidial_sdk-0.42.0.dev1 → aidial_sdk-0.43.0.dev1}/pyproject.toml

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