app.py

import streamlit as st
from datasets import load_dataset
import numpy as np
from sentence_transformers import SentenceTransformer
import time
from datetime import datetime
import json
import os
import pandas as pd

# ===================================================================
# 1. НАСТРОЙКИ И ЗАГРУЗКА ДАННЫХ
# ===================================================================

MODEL_NAME = "BAAI/bge-large-en-v1.5"
DATASET_NAME = "wikimedia/wikipedia"
LANGUAGE = "20231101.en"
ARTICLE_LIMIT = 1000

ADMIN_USER = "admin"
ADMIN_PASS = "hfpassword21"

LOG_FILE = "query_logs.json"

# Настройка страницы
st.set_page_config(
    page_title="Wikipedia Assistant",
    page_icon="🌖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Загружаем логи
if os.path.exists(LOG_FILE):
    with open(LOG_FILE, "r") as f:
        query_logs = json.load(f)
else:
    query_logs = []

# ===================================================================
# 2. ЗАГРУЗКА ДАННЫХ (с кэшированием)
# ===================================================================

@st.cache_resource
def load_wikipedia():
    """Загружает датасет Википедии"""
    with st.spinner("📚 Загружаю Википедию..."):
        dataset = load_dataset(DATASET_NAME, LANGUAGE, split="train", streaming=True)
        articles = []
        for i, row in enumerate(dataset):
            if i >= ARTICLE_LIMIT:
                break
            articles.append({
                "id": row.get("id", i),
                "title": row.get("title", "Без названия"),
                "text": row.get("text", "")[:3000],
                "url": row.get("url", "")
            })
        return articles

@st.cache_resource
def load_embedder():
    """Загружает модель для эмбеддингов"""
    with st.spinner("🧠 Загружаю модель..."):
        return SentenceTransformer(MODEL_NAME)

@st.cache_resource
def create_embeddings(articles, embedder):
    """Создаёт эмбеддинги статей"""
    with st.spinner("🔢 Создаю эмбеддинги статей..."):
        texts = [f"{a['title']}\n\n{a['text']}" for a in articles]
        return embedder.encode(texts, normalize_embeddings=True)

# Загружаем всё
articles = load_wikipedia()
embedder = load_embedder()
embeddings = create_embeddings(articles, embedder)

# ===================================================================
# 3. ФУНКЦИИ ПОИСКА И ЛОГИРОВАНИЯ
# ===================================================================

def search_wikipedia(query):
    """Ищет ответ на вопрос в Википедии"""
    if not query:
        return None

    start_time = time.time()
    
    # Поиск
    query_vector = embedder.encode([query], normalize_embeddings=True)[0]
    scores = embeddings @ query_vector
    top_indices = np.argsort(-scores)[:3]

    results = []
    for idx in top_indices:
        score = float(scores[int(idx)])
        if score > 0.3:
            article = articles[int(idx)]
            results.append({
                "title": article['title'],
                "score": score,
                "text": article['text'][:1000],
                "url": article['url']
            })
    
    # Логируем запрос
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "query": query,
        "results_count": len(results),
        "response_time": round(time.time() - start_time, 2)
    }
    query_logs.append(log_entry)
    
    # Сохраняем логи
    with open(LOG_FILE, "w") as f:
        json.dump(query_logs[-100:], f)
    
    return results

def get_admin_stats():
    """Собирает статистику для админ-панели"""
    total_queries = len(query_logs)
    
    if total_queries > 0:
        avg_time = sum(q["response_time"] for q in query_logs) / total_queries
        popular_queries = sorted(query_logs, key=lambda x: x["results_count"], reverse=True)[:5]
    else:
        avg_time = 0
        popular_queries = []
    
    return {
        "total_queries": total_queries,
        "avg_time": avg_time,
        "popular_queries": popular_queries,
        "articles_count": len(articles),
        "model_name": MODEL_NAME
    }

def clear_logs():
    """Очищает логи"""
    global query_logs
    query_logs = []
    with open(LOG_FILE, "w") as f:
        json.dump(query_logs, f)
    return True

# ===================================================================
# 4. ИНТЕРФЕЙС
# ===================================================================

# --- БОКОВАЯ ПАНЕЛЬ (АДМИНКА) ---
with st.sidebar:
    st.image("https://cdn-icons-png.flaticon.com/512/4248/4248455.png", width=80)
    st.title("👑 Админ-панель")
    
    # Вход
    if "logged_in" not in st.session_state:
        st.session_state.logged_in = False
    
    if not st.session_state.logged_in:
        with st.form("login_form"):
            username = st.text_input("👤 Логин", placeholder="admin")
            password = st.text_input("🔑 Пароль", type="password", placeholder="hfpassword21")
            submitted = st.form_submit_button("🔑 Войти")
            
            if submitted:
                if username == ADMIN_USER and password == ADMIN_PASS:
                    st.session_state.logged_in = True
                    st.success("✅ Доступ разрешён!")
                    st.rerun()
                else:
                    st.error("❌ Неверный логин или пароль")
    else:
        st.success("✅ Вы вошли как администратор")
        
        # Кнопка выхода
        if st.button("🚪 Выйти"):
            st.session_state.logged_in = False
            st.rerun()
        
        st.divider()
        
        # Статистика
        st.subheader("📊 Статистика")
        stats = get_admin_stats()
        
        col1, col2 = st.columns(2)
        with col1:
            st.metric("Всего запросов", stats["total_queries"])
            st.metric("Загружено статей", stats["articles_count"])
        with col2:
            st.metric("Ср. время ответа", f"{stats['avg_time']:.2f}с")
            st.metric("Модель", stats["model_name"].split("/")[-1])
        
        # Популярные запросы
        if stats["popular_queries"]:
            st.subheader("🔥 Топ-5 запросов")
            for i, q in enumerate(stats["popular_queries"], 1):
                st.write(f"{i}. **{q['query']}** (найдено: {q['results_count']})")
        
        st.divider()
        
        # Логи
        st.subheader("📋 Последние запросы")
        if query_logs:
            # Показываем последние 10 в виде таблицы
            df = pd.DataFrame(query_logs[-10:])
            df["timestamp"] = pd.to_datetime(df["timestamp"]).dt.strftime("%H:%M:%S")
            st.dataframe(
                df[["timestamp", "query", "results_count", "response_time"]],
                column_config={
                    "timestamp": "Время",
                    "query": "Запрос",
                    "results_count": "Результатов",
                    "response_time": "Время (с)"
                },
                use_container_width=True,
                hide_index=True
            )
            
            if st.button("🗑️ Очистить логи", type="secondary"):
                clear_logs()
                st.success("Логи очищены!")
                st.rerun()
        else:
            st.info("📭 Логов пока нет")

# --- ОСНОВНАЯ ЧАСТЬ ---
st.title("🌖 Wikipedia Assistant")
st.markdown("Задай вопрос — я найду ответ в Википедии!")

# Информация о загрузке
st.success(f"✅ Загружено {len(articles)} статей")

# Поиск
query = st.text_input(
    "🔍 Что хочешь узнать?",
    placeholder="Например: Как измеряют расстояние до галактик?",
    key="query_input"
)

col1, col2 = st.columns([1, 5])
with col1:
    search_clicked = st.button("🔎 Найти", type="primary", use_container_width=True)

# Выполняем поиск
if query and (search_clicked or query != st.session_state.get("last_query", "")):
    st.session_state.last_query = query
    
    with st.spinner("🔎 Ищу ответ..."):
        results = search_wikipedia(query)
    
    if results:
        for i, result in enumerate(results, 1):
            with st.expander(f"#{i} {result['title']} (сходство: {result['score']:.2f})", expanded=i==1):
                st.write(result['text'] + "...")
                if result['url']:
                    st.link_button("🔗 Читать на Википедии", result['url'])
    else:
        st.warning("😕 Не нашёл подходящей статьи. Попробуй уточнить вопрос.")
        
elif not query:
    st.info("💡 Напиши вопрос, и я найду ответ в Википедии")

# ===================================================================
# 5. ДОПОЛНИТЕЛЬНЫЕ ВОЗМОЖНОСТИ
# ===================================================================

# Подвал
st.divider()
st.caption("🌖 Wikipedia Assistant | Использует BAAI/bge-large-en-v1.5")

# Кнопка внизу для переключения темы
if st.button("🎨 Сменить тему"):
    st.session_state.theme = "dark" if st.session_state.get("theme") != "dark" else "light"
    st.rerun()