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

prototype/streamlit_app.py
117 строк5 KB
Георгий Пичко
В прототип добавлен чат
13 май 2026, 12:38
13 май 2026, 12:380416b5d
"""Веб-интерфейс: индексация, поиск по заметкам и чат Ollama."""

from __future__ import annotations

import os
import sys
from pathlib import Path

import requests
import streamlit as st

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

from chat.chatbot import send_chat, ollama_base_url
from search_engine import ensure_index_matches_data, get_model, search_notes


@st.cache_resource
def load_model():
    return get_model()


def main() -> None:
    st.set_page_config(page_title="Поиск по заметкам и чат", layout="wide")
    st.title("Семантический поиск по заметкам")
    st.caption(
        "Файлы `.txt` из папки `data/`, модель `all-MiniLM-L6-v2`, ChromaDB. "
        "Чат — через Ollama (`OLLAMA_HOST`, по умолчанию localhost:11434)."
    )

    model = load_model()

    with st.spinner("Проверка отпечатка data/ и при необходимости переиндексация..."):
        info = ensure_index_matches_data(model=model)

    if info.get("indexed"):
        detail = info.get("detail")
        if detail == "no_txt_files":
            st.warning("В папке `data/` нет `.txt` файлов (индекс очищен).")
        elif detail == "empty_after_chunking":
            st.warning("Текст в файлах пустой после обработки.")
        else:
            st.success(
                f"Индекс обновлён: файлов {info['files']}, фрагментов {info['chunks']}."
            )
    elif info.get("reason") == "up_to_date":
        st.info(
            f"Файлы в `data/` не менялись — в индексе {info.get('previous_chunks', 0)} фрагментов."
        )

    tab_search, tab_chat = st.tabs(["Поиск по заметкам", "Чат Ollama"])

    with tab_search:
        top_k = st.sidebar.slider("Сколько результатов (поиск)", 1, 15, 5)
        query = st.text_input(
            "Ваш вопрос",
            placeholder="О чём в заметках говорится про ...",
            key="search_query",
        )

        if st.button("Найти", type="primary", key="search_btn") and query.strip():
            out = search_notes(query.strip(), top_k=top_k, model=model)
            if "error" in out:
                st.error(out["error"])
            else:
                hits = out.get("hits") or []
                if not hits:
                    st.info("Ничего не найдено.")
                else:
                    for i, h in enumerate(hits, 1):
                        dist = h.get("distance")
                        dist_s = f"{dist:.4f}" if dist is not None else "n/a"
                        with st.expander(f"#{i} — {h.get('source', '?')} (расстояние {dist_s})"):
                            st.code((h.get("text") or "").strip() or "(пусто)", language=None)

    with tab_chat:
        st.subheader("Чат с моделью через Ollama")
        st.caption(f"Сервер: `{ollama_base_url()}` — переменная окружения `OLLAMA_HOST`.")

        if "chat_messages" not in st.session_state:
            st.session_state.chat_messages = []

        default_model = os.environ.get("OLLAMA_MODEL", "llama3.2:1b")
        chat_model = st.text_input("Модель Ollama", value=default_model, key="ollama_model_input")

        if st.button("Очистить историю"):
            st.session_state.chat_messages = []
            st.rerun()

        for msg in st.session_state.chat_messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

        if prompt := st.chat_input("Сообщение для модели..."):
            model_name = chat_model.strip()
            if not model_name:
                st.warning("Укажите имя модели Ollama.")
                st.stop()
            st.session_state.chat_messages.append({"role": "user", "content": prompt})
            try:
                with st.spinner("Ответ..."):
                    answer = send_chat(st.session_state.chat_messages, model_name)
                st.session_state.chat_messages.append({"role": "assistant", "content": answer})
            except requests.HTTPError as e:
                st.session_state.chat_messages.pop()
                body = (e.response.text or "")[:800]
                st.error(f"HTTP {e.response.status_code}: {body}")
            except requests.RequestException as e:
                st.session_state.chat_messages.pop()
                st.error(f"Ошибка запроса: {e}")
            st.rerun()


if __name__ == "__main__":
    main()