spb/doc-api

1#!/usr/bin/env python32"""Collections end-to-end on api.x.ai with the inference key (create → upload → add document → poll → search → file_search tool → cleanup)."""3from __future__ import annotations4import json, sys, time5from pathlib import Path6ROOT = Path(__file__).resolve().parents[2]7sys.path.insert(0, str(ROOT))8from scripts import live  # noqa: E4029src = open(ROOT / "tmp/xai-inference/probe.py").read()10exec(src.split("# ---------------- (a) chat minimal")[0])  # helpers: X, save, sse, OUT, M11exec(src[src.index('def multipart('):src.index('FILE_STATE = {}')])12def J(b):13    return b if isinstance(b, dict) else {"_raw": (b.decode("utf-8","replace") if isinstance(b, bytes) else str(b))}14_X = X15def X(*a, **k):16    st, b, h = _X(*a, **k)17    return st, J(b), h18def txt(b):19    return next((c["text"] for o in b.get("output", []) if o.get("type") == "message" for c in o["content"] if c["type"] == "output_text"), None)20out = {}21st, c, _ = X("POST", "/v1/collections", {"collection_name": "atlas-probe-2", "field_definitions": [{"key": "type", "required": False, "inject_into_chunk": False, "unique": False, "description": "doc type"}]}, "probe3 collections create", 0)22out["create"] = {"status": st, "collection_id": c.get("collection_id")}23cid = c["collection_id"]24content = b"ATLAS collections probe. The secret fruit is MANGO. " * 425body, ct = multipart([("purpose", "assistants")], "file", "atlas-coll.txt", content)26st, f, _ = X("POST", "/v1/files", None, "probe3 upload for collection", 0, data=body, content_type=ct); fid = f["id"]27st, b, _ = X("POST", f"/v1/collections/{cid}/documents/{fid}", {"collection_id": cid, "file_id": fid, "fields": {"type": "note"}}, "probe4 add document (body collection_id+file_id)", 0)28save("collections_add_document_4", {"status": st, "body": b}); out["add_document_body_file_id"] = {"status": st, "body": b}29if st != 200:30    body2, ct2 = multipart([("name", "atlas-coll.txt"), ("content_type", "text/plain"), ("fields", json.dumps({"type": "note"}))], "data", "atlas-coll.txt", content)31    st, b, _ = X("POST", f"/v1/collections/{cid}/documents", None, "probe3 upload document multipart", 0, data=body2, content_type=ct2)32    save("collections_upload_document", {"status": st, "body": b}); out["upload_document_multipart"] = {"status": st, "body": b}33    if st == 200 and isinstance(b, dict):34        fid = (b.get("file_metadata") or {}).get("file_id", fid)35status = None; polls = 036for i in range(25):37    st, d, _ = X("GET", f"/v1/collections/{cid}/documents/{fid}", None, "probe3 get document", 0); polls += 138    status = d.get("status") if st == 200 else d39    if status in ("DOCUMENT_STATUS_PROCESSED", "DOCUMENT_STATUS_FAILED") or st != 200: break40    time.sleep(3)41save("collections_get_document_4", {"status": st, "body": d}); out["get_document"] = {"status": st, "doc_status": status, "polls": polls, "body": d}42st, l, _ = X("GET", f"/v1/collections/{cid}/documents", None, "probe3 list documents", 0)43save("collections_list_documents_4", {"status": st, "body": l}); out["list_documents"] = {"status": st, "body": l}44st, bg, _ = X("GET", f"/v1/collections/{cid}/documents:batchGet?file_ids[]={fid}", None, "probe3 documents batchGet", 0)45out["batchGet"] = {"status": st, "body": bg}46st, s, _ = X("POST", "/v1/documents/search", {"query": "secret fruit", "source": {"collection_ids": [cid]}, "limit": 2, "retrieval_mode": {"type": "hybrid"}}, "probe3 documents/search", 0.001)47save("documents_search_4", {"status": st, "body": s}); out["documents_search"] = {"status": st, "body": s}48st, r, _ = X("POST", "/v1/responses", {"model": M, "input": "What is the secret fruit in my documents? One word.", "tools": [{"type": "file_search", "vector_store_ids": [cid], "max_num_results": 3}], "include": ["file_search_call.results"], "max_turns": 2, "store": False, "max_output_tokens": 4000}, "probe3 responses file_search tool", 0.01, timeout=300)49save("responses_tool_file_search_4", {"status": st, "body": r})50outs = r.get("output", []) if st == 200 else []51out["file_search_tool"] = {"status": st, "output_types": [o.get("type") for o in outs], "tool_items": [o for o in outs if o.get("type") not in ("message", "reasoning")], "text": txt(r) if st == 200 else r, "annotations": next((c.get("annotations") for o in outs if o.get("type") == "message" for c in o["content"]), None), "usage": (r.get("usage") or {}) if st == 200 else None}52st, cl, _ = X("GET", "/v1/collections?limit=5", None, "probe3 collections list", 0); out["list_collections"] = {"status": st, "n": len(cl.get("collections", [])) if st == 200 else cl}53st, up, _ = X("PUT", f"/v1/collections/{cid}", {"collection_id": cid, "collection_name": "atlas-probe-2b"}, "probe4 collections update", 0); out["update"] = {"status": st, "name": up.get("collection_name") if st == 200 else up}54st, dd, _ = X("DELETE", f"/v1/collections/{cid}/documents/{fid}", None, "probe3 remove document", 0); out["remove_document"] = {"status": st, "body": dd}55st, fg, _ = X("GET", f"/v1/files/{fid}", None, "probe3 file after remove from collection", 0); out["file_after_remove"] = {"status": st}56st, dc, _ = X("DELETE", f"/v1/collections/{cid}", None, "probe3 delete collection", 0); out["delete_collection"] = {"status": st, "body": dc}57st, fd, _ = X("DELETE", f"/v1/files/{fid}", None, "probe3 delete file", 0); out["delete_file"] = {"status": st, "body": fd}58(OUT / "summary4.json").write_text(live.mask(json.dumps(out, indent=1, default=str)))59o = json.loads(json.dumps(out, default=str))60for it in o["file_search_tool"].get("tool_items", []) if isinstance(o["file_search_tool"], dict) else []:61    if "results" in it: it["results"] = [str(x)[:200] for x in it["results"]]62print(json.dumps(o, default=str)[:7000])63