pseudo_agent.py

import re

from llm import get_llm_answer, get_llm_completion
from retrieval_impl import RETRIEVAL
from vocabulary.parse_vocabulary import VOCABULARY_MANAGER

MAX_ITERATIONS = 3
DEFAULT_YEAR_FROM = 1918
DEFAULT_YEAR_TO = 2026
DEFAULT_TOP_K = 20
EXPANDED_TOP_K = 30

NEED_MORE_RE = re.compile(r"^NEED_MORE:\s*(.+)$", re.MULTILINE)
MULTI_RESULT_HINTS = ("факт", "пример", "несколько", "пять", "четыре", "три", "подробн")


def _wants_more_results(query: str) -> bool:
    q = query.lower()
    return any(hint in q for hint in MULTI_RESULT_HINTS)


def _merge_indices(existing: list, new: list) -> list:
    seen = set(existing)
    merged = list(existing)
    for idx in new:
        if idx not in seen:
            merged.append(idx)
            seen.add(idx)
    return merged


def _format_context(indices: list) -> str:
    if not indices:
        return ""

    df = RETRIEVAL.paragraphs_df.iloc[indices]
    chunks = []
    for _, row in df.iterrows():
        chunks.append(
            f"""Название: {row.summary}
Период: {row.start_year}-{row.end_year}

{row.text}"""
        )
    return "\n\n---\n\n".join(chunks)


def _find_title_boost_query(query: str) -> str | None:
    """Если запрос пересекается с названием раздела — вернуть его для приоритетного поиска."""
    q = query.lower()
    best_match = None
    best_len = 0

    for summary in RETRIEVAL.paragraphs_df["summary"].unique():
        summary_lower = summary.lower()
        if summary_lower in q or q in summary_lower:
            if len(summary_lower) > best_len:
                best_match = summary
                best_len = len(summary_lower)
                continue

        words = [w for w in re.findall(r"\w+", summary_lower) if len(w) > 3]
        if words and all(w in q for w in words[:2]):
            if len(summary_lower) > best_len:
                best_match = summary
                best_len = len(summary_lower)

    return best_match


def _search(
    query: str,
    *,
    year_from: int,
    year_to: int,
    top_k: int,
) -> tuple[list, str]:
    _, _, indices, status = RETRIEVAL.perform_search(
        query=query,
        top_k=top_k,
        year_from=year_from,
        year_to=year_to,
    )
    return list(indices), status


def _build_judge_prompt(user_query: str, context: str) -> str:
    return f"""Ты оцениваешь, достаточно ли архивных материалов для ответа на вопрос пользователя.

ВОПРОС:
{user_query}

НАЙДЕННЫЕ МАТЕРИАЛЫ:
{context if context else "(пусто — материалов нет)"}

Правила:
- Смотри ТОЛЬКО на материалы, не на свои знания.
- Корпус — только про общественный транспорт Рязани.
- Если материалов хватает для ответа — ответь ровно одной строкой:
SUFFICIENT
- Если материалов мало, они не по теме или их нет — ответь ровно одной строкой:
NEED_MORE: <короткий поисковый запрос на русском>
- Никакого другого текста.

Примеры:
Вопрос: "история маршрутов с буквами в номерах"
Материалы: [текст про буквы в номерах маршрутов]
→ SUFFICIENT

Вопрос: "пять интересных фактов об истории транспорта"
Материалы: [один короткий абзац]
→ NEED_MORE: интересные факты история транспорта Рязань"""


def _parse_judge_response(text: str) -> tuple[str, str | None]:
    text = text.strip()

    for line in text.splitlines():
        line = line.strip()
        if line == "SUFFICIENT":
            return "sufficient", None
        match = NEED_MORE_RE.match(line)
        if match:
            return "need_more", match.group(1).strip()

    if "SUFFICIENT" in text:
        return "sufficient", None

    match = NEED_MORE_RE.search(text)
    if match:
        return "need_more", match.group(1).strip()

    return "sufficient", None


class PseudoAgent:
    def __init__(
        self,
        *,
        year_from: int = DEFAULT_YEAR_FROM,
        year_to: int = DEFAULT_YEAR_TO,
        max_iterations: int = MAX_ITERATIONS,
    ):
        self.year_from = year_from
        self.year_to = year_to
        self.max_iterations = max_iterations

    def run(self, query: str):
        """Генератор: прогресс поиска/судьи и финальный ответ."""
        query = query.strip()
        if not query:
            yield "Введите вопрос"
            return

        parts = ["**Псевдо-агент:** поиск материалов...\n"]
        yield "\n\n".join(parts)

        all_indices: list = []
        search_query = query
        top_k = EXPANDED_TOP_K if _wants_more_results(query) else DEFAULT_TOP_K

        title_boost = _find_title_boost_query(query)
        if title_boost:
            boost_indices, boost_status = _search(
                title_boost,
                year_from=self.year_from,
                year_to=self.year_to,
                top_k=top_k,
            )
            all_indices = _merge_indices(all_indices, boost_indices)
            parts.append(f"**Совпадение с разделом:** `{title_boost}` — {boost_status}")
            yield "\n\n".join(parts)

        for iteration in range(1, self.max_iterations + 1):
            parts.append(f"### Итерация {iteration}")

            indices, status = _search(
                search_query,
                year_from=self.year_from,
                year_to=self.year_to,
                top_k=top_k,
            )
            all_indices = _merge_indices(all_indices, indices)

            parts.append(f"**Поиск:** `{search_query}` (top_k={top_k})")
            parts.append(f"**Результат:** {status}, всего разделов: {len(all_indices)}")
            yield "\n\n".join(parts)

            context = _format_context(all_indices)

            if not context:
                if iteration < self.max_iterations:
                    search_query = query
                    top_k = EXPANDED_TOP_K
                    parts.append("**Судья:** (пропущен — материалов нет, повторяю поиск)")
                    yield "\n\n".join(parts)
                    continue
                break

            judge_prompt = _build_judge_prompt(query, context)
            judge_raw = get_llm_completion(judge_prompt, max_tokens=100, temperature=0.0)
            decision, next_query = _parse_judge_response(judge_raw)

            parts.append(f"**Судья:** `{judge_raw.strip()}`")
            yield "\n\n".join(parts)

            if decision == "sufficient" or iteration == self.max_iterations:
                break

            if next_query:
                search_query = next_query
            else:
                search_query = query

        context = _format_context(all_indices)
        if not context:
            yield "\n\n".join(parts) + "\n\n## Ответ\n\nВ корпусе не нашёл материалов по этому запросу. Попробуйте уточнить вопрос."
            return

        parts.append("## Ответ\n\n")
        yield "\n\n".join(parts)

        answer_prompt = VOCABULARY_MANAGER.wrap_prompt(context, query)
        answer_text = ""
        for chunk in get_llm_answer(answer_prompt):
            answer_text += chunk
            yield "\n\n".join(parts[:-1]) + parts[-1] + answer_text


PSEUDO_AGENT = PseudoAgent()