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

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

from io import BytesIO
from pathlib import Path
from typing import BinaryIO, Iterator

import pandas as pd
from loguru import logger

from config import (
    CATEGORY_COLS,
    CHECKPOINT_DIR,
    CHUNK_SIZE,
    EXCEL_COLUMN_MAP,
    EXCEL_HEADERS,
    TEXT_SOURCE_COLS,
)


class DataError(ValueError):
    pass


def _coalesce_incident_text(df: pd.DataFrame) -> pd.DataFrame:
    parts = [df[c].astype(str).str.strip() for c in TEXT_SOURCE_COLS if c in df.columns]
    if not parts:
        df["incident_text"] = ""
        return df

    combined = parts[0]
    for part in parts[1:]:
        combined = combined.where(combined.str.len() > 5, part)

    bad = {"", "nan", "none", "'", '"', "''"}
    combined = combined.where(~combined.str.lower().isin(bad), "")
    df["incident_text"] = combined
    df = df.drop(columns=[c for c in TEXT_SOURCE_COLS if c in df.columns])
    return df


def load_excel(source: str | Path | BinaryIO | bytes, *, nrows: int | None = None) -> pd.DataFrame:
    try:
        buffer = BytesIO(source) if isinstance(source, bytes) else source
        df = pd.read_excel(
            buffer,
            usecols=list(EXCEL_HEADERS),
            dtype=str,
            nrows=nrows,
            engine="openpyxl",
        )
    except Exception as exc:
        raise DataError(f"Ошибка чтения Excel: {exc}") from exc

    missing = [h for h in EXCEL_HEADERS if h not in df.columns and h != "Номер инцидента"]
    if missing:
        raise DataError(f"В Excel не найдены столбцы: {', '.join(missing)}")
    df = df.rename(columns={k: v for k, v in EXCEL_COLUMN_MAP.items() if k in df.columns})
    if "incident_id" not in df.columns:
        df["incident_id"] = ""
    df = _coalesce_incident_text(df)
    df = df[df["incident_text"].str.len() > 0]
    return optimize_dtypes(df.dropna(how="all").reset_index(drop=True))


def optimize_dtypes(df: pd.DataFrame) -> pd.DataFrame:
    try:
        for col in CATEGORY_COLS:
            s = df[col].astype(str).str.strip().replace({"": "Не указано", "nan": "Не указано"})
            df[col] = s.astype("category")
        df["incident_text"] = df["incident_text"].astype(str).str.strip()
        if "incident_id" in df.columns:
            df["incident_id"] = df["incident_id"].astype(str).str.strip()
    except Exception as exc:
        logger.warning("Частичная ошибка оптимизации типов: {}", exc)
    return df


def iter_chunks(df: pd.DataFrame, size: int = CHUNK_SIZE) -> Iterator[pd.DataFrame]:
    for start in range(0, len(df), size):
        yield df.iloc[start : start + size].copy()


def save_checkpoint(df: pd.DataFrame, name: str) -> Path:
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    path = CHECKPOINT_DIR / f"{name}.parquet"
    df.to_parquet(path, index=False)
    return path


def merge_llm(df: pd.DataFrame, results: list[dict]) -> pd.DataFrame:
    out = df.copy()
    n = min(len(df), len(results))
    problems = [bool(results[i].get("is_problem", False)) for i in range(n)]
    severities = [int(results[i].get("severity", 1)) for i in range(n)]
    issues = [str(results[i].get("core_issue", "не определено")) for i in range(n)]
    pad = len(df) - n
    problems.extend([False] * pad)
    severities.extend([1] * pad)
    issues.extend(["ошибка анализа"] * pad)
    out["is_problem"] = problems
    out["severity"] = severities
    out["core_issue"] = issues
    return out