feat: support DALL-E 3 (#26) (#30)

Author: vladisavvv

README.md

@@ -57,6 +57,9 @@ Copy `.env.example` to `.env` and customize it for your environment:
 |WEB_CONCURRENCY|1|Number of workers for the server|
 |AZURE_API_VERSION|2023-03-15-preview|The version API for requests to Azure OpenAI API|
 |MODEL_ALIASES|{}|Mapping request's deployment_id to [model name of tiktoken](https://github.com/openai/tiktoken/blob/main/tiktoken/model.py) for correct calculate of tokens. Example: `{"gpt-35-turbo":"gpt-3.5-turbo-0301"}`|
+|DIAL_USE_FILE_STORAGE|False|Save image model artifacts to DIAL File storage (DALL-E images are uploaded to the files storage and its base64 encodings are replaced with links to the storage)|
+|DIAL_URL||URL of the core DIAL server (required when DIAL_USE_FILE_STORAGE=True)|
+|DIAL_API_KEY||API Key for DIAL File storage (required when DIAL_USE_FILE_STORAGE=True)|
 
 ### Docker
 

aidial_adapter_openai/app.py

@@ -9,14 +9,19 @@
 from openai import ChatCompletion, Embedding, error
 from openai.openai_object import OpenAIObject
 
+from aidial_adapter_openai.images import text_to_image_chat_completion
 from aidial_adapter_openai.openai_override import OpenAIException
+from aidial_adapter_openai.utils.deployment_classifier import (
+    is_text_to_image_deployment,
+)
 from aidial_adapter_openai.utils.exceptions import HTTPException
 from aidial_adapter_openai.utils.log_config import LogConfig
 from aidial_adapter_openai.utils.parsers import (
     ApiType,
     parse_body,
     parse_upstream,
 )
+from aidial_adapter_openai.utils.storage import FileStorage
 from aidial_adapter_openai.utils.streaming import generate_stream
 from aidial_adapter_openai.utils.tokens import discard_messages
 from aidial_adapter_openai.utils.versions import compare_versions
@@ -26,6 +31,22 @@
 model_aliases: Dict[str, str] = json.loads(os.getenv("MODEL_ALIASES", "{}"))
 azure_api_version = os.getenv("AZURE_API_VERSION", "2023-03-15-preview")
 
+dial_use_file_storage = (
+    os.getenv("DIAL_USE_FILE_STORAGE", "false").lower() == "true"
+)
+
+file_storage = None
+if dial_use_file_storage:
+    dial_url = os.getenv("DIAL_URL")
+    dial_api_key = os.getenv("DIAL_API_KEY")
+
+    if not dial_url or not dial_api_key:
+        raise ValueError(
+            "DIAL_URL and DIAL_API_KEY environment variables must be initialized if DIAL_USE_FILE_STORAGE is true"
+        )
+
+    file_storage = FileStorage(dial_url, "dalle", dial_api_key)
+
 
 async def handle_exceptions(call):
     try:
@@ -46,10 +67,16 @@ async def chat_completion(deployment_id: str, request: Request):
 
     is_stream = data.get("stream", False)
     openai_model_name = model_aliases.get(deployment_id, deployment_id)
-    dial_api_key = request.headers["X-UPSTREAM-KEY"]
+    api_key = request.headers["X-UPSTREAM-KEY"]
+    upstream_endpoint = request.headers["X-UPSTREAM-ENDPOINT"]
+
+    if is_text_to_image_deployment(deployment_id):
+        return await text_to_image_chat_completion(
+            data, upstream_endpoint, api_key, is_stream, file_storage
+        )
 
     api_base, upstream_deployment = parse_upstream(
-        request.headers["X-UPSTREAM-ENDPOINT"], ApiType.CHAT_COMPLETION
+        upstream_endpoint, ApiType.CHAT_COMPLETION
     )
 
     api_version = azure_api_version
@@ -87,7 +114,7 @@ async def chat_completion(deployment_id: str, request: Request):
     response = await handle_exceptions(
         ChatCompletion().acreate(
             engine=upstream_deployment,
-            api_key=dial_api_key,
+            api_key=api_key,
             api_base=api_base,
             api_type="azure",
             api_version=api_version,
@@ -127,15 +154,15 @@ async def chat_completion(deployment_id: str, request: Request):
 async def embedding(deployment_id: str, request: Request):
     data = await parse_body(request)
 
-    dial_api_key = request.headers["X-UPSTREAM-KEY"]
+    api_key = request.headers["X-UPSTREAM-KEY"]
     api_base, upstream_deployment = parse_upstream(
         request.headers["X-UPSTREAM-ENDPOINT"], ApiType.EMBEDDING
     )
 
     return await handle_exceptions(
         Embedding().acreate(
             deployment_id=upstream_deployment,
-            api_key=dial_api_key,
+            api_key=api_key,
             api_base=api_base,
             api_type="azure",
             api_version=azure_api_version,

aidial_adapter_openai/images.py

@@ -0,0 +1,147 @@
+from typing import Any, AsyncGenerator, Optional
+
+import aiohttp
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+
+from aidial_adapter_openai.utils.exceptions import HTTPException
+from aidial_adapter_openai.utils.storage import FileStorage, upload_base64_file
+from aidial_adapter_openai.utils.streaming import (
+    END_CHUNK,
+    build_chunk,
+    chunk_format,
+    generate_id,
+)
+
+IMG_USAGE = {
+    "completion_tokens": 0,
+    "prompt_tokens": 1,
+    "total_tokens": 1,
+}
+
+
+async def generate_image(api_url: str, api_key: str, user_prompt: str) -> Any:
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            api_url,
+            json={"prompt": user_prompt, "response_format": "b64_json"},
+            headers={"api-key": api_key},
+        ) as response:
+            status_code = response.status
+
+            data = await response.json()
+
+            if status_code == 200:
+                return data
+            else:
+                return JSONResponse(content=data, status_code=status_code)
+
+
+def build_custom_content(base64_image: str, revised_prompt: str) -> Any:
+    return {
+        "custom_content": {
+            "attachments": [
+                {"title": "Revised prompt", "data": revised_prompt},
+                {"title": "Image", "type": "image/png", "data": base64_image},
+            ]
+        }
+    }
+
+
+async def generate_stream(
+    id: str, created: str, custom_content: Any
+) -> AsyncGenerator[Any, Any]:
+    yield chunk_format(
+        build_chunk(id, None, {"role": "assistant"}, created, True)
+    )
+
+    yield chunk_format(build_chunk(id, None, custom_content, created, True))
+
+    yield chunk_format(
+        build_chunk(id, "stop", {}, created, True, usage=IMG_USAGE)
+    )
+
+    yield END_CHUNK
+
+
+def get_user_prompt(data: Any):
+    if (
+        "messages" not in data
+        or len(data["messages"]) == 0
+        or "content" not in data["messages"][-1]
+        or not data["messages"][-1]
+    ):
+        raise HTTPException(
+            "Your request is invalid", 400, "invalid_request_error"
+        )
+
+    return data["messages"][-1]["content"]
+
+
+async def move_attachments_data_to_storage(
+    custom_content: Any, file_storage: FileStorage
+):
+    for attachment in custom_content["custom_content"]["attachments"]:
+        if (
+            "data" not in attachment
+            or "type" not in attachment
+            or not attachment["type"].startswith("image/")
+        ):
+            continue
+
+        file_metadata = await upload_base64_file(
+            file_storage, attachment["data"], attachment["type"]
+        )
+        image_url = file_metadata["path"] + "/" + file_metadata["name"]
+
+        del attachment["data"]
+        attachment["url"] = image_url
+
+
+async def text_to_image_chat_completion(
+    data: Any,
+    upstream_endpoint: str,
+    api_key: str,
+    is_stream: bool,
+    file_storage: Optional[FileStorage],
+) -> Response:
+    if data.get("n", 1) > 1:
+        raise HTTPException(
+            status_code=422,
+            message="The deployment doesn't support n > 1",
+            type="invalid_request_error",
+        )
+
+    api_url = upstream_endpoint + "?api-version=2023-12-01-preview"
+    user_prompt = get_user_prompt(data)
+    model_response = await generate_image(api_url, api_key, user_prompt)
+
+    if isinstance(model_response, JSONResponse):
+        return model_response
+
+    base64_image = model_response["data"][0]["b64_json"]
+    revised_prompt = model_response["data"][0]["revised_prompt"]
+
+    id = generate_id()
+    created = model_response["created"]
+
+    custom_content = build_custom_content(base64_image, revised_prompt)
+
+    if file_storage is not None:
+        await move_attachments_data_to_storage(custom_content, file_storage)
+
+    if not is_stream:
+        return JSONResponse(
+            content=build_chunk(
+                id,
+                "stop",
+                {**custom_content, "role": "assistant"},
+                created,
+                False,
+                usage=IMG_USAGE,
+            )
+        )
+    else:
+        return StreamingResponse(
+            generate_stream(id, created, custom_content),
+            media_type="text/event-stream",
+        )

aidial_adapter_openai/utils/deployment_classifier.py

@@ -0,0 +1,2 @@
+def is_text_to_image_deployment(deployment_id: str):
+    return deployment_id.lower() == "dalle3"

aidial_adapter_openai/utils/storage.py

@@ -0,0 +1,70 @@
+import base64
+import hashlib
+import io
+from typing import TypedDict
+
+import aiohttp
+
+from aidial_adapter_openai.utils.log_config import logger
+
+
+class FileMetadata(TypedDict):
+    name: str
+    type: str
+    path: str
+    contentLength: int
+    contentType: str
+
+
+class FileStorage:
+    base_url: str
+    api_key: str
+
+    def __init__(self, dial_url: str, base_dir: str, api_key: str):
+        self.base_url = f"{dial_url}/v1/files/{base_dir}"
+        self.api_key = api_key
+
+    def auth_headers(self) -> dict[str, str]:
+        return {"api-key": self.api_key}
+
+    @staticmethod
+    def to_form_data(
+        filename: str, content_type: str, content: bytes
+    ) -> aiohttp.FormData:
+        data = aiohttp.FormData()
+        data.add_field(
+            "file",
+            io.BytesIO(content),
+            filename=filename,
+            content_type=content_type,
+        )
+        return data
+
+    async def upload(
+        self, filename: str, content_type: str, content: bytes
+    ) -> FileMetadata:
+        async with aiohttp.ClientSession() as session:
+            data = FileStorage.to_form_data(filename, content_type, content)
+            async with session.post(
+                self.base_url,
+                data=data,
+                headers=self.auth_headers(),
+            ) as response:
+                response.raise_for_status()
+                meta = await response.json()
+                logger.debug(
+                    f"Uploaded file: path={self.base_url}, file={filename}, metadata={meta}"
+                )
+                return meta
+
+
+def _hash_digest(string: str) -> str:
+    return hashlib.sha256(string.encode()).hexdigest()
+
+
+async def upload_base64_file(
+    storage: FileStorage, data: str, content_type: str
+) -> FileMetadata:
+    filename = _hash_digest(data)
+    content: bytes = base64.b64decode(data)
+    return await storage.upload(filename, content_type, content)

aidial_adapter_openai/utils/streaming.py

@@ -17,6 +17,33 @@ def chunk_format(data: str | Mapping[str, Any]):
         return "data: " + json.dumps(data, separators=(",", ":")) + "\n\n"
 
 
+def generate_id():
+    return "chatcmpl-" + str(uuid4())
+
+
+def build_chunk(
+    id: str,
+    finish_reason: Optional[str],
+    delta: Any,
+    created: str,
+    is_stream,
+    **extra
+):
+    return {
+        "id": id,
+        "object": "chat.completion.chunk" if is_stream else "chat.completion",
+        "created": created,
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": finish_reason,
+                "delta": delta,
+            }
+        ],
+        **extra,
+    }
+
+
 END_CHUNK = chunk_format("[DONE]")
 
 
@@ -83,7 +110,7 @@ async def generate_stream(
 
             yield chunk_format(
                 {
-                    "id": "chatcmpl-" + str(uuid4()),
+                    "id": generate_id(),
                     "object": "chat.completion.chunk",
                     "created": str(int(time())),
                     "model": deployment,