02_wait_ingestion.py 988 Bytes
from __future__ import annotations

import sys

import _bootstrap  # noqa: F401

from weknora_eval.api import client_from_config
from weknora_eval.config import load_config
from weknora_eval.loaders import read_jsonl, setup_logging, write_jsonl


def main() -> int:
    setup_logging()
    config = load_config()
    client = client_from_config(config)
    uploads = read_jsonl("data/exported/knowledge_uploads.jsonl", missing_ok=True)
    knowledge_ids = {row["knowledge_id"] for row in uploads if row.get("knowledge_id")} or None
    result = client.wait_ingestion_completed(knowledge_ids=knowledge_ids)
    knowledge = client.list_knowledge()
    write_jsonl("data/exported/knowledge.jsonl", knowledge)

    print(
        "Ingestion status: "
        f"completed={len(result['completed'])} failed={len(result['failed'])} "
        f"pending={len(result['pending'])}"
    )
    return 1 if result["failed"] or result["pending"] else 0


if __name__ == "__main__":
    sys.exit(main())