Merge pull request #6 from weAreMusicAI/DEV-417-user-agent

DEV-417 Add custom User-Agent header and copyResultsTo method to Python SDK

Author: juliofortunato

README.md

@@ -81,26 +81,60 @@ song_url = music_ai.upload_file(file_path)
 
 ### Add a job
 
-Creates a new job and returns its correponding ID.
+Creates a new job and returns its corresponding ID.
 
 ```python
-def add_job(job_name: str, workflow_slug: str, params: Dict[str, str]) -> str
+def add_job(
+    job_name: str,
+    workflow_slug: str,
+    params: Dict[str, Any],
+    **options
+) -> str
 ```
 
 #### Example
 
 ```python
 song_url = "https://your-website.com/song.mp3"
+job_id = music_ai.add_job("job-1", "music-ai/isolate-drums", {
+    "inputUrl": song_url,
+})
+```
+
+Check the [documentation](https://music.ai/docs) for all the existing workflows and expected correspondent parameters.
+
+#### Custom storage
+
+You can optionally store outputs in your own storage by providing upload URLs. To do that, use the `copy_results_to` option, defining one upload URL for each output of the workflow.
+
+```python
 job_id = music_ai.add_job(
     "job-1",
-    "music-ai/generate-chords",
+    "music-ai/isolate-drums",
     {
-        "inputUrl": song_url
+        "inputUrl": song_url,
+    },
+    copy_results_to={
+        "Kick drum": "https://example.com/my-upload-url-1",
+        "Snare drum": "https://example.com/my-upload-url-2"
     }
 )
 ```
 
-Check the [documentation](https://music.ai/docs) for all the existing workflows and expected correspondent parameters.
+The example above uses the `music-ai/isolate-drums` workflow, which has 3 outputs, Kick drum, Snare drum, and Other. We have provided upload URLs for the first two. Since we haven't provided a URL for the third output, it will be stored in Music AI's storage, as usual.
+
+The JSON below contains the data for the job created above. Please note that Music AI doesn't provide download URLs for the outputs directed to your custom storage.
+
+```json
+{
+  // ...
+  "result": {
+    "Kick drum": "[custom storage]",
+    "Snare drum": "[custom storage]",
+    "Other": "https://cdn.music.ai/example/vocals.wav"
+  }
+}
+```
 
 ### Get a job
 

musicai_sdk/client.py

@@ -6,18 +6,32 @@
 import requests
 from requests.exceptions import HTTPError
 
-from .utils import extract_file_extension_from_url, extract_name_from_url
+from .utils import (
+    extract_file_extension_from_url,
+    extract_name_from_url,
+    get_user_agent,
+)
 
 
 class MusicAiClient:
-    def __init__(self, api_key, job_monitor_interval=2, save_output_to_folder=True):
+    def __init__(
+        self,
+        api_key,
+        job_monitor_interval=2,
+        save_output_to_folder=True,
+        user_agent=None,
+    ):
         self.api_key = api_key
         self.base_url = "https://api.music.ai/api"
         self.job_monitor_interval = job_monitor_interval
         self.save_output_to_folder = save_output_to_folder
+        self.user_agent = user_agent or get_user_agent()
 
     def get_headers(self):
-        return {"Authorization": self.api_key}
+        return {
+            "Authorization": self.api_key,
+            "User-Agent": self.user_agent,
+        }
 
     def upload_file(self, file_path):
         response = requests.get(f"{self.base_url}/upload", headers=self.get_headers())
@@ -79,8 +93,19 @@ def list_jobs(self, filters=None):
 
         return response.json()
 
-    def add_job(self, job_name, workflow_slug, params):
-        data = {"name": job_name, "workflow": workflow_slug, "params": params}
+    def add_job(self, job_name, workflow_slug, params, **options):
+        data = {
+            "name": job_name,
+            "workflow": workflow_slug,
+            "params": params,
+        }
+        
+        if "metadata" in options:
+            data["metadata"] = options["metadata"]
+            
+        if "copy_results_to" in options:
+            data["copyResultsTo"] = options["copy_results_to"]
+            
         response = requests.post(
             f"{self.base_url}/job", headers=self.get_headers(), json=data
         )

musicai_sdk/utils.py

@@ -1,4 +1,6 @@
 from urllib.parse import urlparse, unquote
+import platform
+from pathlib import Path
 
 
 def extract_name_from_url(file_url):
@@ -10,3 +12,21 @@ def extract_name_from_url(file_url):
 def extract_file_extension_from_url(url):
     name_from_url = extract_name_from_url(url)
     return name_from_url.split(".")[-1]
+
+
+def get_version():
+    package_path = Path(__file__).parent.parent / "pyproject.toml"
+    with open(package_path, "r") as f:
+        content = f.read()
+
+        for line in content.splitlines():
+            if line.startswith("version = "):
+                return line.split("=")[1].strip().strip("\"'")
+    return "unknown"
+
+
+def get_user_agent(environment="SDK-Python", version=None):
+    if version is None:
+        version = get_version()
+    platform_name = platform.system() or "unknown"
+    return f"Music.AI/{environment}/{version} ({platform_name})"

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "musicai-sdk"
-version = "1.0.0"
-description = "Python sdk for music.ai"
+version = "1.0.1"
+description = "Python SDK for Music.AI"
 authors = ["Music.ai Support <support@music.ai>"]
 license = "MIT"
 readme = "README.md"

tests/test_client.py

@@ -3,6 +3,7 @@
 import shutil
 import tempfile
 import uuid
+import requests
 
 import pytest
 
@@ -71,6 +72,38 @@ def test_add_job(self, client):
         assert job_id is not None
         assert re.match(UUID_REGEX, job_id)
 
+    def test_add_job_with_custom_storage(self, client):
+        """Test adding a job with custom storage"""
+        # Arrange
+        upload_response = requests.get(f"{client.base_url}/upload", headers=client.get_headers())
+        upload_response.raise_for_status()
+        upload_url = upload_response.json()["uploadUrl"]
+        
+        # Act
+        result = client.add_job(
+            "sdk-test-custom-storage",
+            "sdk-test",
+            {
+                "file": "https://music.ai/demo.ogg",
+            },
+            copy_results_to={
+                "file": upload_url
+            }
+        )
+        
+        # Assert
+        assert isinstance(result, dict)
+        job_id = result.get("id")
+        assert job_id is not None
+        assert re.match(UUID_REGEX, job_id)
+        
+        # Additional verification for custom storage
+        completed_job = client.wait_for_job_completion(job_id)
+        assert completed_job["status"] == "SUCCEEDED"
+        assert "result" in completed_job
+        assert "file" in completed_job["result"]
+        assert completed_job["result"]["file"] == "[custom storage]"
+    
     def test_get_job(self, client):
         """Test getting job details"""
         # Arrange
@@ -208,4 +241,4 @@ def test_get_application_info(self, client):
         assert isinstance(result, dict)
         assert "id" in result
         assert re.match(UUID_REGEX, result["id"])
-        assert "name" in result 
+        assert "name" in result