app.py - master - frein/ckm

app.py
1 049 строк45 KB
frein
first_commit
17 мар 2026, 12:11
17 мар 2026, 12:1130c6069
import os
import re
import json
import math
import uuid
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional

import gradio as gr
from llm.config import GigaConfig
from llm.make_llm import make_llm
from tools.tariff_tool import TariffTool, QAPair


# -----------------------------
# Utilities
# -----------------------------


def normalize_text(text: str) -> str:
    return re.sub(r"\s+", " ", text.strip())


def tokenize_ru(text: str) -> List[str]:
    text = text.lower()
    # Keep Cyrillic letters and numbers, split on others
    tokens = re.split(r"[^а-яa-z0-9%]+", text)
    return [t for t in tokens if t]


# -----------------------------
# Document store and loaders
# -----------------------------


@dataclass
class Document:
    doc_id: str
    source: str
    text: str


class DocumentStore:
    def __init__(self):
        self.docs: Dict[str, Document] = {}

    def add_document(self, source: str, text: str) -> str:
        doc_id = str(uuid.uuid4())
        self.docs[doc_id] = Document(doc_id=doc_id, source=source, text=text)
        return doc_id

    def get_all(self) -> List[Document]:
        return list(self.docs.values())

    def load_from_folder(self, folder: str):
        if not os.path.isdir(folder):
            return
        for name in os.listdir(folder):
            path = os.path.join(folder, name)
            if not os.path.isfile(path):
                continue
            ext = os.path.splitext(name)[1].lower()
            if ext in [".txt", ".md"]:
                try:
                    with open(path, "r", encoding="utf-8", errors="ignore") as f:
                        text = f.read()
                except Exception:
                    text = ""
                self.add_document(source=path, text=text)
            elif ext == ".pdf":
                text = self._extract_pdf_text(path)
                self.add_document(source=path, text=text)

    @staticmethod
    def _extract_pdf_text(path: str) -> str:
        # Try PyPDF2 if available; otherwise fallback
        try:
            import PyPDF2  # type: ignore
            text_parts = []
            with open(path, "rb") as f:
                reader = PyPDF2.PdfReader(f)
                for page in reader.pages:
                    try:
                        text_parts.append(page.extract_text() or "")
                    except Exception:
                        pass
            return "\n".join(text_parts)
        except Exception:
            return f"[Парсинг PDF не доступен без PyPDF2] Источник: {os.path.basename(path)}"


# -----------------------------
# Simple Embedding Provider (local hashing)
# -----------------------------


class LocalHashingEmbedding:
    def __init__(self, dim: int = 256):
        self.dim = dim

    def embed(self, text: str) -> List[float]:
        vec = [0.0] * self.dim
        for tok in tokenize_ru(text):
            h = hash(tok) % self.dim
            vec[h] += 1.0
        # L2 normalize
        norm = math.sqrt(sum(v * v for v in vec)) or 1.0
        return [v / norm for v in vec]


class GigaLCEmbedder:
    def __init__(self, cfg: GigaConfig):
        try:
            from langchain_gigachat import GigaChatEmbeddings  # type: ignore
        except Exception as e:
            raise RuntimeError("langchain-gigachat не установлен или недоступен") from e
        kwargs = dict(credentials=cfg.api_key, verify_ssl_certs=cfg.verify_ssl_certs, scope=cfg.scope)
        if cfg.embed_model:
            kwargs["model"] = cfg.embed_model
        self._emb = GigaChatEmbeddings(**kwargs)

    def embed(self, text: str) -> List[float]:
        vecs = self._emb.embed_documents([text])
        return vecs[0] if vecs else []


def cosine(a: List[float], b: List[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


# -----------------------------
# Simple BM25 implementation
# -----------------------------


class BM25:
    def __init__(self, docs: List[str], k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.docs = [tokenize_ru(d) for d in docs]
        self.N = len(docs)
        self.avgdl = sum(len(d) for d in self.docs) / self.N if self.N else 0.0
        self.df: Dict[str, int] = {}
        for d in self.docs:
            for t in set(d):
                self.df[t] = self.df.get(t, 0) + 1
        self.idf: Dict[str, float] = {}
        for t, df in self.df.items():
            # IDF with added constants to avoid negative values on very common terms
            self.idf[t] = math.log((self.N - df + 0.5) / (df + 0.5) + 1)

    def score(self, query: str, idx: int) -> float:
        if not self.docs:
            return 0.0
        q = tokenize_ru(query)
        doc = self.docs[idx]
        tf: Dict[str, int] = {}
        for t in doc:
            tf[t] = tf.get(t, 0) + 1
        score = 0.0
        dl = len(doc)
        for t in q:
            if t not in tf:
                continue
            idf = self.idf.get(t, 0.0)
            denom = tf[t] + self.k1 * (1 - self.b + self.b * dl / (self.avgdl or 1.0))
            score += idf * (tf[t] * (self.k1 + 1)) / (denom or 1.0)
        return score


# -----------------------------
# Vector index
# -----------------------------


@dataclass
class Chunk:
    doc_id: str
    chunk_id: str
    source: str
    text: str
    embedding: List[float] = field(default_factory=list)


class VectorIndex:
    def __init__(self, embedder: LocalHashingEmbedding):
        self.embedder = embedder
        self.chunks: List[Chunk] = []

    def build(self, docs: List[Document], chunk_size: int = 800, chunk_overlap: int = 100):
        self.chunks = []
        for d in docs:
            text = normalize_text(d.text)
            start = 0
            while start < len(text):
                end = min(len(text), start + chunk_size)
                chunk_text = text[start:end]
                if not chunk_text.strip():
                    break
                c = Chunk(
                    doc_id=d.doc_id,
                    chunk_id=str(uuid.uuid4()),
                    source=d.source,
                    text=chunk_text,
                )
                c.embedding = self.embedder.embed(c.text)
                self.chunks.append(c)
                if end == len(text):
                    break
                start = end - chunk_overlap

    def search(self, query: str, top_k: int = 5) -> List[Tuple[Chunk, float]]:
        if not self.chunks:
            return []
        q_emb = self.embedder.embed(query)
        scored = [(c, cosine(q_emb, c.embedding)) for c in self.chunks]
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

    def load_chunks(self, chunks: List[Chunk]):
        self.chunks = chunks


class HybridRetriever:
    def __init__(self, vector_index: VectorIndex, bm25: Optional[BM25] = None, w_vec: float = 0.5, w_bm25: float = 0.5):
        self.vector_index = vector_index
        self.bm25 = bm25
        self.w_vec = w_vec
        self.w_bm25 = w_bm25

    def search(self, query: str, top_k: int = 5) -> List[Tuple[Chunk, float]]:
        vec_res = self.vector_index.search(query, top_k=top_k * 2)
        if not self.bm25:
            return vec_res[:top_k]

        # Build BM25 scores per chunk using chunk text
        bm25_scores = []
        docs_texts = [c.text for c, _ in vec_res]
        bm25_local = BM25(docs_texts)
        for i, (c, _) in enumerate(vec_res):
            bm25_scores.append((c, bm25_local.score(query, i)))

        # Normalize scores
        def normalize(scores: List[Tuple[Chunk, float]]) -> Dict[str, float]:
            if not scores:
                return {}
            vals = [s for _, s in scores]
            mx = max(vals) if vals else 1.0
            mn = min(vals) if vals else 0.0
            denom = (mx - mn) or 1.0
            return {c.chunk_id: (s - mn) / denom for c, s in scores}

        vec_norm = normalize(vec_res)
        bm_norm = normalize(bm25_scores)

        combined = []
        for c, _ in vec_res:
            score = self.w_vec * vec_norm.get(c.chunk_id, 0.0) + self.w_bm25 * bm_norm.get(c.chunk_id, 0.0)
            combined.append((c, score))
        combined.sort(key=lambda x: x[1], reverse=True)
        return combined[:top_k]


# -----------------------------
# Tariff parameter extraction
# -----------------------------


def extract_tariff_params(text: str) -> Dict[str, str]:
    t = text.lower()
    params: Dict[str, str] = {}

    # Monthly subscription cost
    m = re.search(r"(\d+[\s\d]*)\s*руб", t)
    if m:
        val = m.group(1).replace(" ", "")
        params["monthly_subscription_cost_rub"] = val

    # Free legal entity payments per month
    m = re.search(r"бесплатн\w*\s*(\d+)\s*плате?ж(ки|ек)?", t)
    if m:
        params["free_legal_payments_per_month"] = m.group(1)

    # Fee per extra transfer (100 рублей)
    m = re.search(r"по\s*(\d+[\s\d]*)\s*руб(?:\.|\b).{0,20}(перевод|платеж)", t)
    if m:
        params["extra_legal_payment_fee_rub"] = m.group(1).replace(" ", "")

    # Free individual payments limit per month (e.g., 500 тысяч)
    m = re.search(r"(\d+[\s\d]*)\s*(тыс|тысяч|млн|миллион)\w*\s*в\s*месяц\s*без\s*комиссии", t)
    if m:
        num = m.group(1).replace(" ", "")
        scale = m.group(2)
        if scale.startswith("тыс"):
            num_val = int(num) * 1000
        elif scale.startswith("мл") or scale.startswith("мил"):
            num_val = int(num) * 1000000
        else:
            num_val = int(num)
        params["free_individual_payments_limit_rub"] = str(num_val)

    # Fees for individuals above limit (2.5% / 3.5%)
    m = re.search(r"(\d+[\.,]?\d*)%\s*карта\s*сбербанка.*?(\d+[\.,]?\d*)%\s*если\s*карта\s*другого\s*банка", t)
    if m:
        params["fee_above_limit_sber_percent"] = m.group(1).replace(",", ".")
        params["fee_above_limit_other_percent"] = m.group(2).replace(",", ".")

    # Cash deposit free limit (e.g., 50 тысяч)
    m = re.search(r"пополнени\w*\s*сч[её]та.*?(\d+[\s\d]*)\s*(тыс|тысяч|млн|миллион)\w*\s*без\s*комиссии", t)
    if m:
        num = m.group(1).replace(" ", "")
        scale = m.group(2)
        if scale.startswith("тыс"):
            num_val = int(num) * 1000
        elif scale.startswith("мл") or scale.startswith("мил"):
            num_val = int(num) * 1000000
        else:
            num_val = int(num)
        params["cash_deposit_free_limit_rub"] = str(num_val)

    # Cash deposit fee above (e.g., 0,4%)
    m = re.search(r"стандартн\w*\s*тариф\w*\s*(\d+[\.,]?\d*)%\s*от\s*суммы", t)
    if m:
        params["cash_deposit_fee_above_percent"] = m.group(1).replace(",", ".")

    # Business card included
    if re.search(r"бизнес-?карта|бизнес\s*карта", t):
        params["business_card_included"] = "да"

    return params


# -----------------------------
# Tariff tool wiring
# -----------------------------


def build_qa_pairs() -> List[QAPair]:
    qa: List[QAPair] = []
    if not STATE.qual_questions:
        return qa
    for idx, q in enumerate(STATE.qual_questions):
        ans = STATE.qa_answers.get(idx, "")
        qa.append(QAPair(question=str(q), answer=str(ans)))
    return qa


# -----------------------------
# Simple dialog manager (rule-based)
# -----------------------------


def detect_yes(msg: str) -> bool:
    m = msg.lower()
    return any(x in m for x in ["да", "актуально", "+", "ок", "хорошо", "угу"])


def detect_no(msg: str) -> bool:
    m = msg.lower()
    return any(x in m for x in ["нет", "неактуально", "позже"])


def is_user_question(msg: str) -> bool:
    m = (msg or "").lower().strip()
    if "?" in m:
        return True
    # common Russian interrogatives and prompt verbs
    q_tokens = [
        "как", "почему", "зачем", "когда", "где", "чем", "сколько",
        "какой", "какая", "какие", "расскажи", "объясни", "сравни",
        "подробнее", "расскажите", "покажи", "покажите",
    ]
    return any(t in m for t in q_tokens)


def extract_profile_updates(msg: str) -> Dict[str, str]:
    m = msg.lower()
    upd: Dict[str, str] = {}

    if "ип" in m:
        upd["org_type"] = "ИП"
    if "ооо" in m:
        upd["org_type"] = "ООО"

    # number of payments (to legal entities)
    m2 = re.search(r"(\d+)[\s\-]*([\d]+)?", m)
    if m2:
        try:
            n1 = int(m2.group(1))
            upd["legal_payments_per_month"] = str(n1)
        except Exception:
            pass

    # payments to individuals
    if re.search(r"физическ|физ\.\s*лиц|физлиц|физикам", m):
        if detect_yes(m):
            upd["payments_to_individuals"] = "да"
        elif detect_no(m):
            upd["payments_to_individuals"] = "нет"

    # cash deposit need
    if re.search(r"пополнен(ие|ия)|наличн", m):
        if detect_yes(m):
            upd["cash_deposit_needed"] = "да"
        elif detect_no(m):
            upd["cash_deposit_needed"] = "нет"

    # preferred city/address hints are ignored in prototype
    return upd


def next_question_by_list() -> Optional[str]:
    if not STATE.qual_questions:
        return None
    idx = len(STATE.qa_answers)
    if idx < len(STATE.qual_questions):
        return STATE.qual_questions[idx]
    return None


def build_recommendation_with_tool() -> str:
    qa = build_qa_pairs()
    rec = STATE.tariff_tool.recommend(qa, extra={"payments_per_month": STATE.derived.get("payments_per_month")})
    parts = [
        f"Исходя из ваших ответов рекомендую тариф: {rec.name}.",
        f"Ориентировочная стоимость: {rec.price_rub:.0f} ₽ в месяц.",
        f"Основание: {rec.basis}.",
        "Готовы оформить?",
    ]
    return " ".join(parts)


# -----------------------------
# App state and logic
# -----------------------------


class AppState:
    def __init__(self):
        self.doc_store = DocumentStore()
        # Embeddings provider state
        self.embedder: object
        self.embedder_kind: str = "local"  # or "giga"
        self.embedder = LocalHashingEmbedding()
        self.vector_index = VectorIndex(self.embedder)
        self.bm25_global: Optional[BM25] = None
        self.retriever = HybridRetriever(self.vector_index, self.bm25_global)
        self.tariff_tool = TariffTool()
        self.profile: Dict[str, str] = {}
        # Qualification
        self.qual_questions: List[str] = []
        self.qual_path: Optional[str] = None
        self.qa_answers: Dict[int, str] = {}
        self.qual_started: bool = False
        # Sources
        self.prime_tariffs_path: Optional[str] = None
        # RAG over rko_rag_corpus only
        self.rag_vector_index = VectorIndex(self.embedder)
        self.rag_bm25: Optional[BM25] = None
        self.rag_retriever = HybridRetriever(self.rag_vector_index, self.rag_bm25)
        self.rag_sources: List[str] = []
        # GigaChat config
        self.giga_cfg: Optional[GigaConfig] = None
        self.giga_client = None  # reserved for HTTP client if needed
        # Chunking params and cache
        self.chunk_size: int = 800
        self.chunk_overlap: int = 100
        self.rebuild_on_start: bool = False

    def build_indexes(self, bm25_enable: bool = True, w_vec: float = 0.6, w_bm25: float = 0.4, use_cache: bool = True):
        docs = self.doc_store.get_all()
        if use_cache and self._try_load_cache_global(docs):
            pass
        else:
            self.vector_index.build(docs, chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap)
            self._save_cache_global(docs)
        if bm25_enable:
            self.bm25_global = BM25([d.text for d in docs])
        else:
            self.bm25_global = None
        self.retriever = HybridRetriever(self.vector_index, self.bm25_global, w_vec=w_vec, w_bm25=w_bm25)

    def build_rag_indexes(self, bm25_enable: bool = True, w_vec: float = 0.6, w_bm25: float = 0.4, use_cache: bool = True):
        texts, paths = get_rko_rag_corpus_texts(self.doc_store)
        # Build temporary document list for chunks
        rag_docs = []
        for i, t in enumerate(texts):
            rag_docs.append(Document(doc_id=str(i), source=paths[i] if i < len(paths) else f"rko_{i}", text=t))
        self.rag_sources = paths
        if use_cache and self._try_load_cache_rag(rag_docs):
            pass
        else:
            self.rag_vector_index.build(rag_docs, chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap)
            self._save_cache_rag(rag_docs)
        if bm25_enable:
            self.rag_bm25 = BM25([d.text for d in rag_docs])
        else:
            self.rag_bm25 = None
        self.rag_retriever = HybridRetriever(self.rag_vector_index, self.rag_bm25, w_vec=w_vec, w_bm25=w_bm25)

    def set_embedder(self, kind: str):
        kind = kind.lower().strip()
        if kind == "giga" and self.giga_cfg is not None:
            self.embedder_kind = "giga"
            self.embedder = GigaLCEmbedder(self.giga_cfg)
        else:
            self.embedder_kind = "local"
            self.embedder = LocalHashingEmbedding()
        # Recreate vector indexes with new embedder reference
        self.vector_index = VectorIndex(self.embedder)
        self.rag_vector_index = VectorIndex(self.embedder)

    def reload_from_data(self, bm25_enable: bool = True, w_vec: float = 0.6, w_bm25: float = 0.4):
        self.doc_store = DocumentStore()
        self.doc_store.load_from_folder("data")
        self.build_indexes(bm25_enable, w_vec, w_bm25, use_cache=not self.rebuild_on_start)
        # Prime tariffs path only
        _, path = get_prime_tariffs_texts(self.doc_store)
        self.prime_tariffs_path = path
        # Build RAG indexes limited to rko_rag_corpus
        self.build_rag_indexes(bm25_enable, w_vec, w_bm25, use_cache=not self.rebuild_on_start)
        # Load qualification questions
        self.qual_questions, self.qual_path = load_qual_questions(self.doc_store)
        self.qa_answers = {}
        self.qual_started = False

    # --------------- Cache helpers ---------------
    def _cache_dir(self) -> str:
        d = os.path.join(".cache")
        os.makedirs(d, exist_ok=True)
        return d

    def _files_fingerprint(self, docs: List[Document]) -> str:
        import hashlib
        parts = []
        for d in docs:
            p = d.source
            try:
                st = os.stat(p)
                parts.append(f"{os.path.abspath(p)}|{st.st_mtime_ns}|{st.st_size}")
            except Exception:
                parts.append(f"{os.path.abspath(p)}|0|0")
        meta = "\n".join(sorted(parts))
        h = hashlib.sha256()
        h.update(meta.encode("utf-8"))
        return h.hexdigest()

    def _fingerprint(self, docs: List[Document]) -> str:
        import hashlib
        embed_model = "local" if self.embedder_kind == "local" else (self.giga_cfg.embed_model or "giga") if self.giga_cfg else "giga"
        base = f"{self.embedder_kind}|{embed_model}|{self.chunk_size}|{self.chunk_overlap}|{self._files_fingerprint(docs)}"
        return hashlib.sha256(base.encode("utf-8")).hexdigest()

    def _cache_path(self, name: str, fp: str) -> str:
        return os.path.join(self._cache_dir(), f"{name}_{fp}.json")

    def _save_chunks_json(self, path: str, chunks: List[Chunk]):
        try:
            data = [
                {
                    "doc_id": c.doc_id,
                    "chunk_id": c.chunk_id,
                    "source": c.source,
                    "text": c.text,
                    "embedding": c.embedding,
                }
                for c in chunks
            ]
            with open(path, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False)
        except Exception:
            pass

    def _load_chunks_json(self, path: str) -> Optional[List[Chunk]]:
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            chunks: List[Chunk] = []
            for it in data:
                chunks.append(Chunk(
                    doc_id=it.get("doc_id", ""),
                    chunk_id=it.get("chunk_id", str(uuid.uuid4())),
                    source=it.get("source", ""),
                    text=it.get("text", ""),
                    embedding=it.get("embedding", []),
                ))
            return chunks
        except Exception:
            return None

    def _try_load_cache_global(self, docs: List[Document]) -> bool:
        fp = self._fingerprint(docs)
        path = self._cache_path("global", fp)
        chunks = self._load_chunks_json(path)
        if chunks is None:
            return False
        self.vector_index.load_chunks(chunks)
        return True

    def _save_cache_global(self, docs: List[Document]):
        fp = self._fingerprint(docs)
        path = self._cache_path("global", fp)
        self._save_chunks_json(path, self.vector_index.chunks)

    def _try_load_cache_rag(self, rag_docs: List[Document]) -> bool:
        fp = self._fingerprint(rag_docs)
        path = self._cache_path("rag", fp)
        chunks = self._load_chunks_json(path)
        if chunks is None:
            return False
        self.rag_vector_index.load_chunks(chunks)
        return True

    def _save_cache_rag(self, rag_docs: List[Document]):
        fp = self._fingerprint(rag_docs)
        path = self._cache_path("rag", fp)
        self._save_chunks_json(path, self.rag_vector_index.chunks)


STATE = AppState()
STATE.derived = {}
STATE.agent_stage = None


def init_load():
    # Load default data from ./data
    STATE.reload_from_data()


def get_prime_tariffs_texts(store: DocumentStore) -> Tuple[List[str], Optional[str]]:
    texts: List[str] = []
    path: Optional[str] = None
    for d in store.get_all():
        base = os.path.splitext(os.path.basename(d.source))[0].lower()
        if base == "prime_tariffs":
            texts.append(d.text)
            path = d.source
    return texts, path


def get_rko_rag_corpus_texts(store: DocumentStore) -> Tuple[List[str], List[str]]:
    texts: List[str] = []
    paths: List[str] = []
    for d in store.get_all():
        base = os.path.splitext(os.path.basename(d.source))[0].lower()
        if base == "rko_rag_corpus":
            texts.append(d.text)
            paths.append(d.source)
    return texts, paths


def load_qual_questions(store: DocumentStore) -> Tuple[List[str], Optional[str]]:
    questions_text = None
    path: Optional[str] = None
    for d in store.get_all():
        base = os.path.splitext(os.path.basename(d.source))[0].lower()
        if base == "список_вопросов":
            questions_text = d.text
            path = d.source
            break
    questions: List[str] = []
    if questions_text:
        for line in questions_text.splitlines():
            line = line.strip()
            if not line:
                continue
            line = re.sub(r"^\s*(\d+\.|[-•])\s*", "", line)
            questions.append(line)
    return questions, path


def ui_refresh_questions():
    STATE.qual_questions, STATE.qual_path = load_qual_questions(STATE.doc_store)
    STATE.qa_answers = {}
    STATE.qual_started = False
    q_md = "\n".join(f"- {q}" for q in STATE.qual_questions) or "Не найдено"
    src = f"Источник вопросов: `{STATE.qual_path}`" if STATE.qual_path else "Источник вопросов: не найден"
    return q_md, src



# removed: stub setter — tariffs are fixed (Тариф1/2/3)


def ui_rebuild_indexes(bm25_enable, w_vec, w_bm25):
    STATE.build_indexes(bool(bm25_enable), float(w_vec), float(w_bm25), use_cache=False)
    STATE.build_rag_indexes(bool(bm25_enable), float(w_vec), float(w_bm25), use_cache=False)
    return "Индексы перестроены (кэш обновлён)"


def start_message() -> str:
    return (
        "Добрый день, меня зовут Юлия, Сбербанк. Ранее вы оставляли заявку на открытие расчетного счета, скажите, актуально для вас?"
    )


def _as_text(content) -> str:
    if isinstance(content, str):
        return content
    # Gradio messages may pass content as list of segments
    if isinstance(content, list):
        parts = []
        for it in content:
            if isinstance(it, dict):
                # Common shapes: {"type":"text","text":"..."}
                txt = it.get("text") or it.get("content") or it.get("value")
                if isinstance(txt, str):
                    parts.append(txt)
                else:
                    parts.append(str(txt))
            else:
                parts.append(str(it))
        return "\n".join([p for p in parts if p])
    return str(content or "")


def _format_history(messages) -> str:
    lines = []
    for m in messages or []:
        if isinstance(m, dict):
            role = m.get("role", "assistant")
            content = _as_text(m.get("content", "")).strip()
            if not content:
                continue
            if role == "user":
                lines.append(f"Клиент: {content}")
            else:
                lines.append(f"Юлия: {content}")
        elif isinstance(m, (list, tuple)) and len(m) == 2:
            u, a = m
            u_text = _as_text(u).strip()
            a_text = _as_text(a).strip()
            if u_text:
                lines.append(f"Клиент: {u_text}")
            if a_text:
                lines.append(f"Юлия: {a_text}")
    return "\n".join(lines)


def _retrieve_context(query: str, k: int = 3) -> str:
    try:
        results = STATE.rag_retriever.search(query, top_k=k)
    except Exception:
        results = []
    parts = []
    for c, score in results:
        snippet = c.text.strip().replace("\n", " ")
        src = os.path.basename(c.source)
        parts.append(f"[{src}] {snippet}")
    return "\n".join(parts)


def agent_reply(messages: List[Dict[str, str]], user_msg: str) -> str:
    m = (user_msg or "").lower()
    # If no LLM configured, fallback to simple flow
    if STATE.giga_cfg is None:
        if any(x in m for x in ["готов", "оформ", "давайте оформ", "оформим", "да, оформим"]):
            return (
                "Отлично! Могу рассказать про оформление: потребуется паспорт руководителя (для ИП — паспорт), ИНН и СНИЛС можно озвучить по номерам. Офис работает по будням с 9:00 до 18:30 и в субботу с 10:00 до 15:30. Подтвердите, пожалуйста, ваш город, и предложу ближайший офис."
            )
        if not STATE.qual_started:
            if detect_no(m):
                return "Поняла, тогда при необходимости буду на связи. Хорошего дня!"
            STATE.qual_started = True
            q = next_question_by_list()
            return q or "Уточните, пожалуйста, детали, чтобы я подобрала тариф."
        idx = len(STATE.qa_answers)
        if STATE.qual_questions and idx < len(STATE.qual_questions):
            STATE.qa_answers[idx] = user_msg.strip()
            q = next_question_by_list()
            if q:
                return q
            return build_recommendation_with_tool()
        return "Готова помочь с подбором тарифа. Подтвердите, пожалуйста, актуальность заявки."

    # LLM-driven flow
    system_rules = (
        "Вы — цифровой клиентский менеджер СберБизнес по имени Юлия. "
        "Ваша задача: собрать ответы на вопросы из чек-листа, ответить на вопросы клиента, подобрать тариф и предложить оформить. "
        "Правила: 1) Задавайте по одному уточняющему вопросу за раз. 2) Используйте список вопросов как чек-лист, но формулируйте естественно. 3) Если клиент задаёт дополнительный вопрос — СНАЧАЛА кратко и по делу ответьте на него (опираясь на контекст), и только затем продолжайте сбор параметров/предложение тарифа. "
        "4) Если на все пункты даны ответы — представьте рекомендованный тариф с указанной стоимостью и спросите о согласии. 5) После согласия клиента предложите рассказать про оформление. 6) Отвечая на вопросы по условиям, опирайтесь на контекст (если есть). "
        "7) Если клиент говорит, что потребности в оформлении сейчас нет или тема не актуальна — всё равно отвечайте на его вопросы, без навязывания; в конце мягко предложите обратиться позже. 8) Пишите на русском, вежливо, кратко и по делу. "
        "В конце сообщения добавляйте управляющий маркер: #STATE:ASK — если ещё есть неотвеченные вопросы; #STATE:RECOMMEND — если готовы рекомендовать; #STATE:PROCEED — если клиент согласился и вы предлагаете оформление. "
        "Также добавьте строку JSON с извлечёнными параметрами: STATE:{\"payments_per_month\": <число или null>}"
    )
    questions_text = "\n".join(f"- {q}" for q in (STATE.qual_questions or [])) or ""
    history_text = _format_history(messages)
    context_text = _retrieve_context(user_msg, k=3)
    _cur_rec = STATE.tariff_tool.recommend(build_qa_pairs())
    tariff_name = _cur_rec.name
    tariff_price = f"{_cur_rec.price_rub:.0f} ₽/мес"

    prompt = (
        f"{system_rules}\n\n"
        f"Чек-лист вопросов:\n{questions_text}\n\n"
        f"Рекомендованный тариф (для финального шага): {tariff_name}, стоимость {tariff_price}.\n\n"
        f"История диалога (формат Клиент/Юлия):\n{history_text}\n\n"
        f"Контекст (справка, если уместно):\n{context_text}\n\n"
        f"Текущее сообщение клиента: {user_msg}\n\n"
        f"Ответьте как Юлия."
    )

    try:
        llm = make_llm(STATE.giga_cfg)
        out = llm.invoke(prompt)
        text = getattr(out, "content", None) or str(out)
        # Parse control state and inline JSON
        state_tag = None
        m_tag = re.search(r"#STATE:(ASK|RECOMMEND|PROCEED)", text, flags=re.IGNORECASE)
        if m_tag:
            state_tag = m_tag.group(1).upper()
            STATE.agent_stage = state_tag
        # Strip control tags and STATE JSON from visible text
        visible = re.sub(r"#STATE:(ASK|RECOMMEND|PROCEED)", "", text, flags=re.IGNORECASE)
        visible = re.sub(r"STATE\s*:\s*\{.*?\}", "", visible, flags=re.IGNORECASE | re.DOTALL)
        visible = visible.strip()
        # Try parse inline STATE JSON
        try:
            m = re.search(r"STATE\s*:\s*(\{.*?\})", text, flags=re.IGNORECASE | re.DOTALL)
            if m:
                import json as _json
                data = _json.loads(m.group(1))
                if isinstance(data, dict) and "payments_per_month" in data:
                    v = data.get("payments_per_month")
                    if isinstance(v, (int, float)):
                        STATE.derived["payments_per_month"] = int(v)
                    elif isinstance(v, str) and v.strip().isdigit():
                        STATE.derived["payments_per_month"] = int(v.strip())
        except Exception:
            pass
        # If LLM indicates RECOMMEND, override with deterministic recommendation
        if STATE.agent_stage == "RECOMMEND" and not is_user_question(user_msg):
            return build_recommendation_with_tool()
        if STATE.agent_stage == "PROCEED":
            return (
                "Отлично! Могу рассказать про оформление: потребуется паспорт руководителя (для ИП — паспорт), ИНН и СНИЛС можно озвучить по номерам. Офис работает по будням с 9:00 до 18:30 и в субботу с 10:00 до 15:30. Подтвердите, пожалуйста, ваш город, и предложу ближайший офис."
            )
        return visible
    except Exception as e:
        # Fallback if LLM fails
        return (
            "Извините, временно не удалось обратиться к LLM. "
            "Давайте продолжим: " + (next_question_by_list() or "Подтвердите, пожалуйста, актуальность заявки.")
        )


def ui_chat(user_message, messages):
    if messages is None or len(messages) == 0:
        messages = [{"role": "assistant", "content": start_message()}]
        if not user_message:
            return messages, ""

    # Append user message
    messages.append({"role": "user", "content": user_message})
    # Record answer into QA buffer (sequentially)
    try:
        idx = len(STATE.qa_answers)
        if STATE.qual_questions and idx < len(STATE.qual_questions):
            STATE.qa_answers[idx] = user_message.strip()
    except Exception:
        pass
    # Retrieve context (not shown to user explicitly, but could be)
    _ = STATE.rag_retriever.search(user_message, top_k=3)

    # Generate agent reply
    reply = agent_reply(messages, user_message)
    messages.append({"role": "assistant", "content": reply})
    return messages, ""


def ui_calc_tariff(messages):
    qa = build_qa_pairs()
    rec = STATE.tariff_tool.recommend(qa, extra={"payments_per_month": STATE.derived.get("payments_per_month")})
    text = (
        f"Рассчитано (заглушка): {rec.name}, ориентировочно {rec.price_rub:.0f} ₽/мес.\n"
        f"Основание: {rec.basis}"
    )
    if messages is None or len(messages) == 0:
        messages = [{"role": "assistant", "content": start_message()}]
    messages.append({"role": "assistant", "content": text})
    return messages


def build_interface():
    with gr.Blocks(title="СберБизнес: Цифровой КМ (прототип)") as demo:
        gr.Markdown("""
        Прототип цифрового клиентского менеджера СберБизнес.
        - Страница настроек: список вопросов квалификации (из файла `список_вопросов` в папке data), настройка поиска (RAG) и заглушки тарифа (название и сумма).
        - Страница чата: диалог с клиентом, поочередные вопросы из списка, рекомендация тарифа по заглушке и предложение оформить.
        """)

        with gr.Tabs():
            with gr.TabItem("LLM / GигаChat"):
                gr.Markdown("Настройка подключения к GigaChat и параметров LLM.")
                with gr.Row():
                    api_key = gr.Textbox(label="API Key", type="password")
                with gr.Row():
                    temperature = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Temperature")
                    max_tokens = gr.Number(value=None, label="Max tokens (опционально)")
                with gr.Row():
                    use_giga_embeddings = gr.Checkbox(value=False, label="Использовать эмбеддинги GigaChat для индексов")
                with gr.Row():
                    save_llm = gr.Button("Сохранить настройки")
                    test_embed = gr.Button("Тест эмбеддингов")
                    test_chat = gr.Button("Тест чата")
                with gr.Row():
                    llm_out = gr.Markdown()

                def on_save_llm(ak, temp, mtok, use_giga):
                    try:
                        cfg = GigaConfig(api_key=str(ak or ""), chat_model="GigaChat-2-Max", embed_model="", temperature=float(temp or 0.2), max_tokens=(int(mtok) if mtok is not None else None), verify_ssl_certs=False, scope="GIGACHAT_API_CORP")
                        STATE.giga_cfg = cfg
                        if use_giga:
                            try:
                                STATE.set_embedder("giga")
                                # quick probe to fail fast on invalid model
                                emb = GigaLCEmbedder(cfg)
                                _ = emb.embed("ping")
                                STATE.build_indexes()
                                STATE.build_rag_indexes()
                                return "Настройки сохранены. Эмбеддинги GigaChat активированы."
                            except Exception as ee:
                                STATE.set_embedder("local")
                                STATE.build_indexes()
                                STATE.build_rag_indexes()
                                return f"Настройки сохранены, но эмбеддинги GigaChat недоступны: {ee}. Использую локальные эмбеддинги."
                        else:
                            STATE.set_embedder("local")
                            STATE.build_indexes()
                            STATE.build_rag_indexes()
                            return "Настройки сохранены. Локальные эмбеддинги активны."
                    except Exception as e:
                        return f"Ошибка сохранения/инициализации: {e}"

                def on_test_embed():
                    if not STATE.giga_cfg:
                        return "GigaChat не сконфигурирован"
                    try:
                        emb = GigaLCEmbedder(STATE.giga_cfg)
                        vec = emb.embed("тест")
                        return f"Эмбеддинги ОК. Размер вектора: {len(vec) if vec else 'неизв.'}"
                    except Exception as e:
                        return f"Ошибка эмбеддингов: {e}"

                def on_test_chat():
                    if not STATE.giga_cfg:
                        return "GigaChat не сконфигурирован"
                    try:
                        llm = make_llm(STATE.giga_cfg)
                        # LangChain GigaChat — ChatModel. Используем invoke с промптом.
                        content = llm.invoke("Скажите привет одним словом")
                        # Ответ может быть str или Message
                        text = getattr(content, "content", None) or str(content)
                        return f"Чат ОК (model=GigaChat-2-Max). Ответ: {text[:200]}"
                    except Exception as e:
                        return f"Ошибка чата: {e}"

                save_llm.click(on_save_llm, inputs=[api_key, temperature, max_tokens, use_giga_embeddings], outputs=[llm_out])
                test_embed.click(on_test_embed, outputs=[llm_out])
                test_chat.click(on_test_chat, outputs=[llm_out])

            with gr.TabItem("Настройки"):
                with gr.Row():
                    qual_md = gr.Markdown(value=("\n".join(f"- {q}" for q in STATE.qual_questions) or "Не найдено"))
                with gr.Row():
                    qual_src = gr.Markdown(value=(f"Источник вопросов: `{STATE.qual_path}`" if STATE.qual_path else "Источник вопросов: не найден"))
                with gr.Row():
                    rag_src = gr.Markdown(value=("RAG источник: " + (", ".join(f"`{p}`" for p in STATE.rag_sources) if STATE.rag_sources else "не найден")))
                with gr.Row():
                    refresh_q_btn = gr.Button("Обновить вопросы из data")
                with gr.Row():
                    rebuild_on_start = gr.Checkbox(value=False, label="Пересобирать индексы при старте (игнорировать кэш)")
                    clear_cache_btn = gr.Button("Сбросить кэш индексов")
                    cache_out = gr.Markdown()
                # Tariff stub is fixed tiers (Тариф1/2/3). No manual override UI.
                with gr.Row():
                    bm25_enable = gr.Checkbox(value=True, label="Включить BM25")
                    w_vec = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Вес эмбеддингов")
                    w_bm25 = gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Вес BM25")
                    rebuild_btn = gr.Button("Перестроить индексы")
                    rebuild_out = gr.Markdown()

                refresh_q_btn.click(ui_refresh_questions, outputs=[qual_md, qual_src])
                def on_toggle_rebuild(val):
                    STATE.rebuild_on_start = bool(val)
                    return f"Пересборка при старте: {'включена' if STATE.rebuild_on_start else 'выключена'}"

                def on_clear_cache():
                    try:
                        d = os.path.join('.cache')
                        if os.path.isdir(d):
                            for name in os.listdir(d):
                                p = os.path.join(d, name)
                                try:
                                    if os.path.isfile(p):
                                        os.remove(p)
                                except Exception:
                                    pass
                        return "Кэш очищен"
                    except Exception as e:
                        return f"Не удалось очистить кэш: {e}"

                rebuild_on_start.change(on_toggle_rebuild, inputs=[rebuild_on_start], outputs=[cache_out])
                clear_cache_btn.click(on_clear_cache, outputs=[cache_out])
                rebuild_btn.click(ui_rebuild_indexes, inputs=[bm25_enable, w_vec, w_bm25], outputs=[rebuild_out])

            with gr.TabItem("Состояние"):
                gr.Markdown("Текущее состояние агента (распознанные параметры и этап). Маркеры #STATE и JSON скрыты из чата и видны только здесь.")
                state_view = gr.Code(label="Agent State (JSON)", language="json")
                refresh_state = gr.Button("Обновить состояние")

                def ui_get_state():
                    data = {
                        "stage": STATE.agent_stage,
                        "derived": {
                            "payments_per_month": STATE.derived.get("payments_per_month"),
                        },
                        "qa_answers_count": len(STATE.qa_answers or {}),
                    }
                    return json.dumps(data, ensure_ascii=False, indent=2)

                refresh_state.click(ui_get_state, outputs=[state_view])

            with gr.TabItem("Чат"):
                chatbot = gr.Chatbot(label="Диалог", height=450)
                user_in = gr.Textbox(label="Сообщение клиента", placeholder="Введите сообщение и нажмите Отправить")
                with gr.Row():
                    send_btn = gr.Button("Отправить", variant="primary")
                    calc_btn = gr.Button("Рассчитать тариф (заглушка)")

                def on_load():
                    return [{"role": "assistant", "content": start_message()}]

                demo.load(on_load, outputs=[chatbot])
                send_btn.click(ui_chat, inputs=[user_in, chatbot], outputs=[chatbot, user_in])
                user_in.submit(ui_chat, inputs=[user_in, chatbot], outputs=[chatbot, user_in])
                calc_btn.click(ui_calc_tariff, inputs=[chatbot], outputs=[chatbot])

    return demo


def main():
    init_load()
    demo = build_interface()
    demo.launch()


if __name__ == "__main__":
    main()