00_create_kb.py 1.86 KB
from __future__ import annotations

import sys
from typing import Any

import _bootstrap  # noqa: F401

from weknora_eval.api import bootstrap_client_from_config
from weknora_eval.config import load_config, require_config
from weknora_eval.envfile import set_env_value
from weknora_eval.loaders import setup_logging, write_json


def main() -> int:
    setup_logging()
    config = load_config()
    client = bootstrap_client_from_config(config)
    weknora = config["weknora"]

    existing_id = str(weknora.get("knowledge_base_id") or "")
    name = str(require_config(config, "weknora.knowledge_base_name"))
    if existing_id and existing_id != "replace-me":
        record = {"id": existing_id, "name": name, "source": "env"}
        write_json("data/exported/knowledge_base.json", record)
        print(f"WEKNORA_KB_ID already set: {existing_id}")
        return 0

    created = client.create_knowledge_base(name=name)
    knowledge_base_id = _extract_knowledge_base_id(created)
    if not knowledge_base_id:
        print(f"Created knowledge base but could not extract id from response: {created}")
        return 1

    set_env_value(".env", "WEKNORA_KB_ID", knowledge_base_id)
    write_json("data/exported/knowledge_base.json", {**created, "source": "create"})
    print(f"WEKNORA_KB_ID={knowledge_base_id}")
    print("Wrote ID to .env and data/exported/knowledge_base.json")
    return 0


def _extract_knowledge_base_id(payload: dict[str, Any]) -> str | None:
    candidates = [payload]
    for key in ("data", "knowledge_base"):
        nested = payload.get(key)
        if isinstance(nested, dict):
            candidates.append(nested)

    for row in candidates:
        for key in ("id", "knowledge_base_id", "kb_id", "uuid"):
            value = row.get(key)
            if value:
                return str(value)
    return None


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