src/bot.py - main - GreyCat/knowledge-assistant-bot

src/bot.py
53 строки2 KB
Vlad
Update bot.py
08 сен 2025, 23:14
08 сен 2025, 23:14f800b4d
import os, tempfile
from pathlib import Path
from dotenv import load_dotenv
from telegram import Update
from telegram.ext import ApplicationBuilder, MessageHandler, CommandHandler, ContextTypes, filters

from .chains import stuff_chain, refine_chain
from .retriever import search_with_meta, ingest_file, build_or_update_index

load_dotenv()
BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
KB_DIR = Path(os.getenv("KB_DIR","knowledge_base")); KB_DIR.mkdir(exist_ok=True)
TOP_K = int(os.getenv("TOP_K","5"))

async def start(update: Update, _: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text("Пришлите PDF/DOCX/TXT, потом задайте вопрос.\nКоманды: /reindex /stats")

async def reindex(update: Update, _: ContextTypes.DEFAULT_TYPE):
    try:
        build_or_update_index()
        await update.message.reply_text("Индекс перестроен.")
    except Exception as e:
        await update.message.reply_text(f"Ошибка: {e}")

async def handle_doc(update: Update, context: ContextTypes.DEFAULT_TYPE):
    file = await context.bot.get_file(update.message.document.file_id)
    with tempfile.TemporaryDirectory() as td:
        dst = Path(td)/update.message.document.file_name
        await context.bot.download_file(file.file_path, dst)
        # кладём копию в KB
        kb_path = KB_DIR/dst.name; kb_path.write_bytes(dst.read_bytes())
        n = ingest_file(kb_path)
    await update.message.reply_text(f"Добавлено фрагментов: {n} из {kb_path.name}")

async def handle_text(update: Update, _: ContextTypes.DEFAULT_TYPE):
    q = update.message.text.strip()
    hits = search_with_meta(q, TOP_K)
    context = "\n\n".join(f"[{i+1}] {h['text']}\nИсточник: {h['source']}" for i,h in enumerate(hits))
    chain = stuff_chain  # или refine_chain
    ans = chain(q, [type("D",(),{"page_content":h["text"]}) for h in hits])
    cites = "\n".join(f"• {h['source']}" for h in hits) or "—"
    await update.message.reply_text(f"{ans}\n\nЦитаты:\n{cites}")

def main():
    app = ApplicationBuilder().token(BOT_TOKEN).build()
    app.add_handler(CommandHandler("start", start))
    app.add_handler(CommandHandler("reindex", reindex))
    app.add_handler(MessageHandler(filters.Document.ALL, handle_doc))
    app.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), handle_text))
    app.run_polling()

if __name__ == "__main__":
    main()