aidial-sdk 0.37.0.dev7__tar.gz → 0.38.0.dev1__tar.gz

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/PKG-INFO

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

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/aidial_sdk/application.py

@@ -283,7 +283,7 @@ class DIALApp(FastAPI):
             "DIAL_SDK_SSE_HEARTBEAT_INTERVAL"
         )
 
-        async def _handler(original_request: Request):
+        async def _handler(original_request: Request) -> Response:
             request = await self._parse_request(
                 ChatCompletionRequest, original_request, deployment_id
             )

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/aidial_sdk/chat_completion/__init__.py

@@ -1,3 +1,4 @@
+from aidial_sdk.chat_completion._cache import CacheBreakpointPath
 from aidial_sdk.chat_completion.base import ChatCompletion
 from aidial_sdk.chat_completion.choice import Choice
 from aidial_sdk.chat_completion.enums import FinishReason, Status
@@ -55,6 +56,7 @@ __all__ = [
     "Attachment",
     "Button",
     "CacheBreakpoint",
+    "CacheBreakpointPath",
     "ChatCompletion",
     "Choice",
     "ConfigurationRequest",

aidial_sdk-0.38.0.dev1/aidial_sdk/chat_completion/_cache.py

@@ -0,0 +1,13 @@
+class CacheBreakpointPath:
+    path: str
+
+    def __init__(self, path: str) -> None:
+        self.path = path
+
+    @classmethod
+    def messages(cls, idx: int):
+        return cls(f"prefix.body.messages[{idx}]")
+
+    @classmethod
+    def tools(cls, idx: int):
+        return cls(f"prefix.body.tools[{idx}]")

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/aidial_sdk/chat_completion/response.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 
 from typing_extensions import assert_never
 
+from aidial_sdk.chat_completion._cache import CacheBreakpointPath
 from aidial_sdk.chat_completion._types import ChunkQueue
 from aidial_sdk.chat_completion.choice import Choice
 from aidial_sdk.chat_completion.chunks import (
@@ -22,6 +23,11 @@ from aidial_sdk.chat_completion.chunks import (
     UsagePerModelChunk,
 )
 from aidial_sdk.chat_completion.request import Request
+from aidial_sdk.deployment._headers import (
+    DIAL_CACHE_BREAKPOINT_PATH,
+    DIAL_CACHE_EXPIRE_AT,
+    DIAL_CACHE_EXTRA_METADATA,
+)
 from aidial_sdk.exceptions import HTTPException as DIALException
 from aidial_sdk.exceptions import RequestValidationError, RuntimeServerError
 from aidial_sdk.utils._cancel_scope import CancelScope
@@ -259,3 +265,18 @@ class Response:
                 "Trying to set a header after start of generation",
             )
         self._headers.append((key, value))
+
+    def set_cache_breakpoint(
+        self,
+        *,
+        cache_breakpoint_path: CacheBreakpointPath,
+        cache_expire_at: str | None = None,
+        cache_metadata: str | None = None,
+    ):
+        self.append_header(
+            DIAL_CACHE_BREAKPOINT_PATH, cache_breakpoint_path.path
+        )
+        if cache_expire_at is not None:
+            self.append_header(DIAL_CACHE_EXPIRE_AT, cache_expire_at)
+        if cache_metadata is not None:
+            self.append_header(DIAL_CACHE_EXTRA_METADATA, cache_metadata)

aidial_sdk-0.38.0.dev1/aidial_sdk/deployment/_headers.py

@@ -0,0 +1,13 @@
+DIAL_CONVERSATION_ID = "X-CONVERSATION-ID"
+DIAL_JOB_TITLE = "X-JOB-TITLE"
+
+DIAL_UPSTREAM_ENDPOINT = "X-UPSTREAM-ENDPOINT"
+DIAL_UPSTREAM_KEY = "X-UPSTREAM-KEY"
+DIAL_UPSTREAM_EXTRA_DATA = "X-UPSTREAM-EXTRA-DATA"
+
+DIAL_APPLICATION_PROPERTIES = "X-DIAL-APPLICATION-PROPERTIES"
+DIAL_APPLICATION_ID = "X-DIAL-APPLICATION-ID"
+
+DIAL_CACHE_BREAKPOINT_PATH = "X-DIAL-CACHE-BREAKPOINT-PATH"
+DIAL_CACHE_EXPIRE_AT = "X-DIAL-CACHE-EXPIRE-AT"
+DIAL_CACHE_EXTRA_METADATA = "X-DIAL-CACHE-EXTRA-METADATA"

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/aidial_sdk/deployment/from_request_mixin.py

@@ -15,6 +15,17 @@ from aidial_sdk._pydantic import (
     StrictStr,
 )
 from aidial_sdk._pydantic._compat import model_validator
+from aidial_sdk.deployment._headers import (
+    DIAL_APPLICATION_ID,
+    DIAL_APPLICATION_PROPERTIES,
+    DIAL_CACHE_BREAKPOINT_PATH,
+    DIAL_CACHE_EXTRA_METADATA,
+    DIAL_CONVERSATION_ID,
+    DIAL_JOB_TITLE,
+    DIAL_UPSTREAM_ENDPOINT,
+    DIAL_UPSTREAM_EXTRA_DATA,
+    DIAL_UPSTREAM_KEY,
+)
 from aidial_sdk.exceptions import InternalServerError, InvalidRequestError
 from aidial_sdk.utils.logging import log_debug
 from aidial_sdk.utils.pydantic import ExtraAllowModel
@@ -39,10 +50,6 @@ class FromRequestMixin(ABC, ExtraAllowModel):
         pass
 
 
-_DIAL_APPLICATION_PROPERTIES_HEADER = "X-DIAL-APPLICATION-PROPERTIES"
-_DIAL_APPLICATION_ID_HEADER = "X-DIAL-APPLICATION-ID"
-
-
 class FromRequestDeploymentMixin(FromRequestMixin):
     headers: HeadersType
     base_url: str | None = None
@@ -53,6 +60,17 @@ class FromRequestDeploymentMixin(FromRequestMixin):
     api_version: StrictStr | None = None
     unreliable_dial_application_properties: dict[str, Any] | None = None
     dial_application_id: str | None = None
+
+    conversation_id: str | None = None
+    job_title: str | None = None
+
+    upstream_endpoint: str | None = None
+    upstream_key: str | None = None
+    upstream_extra_data: str | None = None
+
+    cache_breakpoint_path: str | None = None
+    cache_extra_metadata: str | None = None
+
     original_request: fastapi.Request = Field(..., exclude=True)
 
     if PYDANTIC_V2:
@@ -70,7 +88,7 @@ class FromRequestDeploymentMixin(FromRequestMixin):
 
         if not self.dial_application_id:
             raise InvalidRequestError(
-                f"The {_DIAL_APPLICATION_ID_HEADER} header isn't set"
+                f"The {DIAL_APPLICATION_ID} header isn't set"
             )
 
         if not self.base_url:
@@ -176,17 +194,15 @@ class FromRequestDeploymentMixin(FromRequestMixin):
         del headers["Authorization"]
 
         application_properties = None
-        props_header = headers.get(_DIAL_APPLICATION_PROPERTIES_HEADER)
+        props_header = headers.get(DIAL_APPLICATION_PROPERTIES)
         if props_header:
             try:
                 application_properties = json.loads(props_header)
             except json.JSONDecodeError:
                 raise InvalidRequestError(
-                    f"The value of {_DIAL_APPLICATION_PROPERTIES_HEADER} header isn't valid JSON"
+                    f"The value of {DIAL_APPLICATION_PROPERTIES} header isn't valid JSON"
                 )
 
-        application_id = headers.get(_DIAL_APPLICATION_ID_HEADER)
-
         return cls(
             **(await cls.get_request_body(request)),
             api_key_secret=SecretStr(api_key),
@@ -202,7 +218,14 @@ class FromRequestDeploymentMixin(FromRequestMixin):
             original_request=request,
             base_url=base_url,
             unreliable_dial_application_properties=application_properties,
-            dial_application_id=application_id,
+            dial_application_id=headers.get(DIAL_APPLICATION_ID),
+            conversation_id=headers.get(DIAL_CONVERSATION_ID),
+            job_title=headers.get(DIAL_JOB_TITLE),
+            cache_breakpoint_path=headers.get(DIAL_CACHE_BREAKPOINT_PATH),
+            cache_extra_metadata=headers.get(DIAL_CACHE_EXTRA_METADATA),
+            upstream_endpoint=headers.get(DIAL_UPSTREAM_ENDPOINT),
+            upstream_key=headers.get(DIAL_UPSTREAM_KEY),
+            upstream_extra_data=headers.get(DIAL_UPSTREAM_EXTRA_DATA),
         )
 
     @staticmethod

{aidial_sdk-0.37.0.dev7 → aidial_sdk-0.38.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-sdk"
-version = "0.37.0.dev7"
+version = "0.38.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"