SPB Git forge

spb/doc-api

Public
2commits 1branches 0releases
15.7 MBsize
maindefault branch
13 days agolast push
Python 88.3% TypeScript 7.6% Shell 4.1%
5.8 KB · 88 lines python
Raw Blame History
1"""Vector store smoke tests: create, attach w/ static chunking + attributes, poll, search with filters, batches, cleanup. ~$0."""2from __future__ import annotations3import time, uuid4import pytest56TAG = "atlas-platform-agent-test"789def _upload(openai, name, text):10    b = f"----atlas{uuid.uuid4().hex}"11    body = (f"--{b}\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nassistants\r\n--{b}\r\n"12            f"Content-Disposition: form-data; name=\"file\"; filename=\"{name}\"\r\nContent-Type: text/plain\r\n\r\n{text}\r\n--{b}--\r\n").encode()13    st, f, _ = openai("POST", "/v1/files", data=body, content_type=f"multipart/form-data; boundary={b}", note="test_vector_stores upload")14    assert st == 200, f15    return f["id"]161718@pytest.fixture(scope="module")19def store(openai):20    st, vs, _ = openai("POST", "/v1/vector_stores", {"name": TAG, "metadata": {"owner": TAG}, "expires_after": {"anchor": "last_active_at", "days": 1}}, note="test_vector_stores create")21    assert st == 200, vs22    files = []23    yield vs, files24    for fid in files:25        openai("DELETE", f"/v1/vector_stores/{vs['id']}/files/{fid}", note="test_vector_stores cleanup")26        openai("DELETE", f"/v1/files/{fid}", note="test_vector_stores cleanup")27    st, d, _ = openai("DELETE", f"/v1/vector_stores/{vs['id']}", note="test_vector_stores cleanup")28    assert st == 200 and d["deleted"] and d["object"] == "vector_store.deleted"293031def test_store_shape(store):32    vs, _ = store33    assert vs["object"] == "vector_store" and vs["id"].startswith("vs_") and vs["status"] == "completed"34    assert vs["file_counts"]["total"] == 0 and vs["usage_bytes"] == 035    assert vs["expires_after"] == {"anchor": "last_active_at", "days": 1} and vs["expires_at"] == vs["created_at"] + 86400363738def test_attach_poll_search_update(openai, store):39    vs, files = store40    fid = _upload(openai, f"{TAG}.txt", "Atlas policy: each passenger may carry at most two woodchucks. Returns within 30 days.")41    files.append(fid)42    st, vf, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/files", {"file_id": fid, "attributes": {"region": "us", "year": 2026, "draft": False},43                       "chunking_strategy": {"type": "static", "static": {"max_chunk_size_tokens": 100, "chunk_overlap_tokens": 20}}}, note="test_vector_stores attach")44    assert st == 200 and vf["object"] == "vector_store.file" and vf["status"] == "in_progress"45    assert vf["chunking_strategy"] == {"type": "static", "static": {"max_chunk_size_tokens": 100, "chunk_overlap_tokens": 20}}46    for _ in range(45):47        st, vf, _ = openai("GET", f"/v1/vector_stores/{vs['id']}/files/{fid}", note="test_vector_stores poll")48        if vf["status"] != "in_progress":49            break50        time.sleep(2)51    assert vf["status"] == "completed", vf52    assert vf["usage_bytes"] > 053    st, page, _ = openai("GET", f"/v1/vector_stores/{vs['id']}/files?filter=completed", note="test_vector_stores list")54    assert st == 200 and [x["id"] for x in page["data"]] == [fid]55    st, content, _ = openai("GET", f"/v1/vector_stores/{vs['id']}/files/{fid}/content", note="test_vector_stores content")56    assert st == 200 and content["object"] == "vector_store.file_content.page" and content["data"][0]["type"] == "text"57    st, res, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/search", {"query": "how many woodchucks per passenger", "max_num_results": 3, "rewrite_query": True,58                        "filters": {"type": "and", "filters": [{"type": "eq", "key": "region", "value": "us"}, {"type": "gte", "key": "year", "value": 2025}]}}, note="test_vector_stores search")59    assert st == 200 and res["object"] == "vector_store.search_results.page" and isinstance(res["search_query"], list)60    assert len(res["data"]) == 1 and res["data"][0]["file_id"] == fid and 0 < res["data"][0]["score"] <= 161    assert res["data"][0]["attributes"]["region"] == "us" and res["data"][0]["content"][0]["type"] == "text"62    st, none, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/search", {"query": "woodchucks", "filters": {"type": "eq", "key": "region", "value": "eu"}}, note="test_vector_stores search no match")63    assert st == 200 and none["data"] == []64    st, upd, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/files/{fid}", {"attributes": {"region": "ca"}}, note="test_vector_stores update attrs")65    assert st == 200 and upd["attributes"] == {"region": "ca"}  # replaced, not merged666768def test_file_batch(openai, store):69    vs, files = store70    fid = _upload(openai, f"{TAG}-2.txt", "Refunds are processed within 5 business days. " * 3)71    files.append(fid)72    st, fb, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/file_batches", {"files": [{"file_id": fid, "attributes": {"topic": "refunds"}, "chunking_strategy": {"type": "auto"}}]}, note="test_vector_stores batch")73    assert st == 200 and fb["object"] in ("vector_store.file_batch", "vector_store.files_batch") and fb["file_counts"]["total"] == 174    for _ in range(45):75        st, fb, _ = openai("GET", f"/v1/vector_stores/{vs['id']}/file_batches/{fb['id']}", note="test_vector_stores batch poll")76        if fb["status"] != "in_progress":77            break78        time.sleep(2)79    assert fb["status"] == "completed" and fb["file_counts"]["completed"] == 180    st, listed, _ = openai("GET", f"/v1/vector_stores/{vs['id']}/file_batches/{fb['id']}/files", note="test_vector_stores batch files")81    assert st == 200 and listed["data"][0]["id"] == fid and listed["data"][0]["attributes"] == {"topic": "refunds"}828384def test_invalid_chunking_rejected(openai, store):85    vs, _ = store86    st, err, _ = openai("POST", f"/v1/vector_stores/{vs['id']}/files", {"file_id": "file-atlasbogus", "chunking_strategy": {"type": "static", "static": {"max_chunk_size_tokens": 50, "chunk_overlap_tokens": 10}}}, note="test_vector_stores invalid chunking (expect 4xx)")87    assert 400 <= st < 500, err88