Upload a folder and search it

import hashlib
import os
import sys
import time

import requests

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

READABLE = {
    "pdf", "docx", "pptx", "xlsx", "html", "htm", "md", "markdown", "mdx", "txt", "text", "csv", "eml",
    "c", "cc", "cpp", "cs", "css", "go", "h", "hpp", "java", "js", "json", "jsx", "kt", "lua", "m", "mjs",
    "php", "proto", "py", "rb", "rs", "scala", "sh", "sql", "swift", "toml", "ts", "tsx", "vue", "xml",
    "yaml", "yml", "zig",
}


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):
    """Send one file in parts. Return the operation that makes it searchable."""
    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(),
    })
    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"]
    return call("POST", f"{API}/uploads/{plan['upload_id']}/complete", json={"parts": sent})["operation"]


def main(drive_id, folder, question):
    # 1. Upload each readable file. Keep the operation of each one.
    operations = {}
    for root, _, names in os.walk(folder):
        for name in names:
            local_path = os.path.join(root, name)
            if name.rsplit(".", 1)[-1].lower() not in READABLE or os.path.getsize(local_path) == 0:
                print(f"skipped   {local_path}")
                continue
            path = "/" + os.path.relpath(local_path, folder).replace(os.sep, "/")
            operations[path] = upload(drive_id, local_path, path)
            print(f"uploaded  {path}")

    # 2. Wait until each operation is done or failed.
    waiting = dict(operations)
    while waiting:
        time.sleep(2)
        for path, operation in list(waiting.items()):
            operation = call("GET", f"{API}/operations/{operation['id']}")
            if operation["status"] == "done":
                print(f"ready     {path}")
                del waiting[path]
            elif operation["status"] == "failed":
                print(f"failed    {path}: {operation['error']['message']}")
                del waiting[path]

    # 3. Ask.
    found = call("POST", f"{API}/search", json={"query": question, "drive_ids": [drive_id], "effort": "medium"})
    if not found["results"]:
        print("Nothing in the drive is near this question.")
    for result in found["results"]:
        location = result["location"]
        place = f"page {location['page']}, " if "page" in location else ""
        print(f"\n{place}{' > '.join(location['heading_path'])}\n{result['text']}")


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], sys.argv[3])

export BYTEROVER_API_KEY="brv_..."
python upload_folder.py drv_4h7k2m9q8r1t5v6w3x0y ./policies "Which policy section covers EU refunds?"

uploaded  /customer-refund-policy.pdf
uploaded  /procedures/supplier-onboarding.pdf
skipped   ./policies/logo.png
ready     /customer-refund-policy.pdf
ready     /procedures/supplier-onboarding.pdf

page 6, Refunds > European Union
Customers in the European Union may cancel an order within 14 days of delivery and receive a full refund.