fix: propagating per request key (#57)

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|

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",

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,27 +22,32 @@ 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
         self.upload_dir = upload_dir
         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. "

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,