app.py

import os
import json
import glob
import spaces
import requests
import numpy as np
import gradio as gr

from pypdf import PdfReader
import docx
from pptx import Presentation
import openpyxl
from sentence_transformers import SentenceTransformer

# 1. Загрузка модели эмбеддингов
print("Loading Embedding Model...")
embedder = SentenceTransformer('cointegrated/rubert-tiny2')

def extract_text_from_file(file_path):
    ext = os.path.splitext(file_path)[1].lower()
    text = ""
    try:
        if ext == ".pdf":
            reader = PdfReader(file_path)
            for page in reader.pages:
                text += (page.extract_text() or "") + "\n"
        elif ext in [".docx", ".doc"]:
            doc = docx.Document(file_path)
            text = "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
        elif ext in [".pptx", ".ppt"]:
            prs = Presentation(file_path)
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text") and shape.text:
                        text += shape.text + "\n"
        elif ext in [".xlsx", ".xls"]:
            wb = openpyxl.load_workbook(file_path, data_only=True)
            for sheet in wb.sheetnames:
                ws = wb[sheet]
                for row in ws.iter_rows(values_only=True):
                    row_str = " | ".join([str(cell) for cell in row if cell is not None])
                    if row_str.strip():
                        text += row_str + "\n"
        elif ext == ".txt":
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                text = f.read()
    except Exception as e:
        print(f"Ошибка чтения {file_path}: {e}")
    return text

def chunk_text(text, source_name, chunk_size=700, overlap=100):
    chunks = []
    for i in range(0, len(text), chunk_size - overlap):
        chunk = text[i:i + chunk_size].strip()
        if len(chunk) > 40:
            chunks.append({"text": chunk, "source": source_name})
    return chunks

def build_vector_index(folder="documents"):
    os.makedirs(folder, exist_ok=True)
    all_chunks = []
    files = glob.glob(f"{folder}/*.*")
    print(f"Найдено файлов в {folder}: {len(files)}")
    
    for f_path in files:
        filename = os.path.basename(f_path)
        raw_text = extract_text_from_file(f_path)
        if raw_text.strip():
            file_chunks = chunk_text(raw_text, source_name=filename)
            all_chunks.extend(file_chunks)
            print(f"Обработан файл {filename} (чанков: {len(file_chunks)})")

    if not all_chunks:
        all_chunks = [{"text": "База знаний пуста. Загрузите файлы в папку documents/.", "source": "system"}]

    texts = [c["text"] for c in all_chunks]
    embeddings = embedder.encode(texts, convert_to_numpy=True, show_progress_bar=False)
    return all_chunks, embeddings

CHUNKS, EMBEDDINGS = build_vector_index("documents")

def retrieve_relevant_context(query, top_k=4):
    if len(CHUNKS) == 0 or EMBEDDINGS is None:
        return ""
    query_emb = embedder.encode([query], convert_to_numpy=True)
    scores = np.dot(EMBEDDINGS, query_emb.T).squeeze()
    if np.ndim(scores) == 0:
        scores = np.array([scores])
    top_indices = np.argsort(scores)[::-1][:top_k]
    
    context_blocks = []
    for idx in top_indices:
        if scores[idx] > 0.25:
            chunk = CHUNKS[idx]
            context_blocks.append(f"[Источник: {chunk['source']}]\n{chunk['text']}")
    return "\n\n---\n\n".join(context_blocks)

SYSTEM_PROMPT = """Ты — «Инвест-Консультант Щёлково» — ведущий эксперт Администрации г.о. Щёлково.
Твоя задача — давать точные и подробные ответы НА ОСНОВЕ ПРЕДОСТАВЛЕННОГО КОНТЕКСТА ИЗ ДОКУМЕНТОВ.

Правила:
1. Отвечай доброжелательно, официально и строго по делу.
2. Используй только факты из найденных документов.
3. В конце ответа обязательно указывай, из каких файлов/источников взята информация."""

HF_TOKEN = os.environ.get("HF_TOKEN")

@spaces.GPU
def respond(message, history):
    if not HF_TOKEN:
        yield "⚠️ Ошибка: В настройках (Settings -> Secrets) не найден секрет HF_TOKEN!"
        return

    user_text = message.get("text", "") if isinstance(message, dict) else str(message)
    context = retrieve_relevant_context(user_text, top_k=4)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    if history:
        for item in history:
            if isinstance(item, (tuple, list)) and len(item) == 2:
                u, a = item
                if u: messages.append({"role": "user", "content": str(u)})
                if a: messages.append({"role": "assistant", "content": str(a)})
            elif isinstance(item, dict):
                r = item.get("role")
                c = item.get("content")
                if r in ["user", "assistant"] and c:
                    messages.append({"role": r, "content": str(c)})

    augmented_prompt = f"НАЙДЕННЫЕ ДОКУМЕНТЫ ИЗ БАЗЫ ЗНАНИЙ:\n{context}\n\nВОПРОС ПОЛЬЗОВАТЕЛЯ: {user_text}"
    messages.append({"role": "user", "content": augmented_prompt})

    url = "https://router.huggingface.co/v1/chat/completions"
    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
    payload = {
        "model": "Qwen/Qwen2.5-Coder-32B-Instruct",
        "messages": messages,
        "max_tokens": 700,
        "temperature": 0.2,
        "stream": True
    }

    full_response = ""
    try:
        res = requests.post(url, headers=headers, json=payload, stream=True, timeout=30)
        if res.status_code == 200:
            for line in res.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith("data: "):
                        data_str = line_str[6:].strip()
                        if data_str == "[DONE]": break
                        try:
                            data_json = json.loads(data_str)
                            if "choices" in data_json and len(data_json["choices"]) > 0:
                                delta = data_json["choices"][0]["delta"].get("content", "")
                                if delta:
                                    full_response += delta
                                    yield full_response
                        except Exception: pass
        else:
            yield f"⚠️ Ошибка сервера ({res.status_code}): {res.text}"
    except Exception as e:
        yield f"⚠️ Ошибка подключения: {str(e)}"

demo = gr.ChatInterface(
    fn=respond,
    title="Инвест-Консультант г.о. Щёлково",
    description="Умный ассистент Отдела по инвестициям с поиском по вашей базе знаний.",
    examples=[
        "Какие субсидии есть на покупку оборудования?",
        "Как получить землю за 1 рубль?",
        "Контакты отдела инвестиций"
    ],
)

if __name__ == "__main__":
    demo.launch()