aidial_adapter_dial/transformer.py

import logging
import re
from typing import Any, Callable, Coroutine, Self

import aiohttp
from pydantic import BaseModel

from aidial_adapter_dial.utils.app_data import AppData
from aidial_adapter_dial.utils.storage import FileStorage

log = logging.getLogger(__name__)


class AttachmentTransformer(BaseModel):
    local_storage: FileStorage
    local_user_bucket: str
    local_appdata: str

    remote_storage: FileStorage
    remote_user_bucket: str

    proxy_mode: bool

    @classmethod
    async def create(
        cls, remote_storage: FileStorage, local_storage: FileStorage
    ) -> Self:
        async with aiohttp.ClientSession() as session:
            local = await local_storage.get_bucket(session)

            local_appdata = local.get("appdata")
            if local_appdata is None:
                raise ValueError(
                    "The local appdata bucket is expected to be set"
                )
            local_user_bucket = AppData.parse(local_appdata).user_bucket

            remote = await remote_storage.get_bucket(session)
            remote_appdata = remote.get("appdata")
            if remote_appdata is None:
                remote_user_bucket = remote["bucket"]
            else:
                remote_user_bucket = AppData.parse(remote_appdata).user_bucket

        proxy_mode = (
            remote_storage.dial_url == local_storage.dial_url
            and remote_storage.api_key == local_storage.api_key
        )

        log.debug(f"proxy_mode: {proxy_mode}")

        return cls(
            remote_storage=remote_storage,
            remote_user_bucket=remote_user_bucket,
            local_storage=local_storage,
            local_user_bucket=local_user_bucket,
            local_appdata=local_appdata,
            proxy_mode=proxy_mode,
        )

    def get_remote_url(self, local_url: str) -> str:
        """
        user/app files:
            if proxy_mode:
                < files/LOCAL_USER_BUCKET/PATH
                > files/REMOTE_USER_BUCKET/PATH
            else:
                < files/LOCAL_USER_BUCKET/PATH
                > files/REMOTE_USER_BUCKET/LOCAL_USER_BUCKET/PATH
        """

        if self.proxy_mode:
            return local_url

        if not local_url.startswith(f"files/{self.local_user_bucket}/"):
            raise ValueError(f"Unexpected local URL: {local_url!r}")

        return f"files/{self.remote_user_bucket}/{local_url.removeprefix('files/')}"

    def get_local_url(self, remote_url: str) -> str:
        """
        user/app files uploaded from local to remote earlier (reverse of get_remote_url):
            if proxy_mode:
                < files/REMOTE_USER_BUCKET/PATH
                > files/LOCAL_USER_BUCKET/PATH
            else:
                < files/REMOTE_USER_BUCKET/LOCAL_USER_BUCKET/PATH
                > files/LOCAL_USER_BUCKET/PATH

        created by remote (user):
            < files/REMOTE_USER_BUCKET/appdata/REMOTE_APP_NAME/PATH
            > files/LOCAL_USER_BUCKET/appdata/LOCAL_APP_NAME/PATH

        created by remote (app):
            < files/REMOTE_APP_BUCKET/PATH
            > This means an application has a bug in it.
                We reject such URLs right away since there is no way
                to read an application file by a user.
        """

        if not remote_url.startswith(f"files/{self.remote_user_bucket}/"):
            raise ValueError(
                f"The remote file ({remote_url!r}) is expected "
                f"to be uploaded to the remote user bucket ({self.remote_user_bucket!r})"
            )

        remote_path = remote_url.removeprefix(
            f"files/{self.remote_user_bucket}/"
        )

        if remote_path.startswith("appdata/"):
            regex = r"appdata/([^/]+)/(.+)"
            match = re.match(regex, remote_path)
            if match is None:
                raise ValueError(f"Invalid remote appdata path: {remote_url!r}")
            _remote_app_name, path = match.groups()
            return f"files/{self.local_appdata}/{path}"

        if not self.proxy_mode:
            if remote_path.startswith(f"{self.local_user_bucket}/"):
                path = remote_path.removeprefix(f"{self.local_user_bucket}/")
                return f"files/{self.local_user_bucket}/{path}"

            raise ValueError(
                f"The remote file ({remote_url!r}) is expected to be uploaded either "
                "to remote appdata path or "
                "to a local user bucket subpath of remote user bucket."
            )
        else:
            return remote_url

    async def modify_request_attachment(self, attachment: dict) -> None:
        if (ref_url := attachment.get("reference_url")) and (
            local_ref_url := self.local_storage.to_dial_url(ref_url)
        ):
            remote_ref_url = self.get_remote_url(local_ref_url)
            attachment["reference_url"] = remote_ref_url

        if (url := attachment.get("url")) and (
            local_url := self.local_storage.to_dial_url(url)
        ):
            remote_url = self.get_remote_url(local_url)

            await download_and_upload_file(
                self.local_storage,
                local_url,
                self.remote_storage,
                remote_url,
                attachment.get("type"),
            )

            attachment["url"] = remote_url

            log.debug(
                f"uploaded from local to remote: from {local_url!r} to {remote_url!r}"
            )

    async def modify_response_attachment(self, attachment: dict) -> None:
        if (ref_url := attachment.get("reference_url")) and (
            remote_ref_url := self.remote_storage.to_dial_url(ref_url)
        ):
            local_ref_url = self.get_local_url(remote_ref_url)
            attachment["reference_url"] = local_ref_url

        if (url := attachment.get("url")) and (
            remote_url := self.remote_storage.to_dial_url(url)
        ):
            local_url = self.get_local_url(remote_url)

            await download_and_upload_file(
                self.remote_storage,
                remote_url,
                self.local_storage,
                local_url,
                attachment.get("type"),
            )
            attachment["url"] = local_url

            log.debug(
                f"uploaded from remote to local: from {remote_url!r} to {local_url!r}"
            )

    async def modify_request(self, request: dict) -> dict:
        if "messages" in request:
            messages = request["messages"]
            for message in messages:
                await modify_message(message, self.modify_request_attachment)
        return request

    async def modify_response_chunk(self, response: dict) -> dict:
        choices = response.get("choices")
        if choices is None:
            return response

        for choice in choices:
            if "delta" in choice:
                await modify_message(
                    choice["delta"], self.modify_response_attachment
                )

        return response

    async def modify_response(self, response: dict) -> dict:
        choices = response.get("choices")
        if choices is None:
            return response

        for choice in choices:
            if "message" in choice:
                await modify_message(
                    choice["message"], self.modify_response_attachment
                )

        return response


async def download_and_upload_file(
    src_storage: FileStorage,
    src_url: str,
    dest_storage: FileStorage,
    dest_url: str,
    content_type: str | None,
):
    log.debug(f"downloading from {src_url!r} and uploading to {dest_url!r}")

    if src_url != dest_url:
        if _is_directory(src_url):
            raise ValueError("Directories aren't yet supported")

        async with aiohttp.ClientSession() as session:
            content = await src_storage.download(src_url, session)
            await dest_storage.upload(dest_url, content_type, content, session)


async def modify_message(
    message: dict,
    modify_attachment: Callable[[dict], Coroutine[Any, Any, None]],
) -> None:
    cc = message.get("custom_content")
    if cc is None:
        return
    attachments = cc.get("attachments")
    if attachments is None:
        return
    for attachment in attachments:
        await modify_attachment(attachment)


def _is_directory(url: str) -> bool:
    return url[-1] == "/"