prototype/search_engine.py - master - nofuti/Pichko-System-Module

prototype/search_engine.py
250 строк8 KB
Георгий Пичко
Папка прототипа
05 май 2026, 06:50
05 май 2026, 06:50f200cdd
"""Индексация .txt в ChromaDB и семантический поиск через sentence-transformers."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

import chromadb
from chromadb.api.models.Collection import Collection
from sentence_transformers import SentenceTransformer

MODEL_NAME = "all-MiniLM-L6-v2"
COLLECTION_NAME = "notes"
CHROMA_DIR = Path(__file__).resolve().parent / "chroma_db"
DATA_DIR = Path(__file__).resolve().parent / "data"
FINGERPRINT_FILE = CHROMA_DIR / "data_fingerprint.sha256"
ENCODE_BATCH = 32


def chunk_text(text: str) -> list[str]:
    """Один фрагмент = один абзац (разделитель — пустая строка, т.е. двойной перевод строки)."""
    text = text.replace("\r\n", "\n")
    if not text.strip():
        return []
    chunks: list[str] = []
    for block in text.split("\n\n"):
        piece = block.strip()
        if piece:
            chunks.append(piece)
    return chunks


def _load_txt_files() -> list[tuple[str, str]]:
    """Возвращает список (относительный_путь, содержимое)."""
    if not DATA_DIR.is_dir():
        return []
    out: list[tuple[str, str]] = []
    for path in sorted(DATA_DIR.glob("*.txt")):
        try:
            content = path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            content = path.read_text(encoding="utf-8", errors="replace")
        rel = path.name
        out.append((rel, content))
    return out


def get_collection() -> Collection:
    CHROMA_DIR.mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=str(CHROMA_DIR))
    return client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )


def get_model() -> SentenceTransformer:
    return SentenceTransformer(MODEL_NAME)


def compute_data_fingerprint() -> str:
    """Хеш списка .txt: имя, время изменения, размер (реагирует на правки и новые/удалённые файлы)."""
    if not DATA_DIR.is_dir():
        return hashlib.sha256(b"NO_DATA_DIR").hexdigest()
    entries: list[list[Any]] = []
    for path in sorted(DATA_DIR.glob("*.txt")):
        try:
            st = path.stat()
        except OSError:
            continue
        entries.append([path.name, st.st_mtime_ns, st.st_size])
    blob = json.dumps(entries, ensure_ascii=True, separators=(",", ":"))
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()


def _read_saved_fingerprint() -> str | None:
    if not FINGERPRINT_FILE.is_file():
        return None
    try:
        text = FINGERPRINT_FILE.read_text(encoding="utf-8").strip()
    except OSError:
        return None
    return text or None


def _write_saved_fingerprint(digest: str) -> None:
    CHROMA_DIR.mkdir(parents=True, exist_ok=True)
    FINGERPRINT_FILE.write_text(digest + "\n", encoding="utf-8")


def _delete_collection() -> None:
    CHROMA_DIR.mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=str(CHROMA_DIR))
    try:
        client.delete_collection(name=COLLECTION_NAME)
    except Exception:
        pass


def reset_index_for_rebuild() -> None:
    """Удаляет отпечаток и коллекцию Chroma — при следующем ensure_index_matches_data будет полная переиндексация."""
    if FINGERPRINT_FILE.is_file():
        try:
            FINGERPRINT_FILE.unlink()
        except OSError:
            pass
    _delete_collection()


def _needs_full_reindex(saved_fp: str | None, current_fp: str, collection_count: int) -> bool:
    if saved_fp != current_fp:
        return True
    has_files = bool(_load_txt_files())
    if has_files and collection_count == 0:
        return True
    if not has_files and collection_count > 0:
        return True
    return False


def _fill_collection_from_disk(collection: Collection, model: SentenceTransformer) -> dict[str, Any]:
    files = _load_txt_files()
    if not files:
        return {"files": 0, "chunks": 0, "detail": "no_txt_files"}

    all_ids: list[str] = []
    all_docs: list[str] = []
    all_meta: list[dict[str, Any]] = []

    for rel_path, content in files:
        parts = chunk_text(content)
        for i, doc in enumerate(parts):
            all_ids.append(f"{rel_path}::{i}")
            all_docs.append(doc)
            all_meta.append({"source": rel_path, "chunk_index": i})

    if not all_docs:
        return {"files": len(files), "chunks": 0, "detail": "empty_after_chunking"}

    embeddings = model.encode(
        all_docs,
        batch_size=ENCODE_BATCH,
        show_progress_bar=len(all_docs) > 8,
        convert_to_numpy=True,
    )

    collection.add(
        ids=all_ids,
        documents=all_docs,
        metadatas=all_meta,
        embeddings=embeddings.tolist(),
    )

    return {"files": len(files), "chunks": len(all_docs), "detail": "ok"}


def ensure_index_matches_data(model: SentenceTransformer | None = None) -> dict[str, Any]:
    """
    Синхронизирует Chroma с папкой data: при изменении файлов или несоответствии БД — полная переиндексация.
    """
    if model is None:
        model = get_model()

    current_fp = compute_data_fingerprint()
    saved_fp = _read_saved_fingerprint()
    collection = get_collection()
    count = collection.count()

    if not _needs_full_reindex(saved_fp, current_fp, count):
        return {"indexed": False, "reason": "up_to_date", "previous_chunks": count}

    _delete_collection()
    collection = get_collection()
    result = _fill_collection_from_disk(collection, model)
    _write_saved_fingerprint(current_fp)

    detail = result.get("detail", "")
    if detail == "no_txt_files":
        return {
            "indexed": True,
            "reason": "reindexed",
            "files": 0,
            "chunks": 0,
            "previous_chunks": count,
            "detail": "no_txt_files",
        }
    if detail == "empty_after_chunking":
        return {
            "indexed": True,
            "reason": "reindexed",
            "files": result["files"],
            "chunks": 0,
            "previous_chunks": count,
            "detail": "empty_after_chunking",
        }

    return {
        "indexed": True,
        "reason": "reindexed",
        "files": result["files"],
        "chunks": result["chunks"],
        "previous_chunks": count,
    }


def index_if_empty(model: SentenceTransformer | None = None) -> dict[str, Any]:
    """
    Устаревшее имя: то же, что ensure_index_matches_data (пересборка при изменении data/).
    """
    return ensure_index_matches_data(model=model)


def search_notes(
    query: str,
    top_k: int = 5,
    model: SentenceTransformer | None = None,
) -> dict[str, Any]:
    collection = get_collection()
    if collection.count() == 0:
        return {"error": "База пуста. Положите .txt в папку data и перезапустите (или удалите chroma_db для переиндексации)."}

    if model is None:
        model = get_model()

    q_emb = model.encode([query], convert_to_numpy=True)[0]

    res = collection.query(
        query_embeddings=[q_emb.tolist()],
        n_results=min(top_k, collection.count()),
        include=["documents", "metadatas", "distances"],
    )

    hits: list[dict[str, Any]] = []
    docs = (res.get("documents") or [[]])[0]
    metas = (res.get("metadatas") or [[]])[0]
    dists = (res.get("distances") or [[]])[0]

    for doc, meta, dist in zip(docs, metas, dists):
        hits.append(
            {
                "text": doc,
                "source": meta.get("source", "?") if meta else "?",
                "chunk_index": meta.get("chunk_index", 0) if meta else 0,
                "distance": float(dist) if dist is not None else None,
            }
        )

    return {"hits": hits}