src/llm.py - ui - Chaizee/ZenithCode_Incident-LLM-analytics

src/llm.py
285 строк10 KB
Chaizee
feat: add llm-analyzer realisation
08 июн 2026, 18:44
08 июн 2026, 18:440796e1d
from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
from typing import Any, Callable
from huggingface_hub import hf_hub_download

from loguru import logger
from pydantic import ValidationError
from llama_cpp import Llama

from config import (
    CACHE_DB,
    DEFAULT_MODEL,
    INCIDENT_PROMPT,
    INCIDENT_SYSTEM,
    LLM_MAX_TOKENS,
    LLM_TEMPERATURE,
    IncidentLLMResponse,
)


def _is_stale_cache_entry(data: dict[str, Any]) -> bool:
    return (
        not data.get("is_problem")
        and int(data.get("severity", 1)) == 1
        and str(data.get("core_issue", "")).strip() in ("не определено", "—", "")
    )


class SQLiteCache:
    def __init__(self, path: str | os.PathLike = CACHE_DB) -> None:
        self.path = str(path)
        os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
        self._init_db()

    def _init_db(self) -> None:
        with sqlite3.connect(self.path) as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS llm_cache (
                    text_hash TEXT PRIMARY KEY,
                    response_json TEXT NOT NULL,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
                """
            )

    @staticmethod
    def _hash(text: str) -> str:
        return hashlib.sha256(text.encode("utf-8")).hexdigest()

    def get(self, text: str) -> dict[str, Any] | None:
        key = self._hash(text)
        try:
            with sqlite3.connect(self.path) as conn:
                row = conn.execute(
                    "SELECT response_json FROM llm_cache WHERE text_hash = ?", (key,)
                ).fetchone()
            if row:
                data = json.loads(row[0])
                if _is_stale_cache_entry(data):
                    return None
                return data
        except Exception as exc:
            logger.warning("Cache read error: {}", exc)
        return None

    def set(self, text: str, data: dict[str, Any]) -> None:
        key = self._hash(text)
        try:
            with sqlite3.connect(self.path) as conn:
                conn.execute(
                    "INSERT OR REPLACE INTO llm_cache (text_hash, response_json) VALUES (?, ?)",
                    (key, json.dumps(data, ensure_ascii=False)),
                )
        except Exception as exc:
            logger.warning("Cache write error: {}", exc)


def parse_response(raw: str) -> IncidentLLMResponse:
    text = (raw or "").strip()
    fence = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
    if fence:
        text = fence.group(1)
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end == -1:
        return IncidentLLMResponse()

    try:
        data = json.loads(text[start : end + 1])
        return IncidentLLMResponse.model_validate(data)
    except (json.JSONDecodeError, ValidationError) as exc:
        logger.debug("Invalid LLM JSON: {} | {}", exc, raw[:120])
        return IncidentLLMResponse()


class LLMEngine:
    def __init__(
        self,
        model_name: str | None = None,
        *,
        use_cache: bool = True,
    ) -> None:
        self.model_name = model_name or os.getenv("LLM_MODEL_NAME", DEFAULT_MODEL)

        try:
            self.model_path = hf_hub_download(
                repo_id="Qwen/Qwen2.5-7B-Instruct-GGUF",
                filename="qwen2.5-7b-instruct-q3_k_m.gguf",
                local_files_only=True
            )
        except Exception:
            self.model_path = "models/qwen2.5-7b-instruct-q3_k_m.gguf"

        self.cache = SQLiteCache() if use_cache else None
        self._model: Llama | None = None
        self._few_shot: list[dict[str, Any]] = []
        self._load_few_shot_from_disk()

    def _load_few_shot_from_disk(self) -> None:
        try:
            from config import FEW_SHOT_FILE
            if FEW_SHOT_FILE.is_file():
                self._few_shot = json.loads(FEW_SHOT_FILE.read_text(encoding="utf-8"))
        except Exception:
            self._few_shot = []

    def set_few_shot(self, examples: list[dict[str, Any]]) -> None:
        self._few_shot = examples[:12]

    @property
    def ready(self) -> bool:
        return self._model is not None

    def load(self) -> None:
        if self._model is not None:
            return
        try:
            logger.info("Загрузка локальной GGUF модели на CPU: {}", self.model_path)
            if not os.path.exists(self.model_path):
                raise FileNotFoundError(f"Файл модели не найден по пути: {self.model_path}")
            
            self._model = Llama(
                model_path=self.model_path,
                n_ctx=2048,
                n_threads=4,
                verbose=False
            )
            logger.info("Модель успешно загружена на CPU!")
        except Exception as exc:
            logger.error("Ошибка загрузки Llama-CPP: {}", exc)
            raise RuntimeError(f"Не удалось загрузить модель с диска: {exc}") from exc

    def _few_shot_block(self) -> str:
        if not self._few_shot:
            return ""
        lines = ["Примеры правильного ответа:"]
        for ex in self._few_shot:
            lines.append(
                f'Текст инцидента: "{ex.get("text", "")[:200]}" → '
                f'{{"is_problem": {str(ex.get("is_problem", False)).lower()}, '
                f'"severity": {ex.get("severity", 1)}, '
                f'"core_issue": "{ex.get("core_issue", "")}"}}'
            )
        return "\n".join(lines) + "\n\n"

    def analyze_one(self, text: str) -> dict[str, Any]:
        if self.cache:
            cached = self.cache.get(text)
            if cached:
                return cached

        if not self.ready:
            return IncidentLLMResponse().model_dump()

        try:
            user_instruction = f"{self._few_shot_block()}{INCIDENT_PROMPT.format(text=(text or '')[:4000])}"

            messages = [
                {"role": "system", "content": INCIDENT_SYSTEM},
                {"role": "user", "content": user_instruction},
            ]

            response = self._model.create_chat_completion(
                messages=messages,
                max_tokens=LLM_MAX_TOKENS,
                temperature=LLM_TEMPERATURE,
            )

            response_only = response["choices"][0]["message"]["content"]
            parsed = parse_response(response_only)
            result = parsed.model_dump()
        except Exception as exc:
            logger.error("analyze_one error: {}", exc)
            return IncidentLLMResponse().model_dump()

        if self.cache and not _is_stale_cache_entry(result):
            self.cache.set(text, result)
        return result

    def analyze_batch(
        self, 
        texts: list[str], 
        on_progress: Callable[[int, int], None] | None = None
    ) -> list[dict[str, Any]]:
        if not texts:
            return []
        
        logger.info("Начало пакетной обработки на CPU (всего элементов: {})", len(texts))
        
        results = []
        total = len(texts)
        
        for i, text in enumerate(texts):
            res = self.analyze_one(text)
            results.append(res)
            
            if on_progress is not None:
                try:
                    on_progress(i + 1, total)
                except Exception as exc:
                    logger.warning("Ошибка вызова on_progress: {}", exc)
                    
        return results

    def analyze_all(
        self,
        texts: list[str],
        on_progress: Callable[[int, int], None] | None = None,
    ) -> list[dict[str, Any]]:
        return self.analyze_batch(texts, on_progress=on_progress)

    def executive_summary(self, results: list[dict[str, Any]] | list[str]) -> str:
        if not results:
            return "Нет данных для формирования отчета."
            
        if not self.ready:
            return _fallback_executive_summary(results)

        try:
            from config import SUMMARY_SYSTEM, SUMMARY_PROMPT

            formatted_data = json.dumps(results, ensure_ascii=False, indent=2)[:6000]
            user_instruction = SUMMARY_PROMPT.format(payload=formatted_data)
            messages = [
                {"role": "system", "content": SUMMARY_SYSTEM},
                {"role": "user", "content": user_instruction},
            ]
            response = self._model.create_chat_completion(
                messages=messages,
                max_tokens=LLM_MAX_TOKENS,
                temperature=LLM_TEMPERATURE,
            )
            
            return response["choices"][0]["message"]["content"].strip()

        except Exception as exc:
            logger.error("executive_summary error: {}", exc)
            return _fallback_executive_summary(results)

    @staticmethod
    def template_executive_summary(results: list[dict[str, Any]] | list[str]) -> str:
        return _fallback_executive_summary(results)


def _fallback_executive_summary(results: list[dict[str, Any]] | list[str]) -> str:
    if not results:
        return "Нет данных для формирования отчета."
    lines = ["Краткая сводка по наиболее проблемным муниципалитетам:"]
    for i, item in enumerate(results, start=1):
        if isinstance(item, dict):
            muni = item.get("municipality", "—")
            count = item.get("problem_count", "—")
            sev = item.get("mean_severity", "—")
            reason = item.get("reason", "")
            issues = item.get("top_core_issues", [])
            tail = f" Типовые сути: {', '.join(issues)}." if issues else ""
            lines.append(f"{i}. {muni}: {count} проблем, средняя тяжесть {sev}.{tail} {reason}".strip())
        else:
            lines.append(f"{i}. {item}")
    return "\n".join(lines)