src/cleaner.py - master - sarus/audio-agent

src/cleaner.py
204 строки9 KB
otcheskiy
key conclusion
15 июн 2026, 23:01
15 июн 2026, 23:012fe452a
"""LLM-очистка сырого транскрипта через локальный OpenAI-compatible endpoint.

[experimental] Не входит в основной пайплайн. Запускается только флагом --clean.
"""

import os
from pathlib import Path

import httpx
from dotenv import load_dotenv

from config import (
    LLM_API_BASE_URL,
    LLM_API_KEY,
    LLM_CHUNK_MAX,
    LLM_CHUNK_MIN,
    LLM_MODEL,
)

load_dotenv(Path(__file__).resolve().parent.parent / ".env")

SYSTEM_PROMPT = """Ты помощник по минимальной правке судебной расшифровки ASR.
Цель: не сделать текст красивым, а минимально исправить очевидные ошибки распознавания речи.

Правила:
1. Исправляй только отдельные очевидные ASR-ошибки (не более 1–2 слов подряд).
2. Если исправление меняет больше 1–2 слов подряд — не делай его.
3. Если исправление требует догадки о смысле всей фразы — не исправляй.
4. [неразборчиво] — только для явно бессмысленного фрагмента, не для любой неясности.
5. Сохраняй порядок текста, объём и все реплики. Не сокращай.
6. Возвращай только текст без комментариев и пояснений.

Запрещено:
- переписывать предложения;
- улучшать стиль;
- делать юридическую реконструкцию;
- менять роли участников;
- менять ФИО, даты, суммы, номера дел, названия организаций.

Разрешено (примеры точечных исправлений):
- конспертиза → экспертиза
- обыскании → взыскании
- основательного обогащения → неосновательного обогащения
- целевые износы → целевые взносы
- ПЗ → ФЗ
- общедалевая → общая долевая
- межеваренной → межевания
- правоходарство → ходатайство

Нельзя (оставляй без изменений):
- «заявление законодательства ответчика» → «заявление представителя ответчика»
- «юристов не привыкал» → «юристов не вызывал»
Такие случаи не исправляй, если для этого нужна реконструкция смысла фразы."""

GLOSSARY = (
    "ходатайство, истец, ответчик, взыскание, неосновательное обогащение, "
    "производство по делу, экспертиза, судебные расходы, ГПК РФ, ГК РФ, "
    "Федеральный закон 217, СНТ, ТСН, ЕГРН, межевой план, целевые взносы, "
    "членские взносы, общая долевая собственность, сервитут"
)


class CleanupError(Exception):
    """Понятная ошибка LLM cleanup без traceback."""


def _llm_settings() -> tuple[str, str, str]:
    api_base = os.getenv("LLM_API_BASE_URL", LLM_API_BASE_URL)
    api_key = os.getenv("LLM_API_KEY", LLM_API_KEY)
    model = os.getenv("LLM_MODEL", LLM_MODEL)
    return api_base.rstrip("/"), api_key, model


def check_local_llm_available() -> None:
    """Проверяет доступность локального LLM-сервера."""
    api_base, api_key, _ = _llm_settings()
    try:
        response = httpx.get(
            f"{api_base}/models",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=httpx.Timeout(5.0),
        )
        if response.status_code != 200:
            raise CleanupError(
                f"Локальный LLM недоступен ({api_base}): ответ {response.status_code}. "
                "Запустите LM Studio (http://localhost:1234/v1) или Ollama "
                "(http://localhost:11434/v1)."
            )
    except httpx.ConnectError:
        raise CleanupError(
            f"Локальный LLM сервер не запущен ({api_base}). "
            "Запустите LM Studio или Ollama и проверьте LLM_API_BASE_URL в .env."
        ) from None
    except httpx.TimeoutException:
        raise CleanupError(
            f"Таймаут подключения к локальному LLM ({api_base}). "
            "Проверьте, что сервер запущен и отвечает."
        ) from None
    except CleanupError:
        raise
    except httpx.HTTPError as exc:
        raise CleanupError(
            f"Не удалось подключиться к локальному LLM ({api_base}): {exc}"
        ) from None


def _split_into_chunks(text: str, min_size: int, max_size: int) -> list[str]:
    if len(text) <= max_size:
        return [text]

    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + max_size, len(text))
        if end < len(text):
            search_from = min(start + min_size, end)
            break_at = -1
            for separator in (". ", ".\n", "? ", "! ", ".\u00a0", "\n"):
                pos = text.rfind(separator, search_from, end)
                if pos > break_at:
                    break_at = pos + len(separator)
            if break_at > start:
                end = break_at
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start = end
    return chunks


def _call_llm(chunk: str, chunk_index: int, total_chunks: int) -> str:
    api_base, api_key, model = _llm_settings()

    user_prompt = (
        f"Глоссарий: {GLOSSARY}\n\n"
        f"Фрагмент {chunk_index + 1} из {total_chunks}.\n"
        f"Минимально исправь только очевидные ASR-ошибки (не более 1–2 слов подряд). "
        f"Не переписывай предложения.\n\n"
        f"{chunk}"
    )

    try:
        response = httpx.post(
            f"{api_base}/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": model,
                "temperature": 0,
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_prompt},
                ],
            },
            timeout=httpx.Timeout(300.0),
        )
        response.raise_for_status()
    except httpx.ConnectError:
        raise CleanupError(
            f"Потеряно соединение с локальным LLM ({api_base}) на чанке "
            f"{chunk_index + 1}/{total_chunks}."
        ) from None
    except httpx.HTTPStatusError as exc:
        raise CleanupError(
            f"Локальный LLM вернул ошибку {exc.response.status_code} на чанке "
            f"{chunk_index + 1}/{total_chunks}."
        ) from None
    except httpx.HTTPError as exc:
        raise CleanupError(
            f"Ошибка запроса к локальному LLM на чанке "
            f"{chunk_index + 1}/{total_chunks}: {exc}"
        ) from None

    data = response.json()
    return data["choices"][0]["message"]["content"].strip()


def clean_transcript(text: str) -> str:
    """Очищает сырой текст транскрипта через локальную LLM по чанкам.

    Args:
        text: Сырой текст от Whisper (без заголовка метаданных).

    Returns:
        Очищенный текст.

    Raises:
        CleanupError: если локальный LLM недоступен или запрос не удался.
    """
    if not text.strip():
        return text

    check_local_llm_available()

    chunks = _split_into_chunks(text, LLM_CHUNK_MIN, LLM_CHUNK_MAX)
    cleaned_chunks: list[str] = []

    for index, chunk in enumerate(chunks):
        print(f"LLM cleanup: чанк {index + 1}/{len(chunks)} ({len(chunk)} симв.)")
        cleaned_chunks.append(_call_llm(chunk, index, len(chunks)))

    return "".join(cleaned_chunks)