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

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

from dataclasses import dataclass, field
from datetime import datetime
from io import BytesIO
from pathlib import Path
from typing import Any, BinaryIO, Callable

import pandas as pd
from loguru import logger

from config import CHUNK_SIZE, OUTPUT_DIR
from src.analytics import build_analytics
from src.data import iter_chunks, load_excel, merge_llm, save_checkpoint
from src.llm import LLMEngine
from src.report import build_excel, save_excel

ProgressFn = Callable[[float, str], None]


@dataclass
class PipelineResult:
    enriched: pd.DataFrame = field(default_factory=pd.DataFrame)
    analytics: dict[str, Any] = field(default_factory=dict)
    executive_summary: str = ""
    excel_bytes: bytes = b""
    logs: list[str] = field(default_factory=list)


def _noop(p: float, m: str) -> None:
    pass


def run_pipeline(
    source: str | Path | BinaryIO | bytes,
    *,
    engine: LLMEngine | None = None,
    skip_llm: bool = False,
    chunk_size: int = CHUNK_SIZE,
    region_name: str = "Омская область",
    progress: ProgressFn | None = None,
) -> PipelineResult:
    report = progress or _noop
    result = PipelineResult()
    engine = engine or LLMEngine()

    report(0.05, "Чтение Excel…")
    try:
        df = load_excel(source)
    except Exception as exc:
        raise ValueError(str(exc)) from exc

    logger.info("Загружено {} строк", len(df))
    result.logs.append(f"Загружено {len(df)} строк")

    if not skip_llm:
        report(0.1, "Загрузка LLM…")
        engine.load()
        if not engine.ready:
            raise RuntimeError("LLM не готова к работе после загрузки.")
        result.logs.append(f"LLM: {engine.model_name} ✓")

    frames: list[pd.DataFrame] = []
    total = len(df)
    processed = 0

    for i, chunk in enumerate(iter_chunks(df, chunk_size), start=1):
        texts = chunk["incident_text"].astype(str).tolist()

        if skip_llm:
            llm_rows = [{"is_problem": False, "severity": 1, "core_issue": "—"} for _ in texts]
        else:

            def batch_prog(done: int, tot: int) -> None:
                frac = (processed + done) / max(total, 1)
                report(0.15 + 0.55 * frac, f"LLM {processed + done}/{total}")

            llm_rows = engine.analyze_all(texts, on_progress=batch_prog)

        enriched = merge_llm(chunk, llm_rows)
        save_checkpoint(enriched, f"chunk_{i:04d}")
        frames.append(enriched)
        processed += len(chunk)
        report(0.15 + 0.55 * (processed / max(total, 1)), f"Чанк {i} готов")

    full = pd.concat(frames, ignore_index=True)
    result.enriched = full

    report(0.75, "Агрегация по районам…")
    analytics = build_analytics(full)
    result.analytics = analytics
    result.logs.append(f"Проблем: {analytics['problems_count']} из {analytics['total_rows']}")
    result.logs.append(f"Критических (тяжесть 5): {analytics.get('critical_count', 0)}")

    report(0.88, "Формирование отчёта…")
    top10, top3 = analytics["top10"], analytics["top3"]
    critical = analytics.get("critical")
    result.excel_bytes = build_excel(top10, top3, critical)

    report(0.94, "Справка для руководства…")
    payload = analytics.get("summary_payload", [])
    if skip_llm:
        result.executive_summary = LLMEngine.template_executive_summary(payload)
    else:
        result.executive_summary = engine.executive_summary(payload)
    if region_name:
        result.executive_summary = f"Регион: {region_name}\n\n{result.executive_summary}"

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out = OUTPUT_DIR / f"report_{stamp}.xlsx"
    try:
        save_excel(top10, top3, out, critical)
        summary_path = OUTPUT_DIR / f"summary_{stamp}.txt"
        summary_path.write_text(result.executive_summary, encoding="utf-8")
        result.logs.append(f"Сохранено: {out.name}")
    except Exception as exc:
        logger.warning("Сохранение на диск: {}", exc)

    report(1.0, "Готово")
    logger.info("Пайплайн завершён")
    return result