Upload files

import hashlib
import os
import time

import requests

API = "https://byterover.ai/api/v1"
HEADERS = {"Authorization": f"Bearer {os.environ['BYTEROVER_API_KEY']}"}


def call(method, url, **kwargs):
    response = requests.request(method, url, headers=HEADERS, **kwargs)
    response.raise_for_status()
    return response.json()


def upload(drive_id, local_path, path):
    # 1. Start: the path in the drive, the size and the SHA-256 of the file.
    digest = hashlib.sha256()
    with open(local_path, "rb") as file:
        for block in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(block)
    plan = call("POST", f"{API}/drives/{drive_id}/uploads", json={
        "path": path,
        "size_bytes": os.path.getsize(local_path),
        "sha256": digest.hexdigest(),
    })

    # 2. Send each part to its address, and keep its ETag.
    sent = []
    parts, cursor = plan["parts"], plan["next_cursor"]
    with open(local_path, "rb") as file:
        while True:
            for part in parts:
                file.seek((part["part_number"] - 1) * plan["part_size"])
                put = requests.put(part["url"], data=file.read(plan["part_size"]))
                put.raise_for_status()
                sent.append({"part_number": part["part_number"], "etag": put.headers["ETag"]})
            if cursor is None:
                break
            page = call("GET", f"{API}/uploads/{plan['upload_id']}/parts", params={"cursor": cursor})
            parts, cursor = page["parts"], page["next_cursor"]

    # 3. Complete: the file becomes a document, and an operation starts.
    written = call("POST", f"{API}/uploads/{plan['upload_id']}/complete", json={"parts": sent})

    # 4. Follow the operation until the document is searchable.
    operation = written["operation"]
    while operation["status"] in ("queued", "running"):
        time.sleep((operation["retry_after_ms"] or 1000) / 1000)
        operation = call("GET", f"{API}/operations/{operation['id']}")
    if operation["status"] == "failed":
        raise RuntimeError(operation["error"]["message"])
    return written["document"]


document = upload("drv_4h7k2m9q8r1t5v6w3x0y", "customer-refund-policy.pdf", "/policies/customer-refund-policy.pdf")
print(document["id"])