spb/doc-api

1"""xAI Files API (cheap: upload/list/get/content/public-url/delete), collections management (LIVE_DISCOVERED on api.x.ai) and Batch API (RUN_BATCH_TESTS)."""2from __future__ import annotations34import base645import struct6import time7import uuid8import zlib910import pytest111213def _multipart(fields, filename, content, ctype="text/plain"):14    bnd = "----atlas" + uuid.uuid4().hex15    parts = [f"--{bnd}\r\nContent-Disposition: form-data; name=\"{k}\"\r\n\r\n{v}\r\n".encode() for k, v in fields]16    parts.append(f"--{bnd}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"{filename}\"\r\nContent-Type: {ctype}\r\n\r\n".encode() + content + b"\r\n" + f"--{bnd}--\r\n".encode())17    return b"".join(parts), f"multipart/form-data; boundary={bnd}"181920def _png(w=32, h=32):21    raw = b"".join(b"\x00" + b"\xff\xff\xff" * w for _ in range(h))22    def chunk(t, d):23        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)24    return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")252627def test_file_upload_list_get_content_delete(xai):28    content = b"atlas test file " * 829    body, ct = _multipart([("expires_after", "3600"), ("purpose", "assistants")], "atlas-test.txt", content)30    st, f, _ = xai("POST", "/v1/files", data=body, content_type=ct, note="test_files upload")31    assert st == 200 and f["object"] == "file" and f["bytes"] == len(content) and f["expires_at"] == f["created_at"] + 360032    fid = f["id"]33    try:34        st, lst, _ = xai("GET", "/v1/files?limit=5", note="test_files list")35        assert st == 200 and "data" in lst and "pagination_token" in lst36        st, got, _ = xai("GET", f"/v1/files/{fid}", note="test_files get")37        assert st == 200 and got["id"] == fid and got["purpose"] == ""  # purpose is not stored38        st, raw, _ = xai("GET", f"/v1/files/{fid}/content", content_type=None, note="test_files content")39        assert st == 200 and raw == content40        st, pu, _ = xai("POST", f"/v1/files/{fid}/public-url", {}, note="test_files public-url txt (400 expected)")41        assert st == 400 and "text/plain" in pu["error"]42    finally:43        st, d, _ = xai("DELETE", f"/v1/files/{fid}", note="test_files delete")44        assert st == 200 and d["deleted"] is True45    st, _, _ = xai("GET", f"/v1/files/{fid}", note="test_files get after delete")46    assert st == 404474849def test_public_url_create_idempotent_revoke(xai):50    body, ct = _multipart([("purpose", "assistants")], "atlas-test.png", _png(), "image/png")51    st, f, _ = xai("POST", "/v1/files", data=body, content_type=ct, note="test_files upload png")52    assert st == 20053    fid = f["id"]54    try:55        st, p1, _ = xai("POST", f"/v1/files/{fid}/public-url", {}, note="test_files public-url")56        assert st == 200 and p1["public_url"].startswith("https://files-cdn.x.ai/")57        st, p2, _ = xai("POST", f"/v1/files/{fid}/public-url", {"expires_after": 3600}, note="test_files public-url again")58        assert st == 200 and p2["public_url"] == p1["public_url"] and p2.get("expires_at")59        st, got, _ = xai("GET", f"/v1/files/{fid}", note="test_files get with public url")60        assert got["public_url"] == p1["public_url"]61        st, rv, _ = xai("POST", f"/v1/files/{fid}/public-url/revoke", content_type=None, note="test_files revoke")62        assert st == 200 and rv["revoked"] is True63        st, rv2, _ = xai("POST", f"/v1/files/{fid}/public-url/revoke", content_type=None, note="test_files revoke again")64        assert st == 200 and rv2["revoked"] is False65    finally:66        xai("DELETE", f"/v1/files/{fid}", note="test_files delete png")676869def test_collections_crud_on_api_x_ai(xai):70    st, col, _ = xai("POST", "/v1/collections", {"collection_name": "atlas-pytest", "field_definitions": []}, note="test_files collections create")71    assert st == 200 and col["collection_id"].startswith("collection_")72    cid = col["collection_id"]73    try:74        st, got, _ = xai("GET", f"/v1/collections/{cid}", note="test_files collections get")75        assert st == 200 and got["collection_name"] == "atlas-pytest"76        st, lst, _ = xai("GET", "/v1/collections?limit=5", note="test_files collections list")77        assert st == 200 and any(c["collection_id"] == cid for c in lst["collections"])78        st, s, _ = xai("POST", "/v1/documents/search", {"query": "x", "source": {"collection_ids": [cid]}}, note="test_files documents search empty collection")79        assert st in (200, 404)80    finally:81        st, _, _ = xai("DELETE", f"/v1/collections/{cid}", note="test_files collections delete")82        assert st == 200838485def test_management_api_requires_management_key(xai):86    st, b, _ = xai("GET", "/v1/collections", management=True, note="test_files management-api with inference key")87    assert st == 401 and "management key" in b["message"].lower()888990def test_embeddings_not_available(xai):91    st, b, _ = xai("GET", "/v1/embedding-models", note="test_files embedding models")92    assert st == 200 and b["models"] == []93    st, b, _ = xai("POST", "/v1/embeddings", {"model": "grok-embedding-small", "input": "hi"}, note="test_files embeddings 404")94    assert st == 404959697@pytest.mark.run_batch_tests98def test_batch_lifecycle(xai, models):99    st, b, _ = xai("POST", "/v1/batches", {"name": "atlas-pytest"}, note="test_files batch create")100    assert st == 200 and b["batch_id"].startswith("batch_")101    bid = b["batch_id"]102    st, _, _ = xai("POST", f"/v1/batches/{bid}/requests", {"batch_requests": [{"batch_request_id": "r1", "batch_request": {"chat_get_completion": {"model": models["xai"], "messages": [{"role": "user", "content": "Reply with OK."}], "max_completion_tokens": 32}}}]}, est_cost_usd=0.0004, note="test_files batch add")103    assert st == 200104    for _ in range(18):105        st, b, _ = xai("GET", f"/v1/batches/{bid}", note="test_files batch poll")106        if b["state"]["num_pending"] == 0 and b["state"]["num_requests"] > 0:107            break108        time.sleep(10)109    assert b["state"]["num_success"] == 1110    st, res, _ = xai("GET", f"/v1/batches/{bid}/results", note="test_files batch results")111    assert st == 200 and res["results"][0]["batch_result"]["response"]["chat_get_completion"]["object"] == "chat.completion"112    st, c, _ = xai("POST", f"/v1/batches/{bid}:cancel", content_type=None, note="test_files batch cancel")113    assert st == 200 and c["cancel_time"]114