scripts/ingest_document.py - master - medmikhr/ctk_dm_bot_tg

scripts/ingest_document.py
80 строк3 KB
DAnaKhristenko
UI с ошибкой на nginx
08 ноя 2025, 01:01
08 ноя 2025, 01:01d926466
"""Utility script to ingest a single document into the vector store."""

from __future__ import annotations

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

ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

logger = logging.getLogger(__name__)


async def _ingest_document(path: Path, collection: str) -> dict[str, object]:
    from db import AsyncSessionFactory
    from services.document_ingestion import DocumentIngestionPipeline

    async with AsyncSessionFactory() as session:
        pipeline = DocumentIngestionPipeline(session=session, collection_name=collection)
        result = await pipeline.ingest_path(path)
        return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Ingest a document into the vector store")
    parser.add_argument("path", help="Absolute path to the document file")
    parser.add_argument(
        "--collection",
        default="ctkdocs",
        help="Target collection name (default: ctkdocs)",
    )
    parser.add_argument("--log-level", default="INFO", help="Logging level (default: INFO)")
    parser.add_argument(
        "--db-host",
        help="Override POSTGRES_HOST (useful when running outside containers)",
    )
    parser.add_argument("--db-port", help="Override POSTGRES_PORT")
    parser.add_argument("--db-name", help="Override POSTGRES_DB")
    parser.add_argument("--db-user", help="Override POSTGRES_USER")
    parser.add_argument("--db-password", help="Override POSTGRES_PASSWORD")

    args = parser.parse_args()

    logging.basicConfig(level=getattr(logging, args.log_level.upper(), logging.INFO))

    overrides = {
        "POSTGRES_HOST": args.db_host,
        "POSTGRES_PORT": args.db_port,
        "POSTGRES_DB": args.db_name,
        "POSTGRES_USER": args.db_user,
        "POSTGRES_PASSWORD": args.db_password,
    }
    for key, value in overrides.items():
        if value:
            os.environ[key] = value

    target_path = Path(args.path).expanduser().resolve()
    if not target_path.exists():
        logger.error("Файл %s не найден", target_path)
        return 1

    try:
        result = asyncio.run(_ingest_document(target_path, args.collection))
    except Exception as exc:  # pragma: no cover - defensive execution path
        logger.exception("Не удалось импортировать документ %s", target_path)
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1

    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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