diff --git a/README.md b/README.md index f476c12..0f162eb 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ Copy `.env.example` to `.env` and customize it for your environment: |DEFAULT_REGION||AWS region e.g. "us-east-1"| |LOG_LEVEL|INFO|Log level. Use DEBUG for dev purposes and INFO in prod| |AIDIAL_LOG_LEVEL|WARNING|AI DIAL SDK log level| -|DIAL_USE_FILE_STORAGE|False|Save model artifacts to DIAL File storage (particularly, Stability images are uploaded to the files storage and their base64 encodings are replaced with links to the storage). The creds for the file storage must be passed in `Authorization` header of the incoming request. The file storage won't be used if the header isn't set.| +|DIAL_USE_FILE_STORAGE|False|Save model artifacts to DIAL File storage (particularly, Stability images are uploaded to the files storage and their base64 encodings are replaced with links to the storage). The creds for the file storage must be passed in `api-key` header of the incoming request. The file storage won't be used if the header isn't set.| |DIAL_URL||URL of the core DIAL server (required when DIAL_USE_FILE_STORAGE=True)| |WEB_CONCURRENCY|1|Number of workers for the server| |TEST_SERVER_URL|http://0.0.0.0:5001|Server URL used in the integration tests| diff --git a/aidial_adapter_bedrock/app.py b/aidial_adapter_bedrock/app.py index 13ef690..21e2ba6 100644 --- a/aidial_adapter_bedrock/app.py +++ b/aidial_adapter_bedrock/app.py @@ -23,7 +23,7 @@ os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") is not None ) -SERVICE_NAME = os.environ.get("OTEL_SERVICE_NAME", "aidial-bedrock") +SERVICE_NAME = os.environ.get("OTEL_SERVICE_NAME", "dial-bedrock") app = DIALApp( description="AWS Bedrock adapter for DIAL API", diff --git a/aidial_adapter_bedrock/dial_api/storage.py b/aidial_adapter_bedrock/dial_api/storage.py index 5765f08..047c0b3 100644 --- a/aidial_adapter_bedrock/dial_api/storage.py +++ b/aidial_adapter_bedrock/dial_api/storage.py @@ -1,6 +1,7 @@ import base64 import hashlib import io +import mimetypes from typing import Mapping, Optional, TypedDict import aiohttp @@ -21,11 +22,16 @@ class FileMetadata(TypedDict): contentType: str +class Bucket(TypedDict): + bucket: str + appdata: str + + class FileStorage: dial_url: str upload_dir: str auth: Auth - bucket: Optional[str] + bucket: Optional[Bucket] def __init__(self, dial_url: str, upload_dir: str, auth: Auth): self.dial_url = dial_url @@ -33,15 +39,15 @@ def __init__(self, dial_url: str, upload_dir: str, auth: Auth): self.auth = auth self.bucket = None - async def _get_bucket(self, session: aiohttp.ClientSession) -> str: + async def _get_bucket(self, session: aiohttp.ClientSession) -> Bucket: if self.bucket is None: async with session.get( f"{self.dial_url}/v1/bucket", headers=self.auth.headers, ) as response: response.raise_for_status() - body = await response.json() - self.bucket = body["bucket"] + self.bucket = await response.json() + log.debug(f"bucket: {self.bucket}") return self.bucket @@ -63,9 +69,12 @@ async def upload( ) -> FileMetadata: async with aiohttp.ClientSession() as session: bucket = await self._get_bucket(session) + + appdata = bucket["appdata"] + ext = mimetypes.guess_extension(content_type) or "" + url = f"{self.dial_url}/v1/files/{appdata}/{self.upload_dir}/{filename}{ext}" + data = FileStorage._to_form_data(filename, content_type, content) - ext = _get_extension(content_type) or "" - url = f"{self.dial_url}/v1/files/{bucket}/{self.upload_dir}/{filename}{ext}" async with session.put( url=url, @@ -89,12 +98,6 @@ def _compute_hash_digest(file_content: str) -> str: return hashlib.sha256(file_content.encode()).hexdigest() -def _get_extension(content_type: str) -> Optional[str]: - if content_type.startswith("image/"): - return "." + content_type[len("image/") :] - return None - - DIAL_USE_FILE_STORAGE = get_env_bool("DIAL_USE_FILE_STORAGE", False) DIAL_URL: Optional[str] = None @@ -110,7 +113,7 @@ def create_file_storage( if not DIAL_USE_FILE_STORAGE or DIAL_URL is None: return None - auth = Auth.from_headers("authorization", headers) + auth = Auth.from_headers("api-key", headers) if auth is None: log.debug( "The request doesn't have required headers to use the DIAL file storage. " diff --git a/aidial_adapter_bedrock/llm/model/stability.py b/aidial_adapter_bedrock/llm/model/stability.py index 3e145ea..2b36ede 100644 --- a/aidial_adapter_bedrock/llm/model/stability.py +++ b/aidial_adapter_bedrock/llm/model/stability.py @@ -102,9 +102,7 @@ class StabilityAdapter(ChatModel): @classmethod def create(cls, client: Bedrock, model: str, headers: Mapping[str, str]): - storage: Optional[FileStorage] = create_file_storage( - "images/stable-diffusion", headers - ) + storage: Optional[FileStorage] = create_file_storage("images", headers) return cls( client=client, model=model,