tech/7_сохранение_qdrant/save_to_qdrant.py

#!/usr/bin/env python3
"""
Шаг 7: сохранение вектора эмбеддинга и payload в коллекцию Qdrant chapter_analyses.

Вход:
  - merged JSON (шаг 5): теги, анализ и метаданные главы/книги (book_id, chapter_id,
    chapter_number, chapter_title, book_title, author — из входного JSON этапа 1,
    подмешанные при мерже через --input-chapter).
  - вектор (файл или stdin): результат шага 6 (embed_cli.py).

Выход: одна точка в коллекции chapter_analyses (upsert по chapter_id).
"""

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any


def env(name: str, default: str) -> str:
    """Читает переменную окружения или default."""
    return os.environ.get(name, default).strip()


def tags_from_merged(merged: dict[str, Any]) -> list[dict[str, Any]]:
    """Преобразует теги из merged JSON в плоский список для payload (tag, category, confidence)."""
    result: list[dict[str, Any]] = []
    tags_block = merged.get("tags") or {}
    for category, items in tags_block.items():
        for item in items if isinstance(items, list) else []:
            if isinstance(item, dict) and "tag" in item:
                result.append({
                    "tag": item["tag"],
                    "category": category,
                    "confidence": item.get("confidence"),
                })
    return result


def payload_from_merged(merged: dict[str, Any], validation_score: float | None = None) -> dict[str, Any]:
    """
    Собирает payload для Qdrant из merged JSON (шаг 5).

    Ожидает в merged поля из входного JSON этапа 1: book_id, chapter_id,
    chapter_number, chapter_title, book_title; опционально author.
    """
    required = ("book_id", "chapter_id", "chapter_number", "chapter_title")
    missing = [k for k in required if not merged.get(k)]
    if missing:
        raise ValueError(
            f"В merged JSON отсутствуют поля (должны быть из входного JSON этапа 1, мерж с --input-chapter): {missing}"
        )
    payload: dict[str, Any] = {
        "bookId": merged["book_id"],
        "chapterId": merged["chapter_id"],
        "chapterNumber": int(merged["chapter_number"]),
        "chapterTitle": merged["chapter_title"],
        "tags": tags_from_merged(merged),
    }
    if validation_score is not None:
        payload["validationScore"] = validation_score
    if merged.get("book_title") is not None:
        payload["title"] = merged["book_title"]
    if merged.get("author") is not None:
        payload["author"] = merged["author"]
    return payload


def load_vector(path: Path | None) -> list[float]:
    """Загружает вектор из файла или stdin (JSON-массив)."""
    if path is None or str(path) == "-":
        data = json.load(sys.stdin)
    else:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
    if not isinstance(data, list):
        raise ValueError("Вектор должен быть JSON-массивом чисел")
    return [float(x) for x in data]


def upsert_point(
    base_url: str,
    collection: str,
    point_id: str,
    vector: list[float],
    payload: dict[str, Any],
) -> dict[str, Any]:
    """Добавляет или обновляет одну точку в коллекции Qdrant."""
    url = f"{base_url.rstrip('/')}/collections/{collection}/points?wait=true"
    body = {
        "points": [
            {
                "id": point_id,
                "vector": vector,
                "payload": payload,
            }
        ]
    }
    req_body = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=req_body,
        headers={"Content-Type": "application/json"},
        method="PUT",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main() -> int:
    """Точка входа CLI."""
    parser = argparse.ArgumentParser(
        description="Шаг 7: сохранить вектор и payload в Qdrant (коллекция chapter_analyses).",
    )
    parser.add_argument(
        "merged_json",
        type=Path,
        help="Путь к merged JSON (шаг 5) — для тегов и др.",
    )
    parser.add_argument(
        "vector",
        nargs="?",
        default=None,
        help="Путь к файлу с вектором (JSON-массив) или '-' для stdin.",
    )
    parser.add_argument(
        "--validation-score",
        type=float,
        default=None,
        help="Оценка валидации (опционально).",
    )
    parser.add_argument(
        "--qdrant-url",
        default=env("QDRANT_URL", "http://localhost:6333"),
        help="URL Qdrant (по умолчанию из QDRANT_URL или localhost:6333).",
    )
    parser.add_argument(
        "--collection",
        default=env("QDRANT_COLLECTION_CHAPTER_ANALYSES", "chapter_analyses"),
        help="Имя коллекции.",
    )
    args = parser.parse_args()

    if not args.merged_json.exists():
        print(f"Ошибка: файл не найден: {args.merged_json}", file=sys.stderr)
        return 1

    try:
        with open(args.merged_json, encoding="utf-8") as f:
            merged = json.load(f)
    except json.JSONDecodeError as e:
        print(f"Ошибка разбора merged JSON: {e}", file=sys.stderr)
        return 1

    try:
        vector = load_vector(args.vector)
    except (ValueError, FileNotFoundError) as e:
        print(f"Ошибка загрузки вектора: {e}", file=sys.stderr)
        return 1

    try:
        payload = payload_from_merged(merged, validation_score=args.validation_score)
    except ValueError as e:
        print(f"Ошибка: {e}", file=sys.stderr)
        return 1

    chapter_id = merged["chapter_id"]

    try:
        result = upsert_point(
            args.qdrant_url,
            args.collection,
            point_id=chapter_id,
            vector=vector,
            payload=payload,
        )
    except urllib.error.HTTPError as e:
        print(f"Ошибка HTTP {e.code}: {e.reason}", file=sys.stderr)
        if e.fp:
            try:
                print(e.fp.read().decode("utf-8")[:500], file=sys.stderr)
            except Exception:
                pass
        return 1
    except urllib.error.URLError as e:
        print(f"Ошибка запроса к Qdrant: {e.reason}", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"Ошибка: {e}", file=sys.stderr)
        return 1

    if result.get("status") == "ok":
        print(f"Точка {chapter_id} сохранена в коллекцию {args.collection}.", flush=True)
        return 0
    print(f"Неожиданный ответ Qdrant: {result}", file=sys.stderr)
    return 1


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