scripts/evaluator_document_store_bootstrap.py - master - Vibek/Agent_GG

scripts/evaluator_document_store_bootstrap.py
142 строки4 KB
Zeron
Добавлен seed для документов и проверки полученных чанков. Обновил README и runbook
07 май 2026, 21:21
07 май 2026, 21:21836b9ba
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path


WORKSPACE_ROOT = Path(__file__).resolve().parents[1]
if str(WORKSPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(WORKSPACE_ROOT))


from agent.evaluator.vectorization.ingestion import (  # noqa: E402
    DEFAULT_EMBEDDING_BATCH_SIZE,
    build_document_store_engine,
    ingest_r_document_embeddings,
)
from agent.evaluator.vectorization.readiness import evaluate_document_store_readiness  # noqa: E402
from agent.llm.gigachat_client import gigachat_embeddings_client  # noqa: E402


logger = logging.getLogger(__name__)


def _build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Check and optionally bootstrap evaluator document store.")
    parser.add_argument(
        "--database-url",
        type=str,
        default=None,
        help="Optional PostgreSQL URL override. Defaults to DATABASE_URL from settings.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_EMBEDDING_BATCH_SIZE,
        help="How many chunk embeddings to request from GigaChat in one batch if ingestion is needed.",
    )
    parser.add_argument(
        "--check-only",
        action="store_true",
        help="Only check readiness. Do not run ingestion even if credentials are configured.",
    )
    return parser


def _print_json(payload: dict[str, object]) -> None:
    print(json.dumps(payload, ensure_ascii=False, indent=2))


def main() -> int:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
    args = _build_arg_parser().parse_args()

    try:
        engine = build_document_store_engine(args.database_url)
        readiness = evaluate_document_store_readiness(engine)
    except Exception as exc:
        logger.exception("[evaluator-bootstrap] readiness check failed: %s", exc)
        _print_json(
            {
                "status": "failed",
                "action": "readiness_check_failed",
                "error": str(exc),
            }
        )
        return 1

    if readiness.is_ready:
        _print_json(
            {
                "status": "ready",
                "action": "skip_ingestion",
                "readiness": readiness.to_dict(),
            }
        )
        return 0

    if args.check_only:
        _print_json(
            {
                "status": "not_ready",
                "action": "check_only",
                "readiness": readiness.to_dict(),
            }
        )
        return 2

    if not gigachat_embeddings_client.is_configured:
        _print_json(
            {
                "status": "not_ready",
                "action": "missing_seed_or_credentials",
                "message": (
                    "Evaluator document store is not ready, seed is missing or incomplete, "
                    "and GigaChat credentials are not configured. Restore seed or configure credentials."
                ),
                "readiness": readiness.to_dict(),
            }
        )
        return 2

    logger.info("[evaluator-bootstrap] document store is not ready (%s), running ingestion", readiness.reason)
    try:
        verification = ingest_r_document_embeddings(
            batch_size=args.batch_size,
            database_url=args.database_url,
        )
    except Exception as exc:
        logger.exception("[evaluator-bootstrap] ingestion failed: %s", exc)
        _print_json(
            {
                "status": "failed",
                "action": "ingestion_failed",
                "error": str(exc),
                "readiness_before_ingestion": readiness.to_dict(),
            }
        )
        return 1

    _print_json(
        {
            "status": "ready",
            "action": "ingestion_completed",
            "verification": {
                "expected_chunk_count": verification.expected_chunk_count,
                "stored_chunk_count": verification.stored_chunk_count,
                "embedded_chunk_count": verification.embedded_chunk_count,
                "expected_doc_ids": list(verification.expected_doc_ids),
                "loaded_doc_ids": list(verification.loaded_doc_ids),
                "chunks_by_doc_id": verification.chunks_by_doc_id,
                "is_ready": verification.is_ready,
            },
        }
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())