Building a Python client

import json
import mimetypes
import os
import time
from pathlib import Path, PurePosixPath

import requests


PHYLO_BASE_URL = os.getenv("PHYLO_BASE_URL", "https://api.phylo.bio/v1")
REQUEST_TIMEOUT_SECONDS = 60
POLL_INTERVAL_SECONDS = 15
TASK_TIMEOUT_SECONDS = 24 * 60 * 60
TERMINAL_TASK_STATUSES = {"completed", "failed", "cancelled"}


class PhyloClient:
    def __init__(self, api_key=None, api_base=None):
        self.api_base = (api_base or PHYLO_BASE_URL).rstrip("/")
        self.session = requests.Session()
        self.session.headers["Authorization"] = (
            f"Bearer {api_key or os.environ['PHYLO_API_KEY']}"
        )

    def _request(self, method, path, **kwargs):
        response = self.session.request(
            method,
            f"{self.api_base}{path}",
            timeout=REQUEST_TIMEOUT_SECONDS,
            **kwargs,
        )
        response.raise_for_status()
        return response.json() if response.content else {}

    def list_projects(self):
        return self._request("GET", "/projects")

    def create_project(self, title, description=None):
        body = {"title": title}
        if description is not None:
            body["description"] = description
        return self._request("POST", "/projects", json=body)

    def create_task(
        self,
        project_id,
        title=None,
        initial_messages=None,
        model="max",
        auto_mode=True,
    ):
        body = {"project_id": project_id, "auto_mode": auto_mode}
        if title is not None:
            body["title"] = title
        if initial_messages is not None:
            body["initial_messages"] = initial_messages
        if model is not None:
            body["model"] = model
        return self._request("POST", "/experimental/tasks", json=body)

    def upload_file(self, project_id, path, mime_type=None):
        path = Path(path)
        mime_type = (
            mime_type
            or mimetypes.guess_type(path.name)[0]
            or "application/octet-stream"
        )
        intent = self._request(
            "POST",
            "/experimental/files",
            json={
                "project_id": project_id,
                "filename": path.name,
                "size_bytes": path.stat().st_size,
                "mime_type": mime_type,
            },
        )
        if intent["upload_type"] != "single":
            raise RuntimeError("This client only handles single-part uploads.")

        with path.open("rb") as file_handle:
            upload = requests.post(
                intent["upload_url"],
                data=intent["fields"],
                files={"file": (path.name, file_handle, mime_type)},
                timeout=REQUEST_TIMEOUT_SECONDS,
            )
        upload.raise_for_status()

        self._request(
            "POST",
            f"/experimental/files/{intent['file_id']}/finalize",
            json={},
        )
        return intent["file_id"]

    def send_message(self, task_id, content, file_ids=None):
        return self._request(
            "POST",
            f"/experimental/tasks/{task_id}/messages",
            json={"content": content, "file_ids": file_ids or []},
        )

    def get_task(self, task_id):
        return self._request("GET", f"/experimental/tasks/{task_id}")

    def wait_for_task(
        self,
        task_id,
        timeout_seconds=TASK_TIMEOUT_SECONDS,
        poll_interval_seconds=POLL_INTERVAL_SECONDS,
    ):
        deadline = time.monotonic() + timeout_seconds
        while time.monotonic() < deadline:
            task = self.get_task(task_id)
            status = task["status"]
            print(f"Task status: {status}")
            if status in TERMINAL_TASK_STATUSES:
                if status != "completed":
                    raise RuntimeError(f"Task ended with status: {status}")
                return task
            time.sleep(poll_interval_seconds)
        raise TimeoutError(f"Task did not finish within {timeout_seconds} seconds")

    def download_results(self, task_id, output_directory):
        output_directory = Path(output_directory).resolve()
        output_directory.mkdir(parents=True, exist_ok=True)

        after_id = None
        downloaded_count = 0
        while True:
            params = {"with_urls": "true"}
            if after_id is not None:
                params["after_id"] = after_id
            page = self._request(
                "GET",
                f"/experimental/tasks/{task_id}/results",
                params=params,
            )

            for result in page["data"]:
                name = result["name"]
                relative_path = PurePosixPath(name)
                if (
                    not relative_path.parts
                    or relative_path.is_absolute()
                    or ".." in relative_path.parts
                    or "\\" in name
                ):
                    raise ValueError(f"Unsafe result path: {name}")

                output_path = output_directory.joinpath(*relative_path.parts)
                output_path.parent.mkdir(parents=True, exist_ok=True)
                if output_directory not in output_path.resolve().parents:
                    raise ValueError(f"Unsafe result path: {name}")
                if output_path.exists():
                    raise FileExistsError(f"Result path already exists: {output_path}")

                response = requests.get(
                    result["download_url"],
                    timeout=REQUEST_TIMEOUT_SECONDS,
                )
                response.raise_for_status()
                output_path.write_bytes(response.content)
                downloaded_count += 1
                print(f"Downloaded: {output_path}")

            if not page["has_more"]:
                break
            after_id = page["last_id"]

        if downloaded_count == 0:
            raise RuntimeError("The task completed without result files.")

        return output_directory


if __name__ == "__main__":
    print(json.dumps(PhyloClient().list_projects(), indent=2))