Merge branch 'development' into feat/supported-extensible-routes

aidial_sdk/chat_completion/__init__.py

@@ -11,10 +11,9 @@
     Message,
     Request,
     Role,
-    Tool,
-    ToolCall,
-    ToolChoice,
 )
+from aidial_sdk.chat_completion.request import Stage as RequestStage
+from aidial_sdk.chat_completion.request import Tool, ToolCall, ToolChoice
 from aidial_sdk.chat_completion.response import Response
 from aidial_sdk.chat_completion.stage import Stage
 from aidial_sdk.deployment.tokenize import (

aidial_sdk/chat_completion/enums.py

@@ -1,14 +1,14 @@
 from enum import Enum
 
 
-class FinishReason(Enum):
+class FinishReason(str, Enum):
     STOP = "stop"
     LENGTH = "length"
     FUNCTION_CALL = "function_call"
     TOOL_CALLS = "tool_calls"
     CONTENT_FILTER = "content_filter"
 
 
-class Status(Enum):
+class Status(str, Enum):
     COMPLETED = "completed"
     FAILED = "failed"

aidial_sdk/chat_completion/request.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Mapping, Optional, Union
 
+from aidial_sdk.chat_completion.enums import Status
 from aidial_sdk.deployment.from_request_mixin import FromRequestDeploymentMixin
 from aidial_sdk.pydantic_v1 import (
     ConstrainedFloat,
@@ -21,7 +22,15 @@ class Attachment(ExtraForbidModel):
     reference_url: Optional[StrictStr] = None
 
 
+class Stage(ExtraForbidModel):
+    name: StrictStr
+    status: Status
+    content: Optional[StrictStr] = None
+    attachments: Optional[List[Attachment]] = None
+
+
 class CustomContent(ExtraForbidModel):
+    stages: Optional[List[Stage]] = None
     attachments: Optional[List[Attachment]] = None
     state: Optional[Any] = None
 
@@ -39,7 +48,7 @@ class ToolCall(ExtraForbidModel):
     function: FunctionCall
 
 
-class Role(Enum):
+class Role(str, Enum):
     SYSTEM = "system"
     USER = "user"
     ASSISTANT = "assistant"

aidial_sdk/embeddings/__init__.py

@@ -0,0 +1,8 @@
+from aidial_sdk.embeddings.base import Embeddings
+from aidial_sdk.embeddings.request import (
+    Attachment,
+    EmbeddingsMultiModalInput,
+    EmbeddingsRequestCustomFields,
+    Request,
+)
+from aidial_sdk.embeddings.response import Embedding, Response, Usage

aidial_sdk/embeddings/request.py

@@ -1,6 +1,6 @@
-from enum import Enum
 from typing import List, Literal, Optional, Union
 
+from aidial_sdk.chat_completion.request import Attachment
 from aidial_sdk.deployment.from_request_mixin import FromRequestDeploymentMixin
 from aidial_sdk.pydantic_v1 import StrictFloat, StrictInt, StrictStr
 from aidial_sdk.utils.pydantic import ExtraForbidModel
@@ -16,18 +16,18 @@ class AzureEmbeddingsRequest(ExtraForbidModel):
     user: Optional[StrictStr] = None
 
 
-class DialEmbeddingsType(str, Enum):
-    SYMMETRIC = "symmetric"
-    DOCUMENT = "document"
-    QUERY = "query"
-
-
 class EmbeddingsRequestCustomFields(ExtraForbidModel):
-    type: Optional[DialEmbeddingsType] = None
+    type: Optional[StrictStr] = None
     instruction: Optional[StrictStr] = None
 
 
+EmbeddingsMultiModalInput = Union[
+    StrictStr, Attachment, List[Union[StrictStr, Attachment]]
+]
+
+
 class EmbeddingsRequest(AzureEmbeddingsRequest):
+    custom_input: Optional[List[EmbeddingsMultiModalInput]] = None
     custom_fields: Optional[EmbeddingsRequestCustomFields] = None
 
 

aidial_sdk/header_propagator.py

@@ -1,6 +1,6 @@
 import types
 from contextvars import ContextVar
-from typing import Optional
+from typing import MutableMapping, Optional
 
 import aiohttp
 import httpx
@@ -79,21 +79,12 @@ async def _on_aiohttp_request_start(
         trace_config_ctx: types.SimpleNamespace,
         params: aiohttp.TraceRequestStartParams,
     ):
-        if not str(params.url).startswith(self._dial_url):
-            return
-
-        api_key_val = self._api_key.get()
-        if api_key_val:
-            params.headers["api-key"] = api_key_val
+        self._modify_headers(str(params.url), params.headers)
 
     def _instrument_requests(self):
         def instrumented_send(wrapped, instance, args, kwargs):
             request: requests.PreparedRequest = args[0]
-            if request.url and request.url.startswith(self._dial_url):
-                api_key_val = self._api_key.get()
-                if api_key_val:
-                    request.headers["api-key"] = api_key_val
-
+            self._modify_headers(request.url or "", request.headers)
             return wrapped(*args, **kwargs)
 
         wrapt.wrap_function_wrapper(requests.Session, "send", instrumented_send)
@@ -102,10 +93,7 @@ def _instrument_httpx(self):
 
         def instrumented_build_request(wrapped, instance, args, kwargs):
             request: httpx.Request = wrapped(*args, **kwargs)
-            if request.url and str(request.url).startswith(self._dial_url):
-                api_key_val = self._api_key.get()
-                if api_key_val:
-                    request.headers["api-key"] = api_key_val
+            self._modify_headers(str(request.url), request.headers)
             return request
 
         wrapt.wrap_function_wrapper(
@@ -115,3 +103,21 @@ def instrumented_build_request(wrapped, instance, args, kwargs):
         wrapt.wrap_function_wrapper(
             httpx.AsyncClient, "build_request", instrumented_build_request
         )
+
+    def _modify_headers(
+        self, url: str, headers: MutableMapping[str, str]
+    ) -> None:
+        if url.startswith(self._dial_url):
+            api_key = self._api_key.get()
+            if api_key:
+                old_api_key = headers.get("api-key")
+                old_authz = headers.get("Authorization")
+
+                if (
+                    old_api_key
+                    and old_authz
+                    and old_authz == f"Bearer {old_api_key}"
+                ):
+                    headers["Authorization"] = f"Bearer {api_key}"
+
+                headers["api-key"] = api_key

tests/header_propagation/client.py

@@ -21,33 +21,35 @@ class Library(str, Enum):
 class Request(BaseModel):
     url: str
     lib: Library
+    headers: dict
 
 
 @app.post("/")
 async def handle(request: Request):
     url = request.url
     lib = request.lib
+    headers = request.headers
 
     if lib == Library.requests:
-        response = requests.get(url)
+        response = requests.get(url, headers=headers)
         status_code = response.status_code
         content = response.json()
 
     elif lib == Library.httpx_async:
         async with httpx.AsyncClient() as client:
-            response = await client.get(url)
+            response = await client.get(url, headers=headers)
             status_code = response.status_code
             content = response.json()
 
     elif lib == Library.httpx_sync:
         with httpx.Client() as client:
-            response = client.get(url)
+            response = client.get(url, headers=headers)
             status_code = response.status_code
             content = response.json()
 
     elif lib == Library.aiohttp:
         async with aiohttp.ClientSession() as session:
-            async with session.get(url) as response:
+            async with session.get(url, headers=headers) as response:
                 status_code = response.status
                 content = await response.json()
 

tests/test_embeddings.py

@@ -5,7 +5,6 @@
 from aidial_sdk import DIALApp
 from tests.applications.simple_embeddings import SimpleEmbeddings
 from tests.utils.endpoint_test import TestCase, run_endpoint_test
-from tests.utils.errors import invalid_request_error
 
 deployment = "test-app"
 app = DIALApp().add_embeddings(deployment, SimpleEmbeddings())
@@ -37,16 +36,6 @@
         {"input": "a", "custom_fields": {"type": "query"}},
         expected_response_1,
     ),
-    TestCase(
-        app,
-        deployment,
-        "embeddings",
-        {"input": "a", "custom_fields": {"type": "hello"}},
-        invalid_request_error(
-            "custom_fields.type",
-            "value is not a valid enumeration member; permitted: 'symmetric', 'document', 'query'",
-        ),
-    ),
     TestCase(
         app,
         deployment,