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

prototype/main.py
87 строк3 KB
Георгий Пичко
В прототип добавлен чат
13 май 2026, 12:38
13 май 2026, 12:380416b5d
"""CLI: индексация при пустой БД, затем цикл вопросов."""
""" cd C:\Users\user\Desktop\AI_Chat\prototype
streamlit run streamlit_app.py"""
from __future__ import annotations

import argparse
import sys

from search_engine import ensure_index_matches_data, get_model, reset_index_for_rebuild, search_notes


def _configure_stdio_utf8() -> None:
    if sys.platform == "win32":
        for stream in (sys.stdout, sys.stderr, sys.stdin):
            if hasattr(stream, "reconfigure"):
                try:
                    stream.reconfigure(encoding="utf-8")
                except (OSError, ValueError):
                    pass


def main() -> None:
    _configure_stdio_utf8()
    parser = argparse.ArgumentParser(description="Семантический поиск по заметкам (.txt в data/)")
    parser.add_argument("-k", "--top-k", type=int, default=5, help="Сколько фрагментов показать")
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Принудительно пересобрать индекс при каждом запуске (игнорировать отпечаток файлов)",
    )
    args = parser.parse_args()

    print("Загрузка модели эмбеддингов (при первом запуске может скачаться)...")
    model = get_model()

    if args.rebuild:
        reset_index_for_rebuild()

    info = ensure_index_matches_data(model=model)
    if info.get("indexed"):
        detail = info.get("detail")
        if detail == "no_txt_files":
            print("Папка data: нет .txt — база очищена. Добавьте заметки в: data/")
        elif detail == "empty_after_chunking":
            print("Файлы есть, но после разбиения не осталось непустого текста.")
        else:
            prev = info.get("previous_chunks", 0)
            print(
                f"Индекс обновлён (файлы в data изменились или база устарела): "
                f"файлов {info['files']}, фрагментов {info['chunks']}"
                + (f"; ранее было фрагментов {prev}." if prev else ".")
            )
    elif info.get("reason") == "up_to_date":
        print(
            f"Данные в data/ без изменений — используется текущий индекс "
            f"({info.get('previous_chunks', 0)} фрагментов)."
        )

    print()
    print("Введите вопрос (пустая строка или Ctrl+Z+Enter — выход).")
    while True:
        try:
            q = input("\nВопрос> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not q:
            break
        out = search_notes(q, top_k=args.top_k, model=model)
        if "error" in out:
            print(out["error"])
            continue
        for i, h in enumerate(out.get("hits") or [], 1):
            dist = h.get("distance")
            dist_s = f"{dist:.4f}" if dist is not None else "n/a"
            src = h.get("source", "?")
            print(f"\n--- [{i}] {src} (расстояние {dist_s}) ---")
            text = (h.get("text") or "").strip()
            preview = text if len(text) <= 1200 else text[:1200] + "\n…"
            print(preview)


if __name__ == "__main__":
    try:
        main()
    except BrokenPipeError:
        sys.exit(0)