zk-data-agent/skills/online-mining-v2/scripts/online_mining_common.py

from __future__ import annotations

"""online-mining-v2 可迁移脚本的共享逻辑。

脚本只做只读 ELK 查询和轻量格式转换，输出统一 JSON，方便被不同 Agent 调用。
"""

import argparse
import importlib.util
import json
import re
import sys
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any


class OnlineMiningError(ValueError):
    """线上挖掘参数或数据错误。"""


SCRIPT_DIR = Path(__file__).resolve().parent
SKILL_DIR = SCRIPT_DIR.parent
REPO_ROOT = SKILL_DIR.parent.parent
ELK_FETCH_DIR = SKILL_DIR.parent / "elk-fetch"


SOURCE_CONFIGS: dict[str, dict[str, Any]] = {
    "main": {
        "profile": "default",
        "index": "arch-flat-nlp-log-f-*",
        "id_field": "request_id",
        "time_field": "timestamp",
    },
    "pre_processing": {
        "profile": "default",
        "index": "pre-processing-info*",
        "id_field": "requestId",
        "time_field": "timestamp",
    },
}


def load_json_payload(argv: list[str] | None = None) -> dict[str, Any]:
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument("--input", "-i", help="JSON 参数文件；不传则从 stdin 读取")
    args = parser.parse_args(argv)
    text = Path(args.input).expanduser().read_text(encoding="utf-8") if args.input else sys.stdin.read()
    if not text.strip():
        raise OnlineMiningError("input JSON is required")
    payload = json.loads(text)
    if not isinstance(payload, dict):
        raise OnlineMiningError("input JSON must be an object")
    return payload


def emit_success(payload: dict[str, Any]) -> None:
    print(json.dumps({"ok": True, **payload}, ensure_ascii=False, separators=(",", ":")))


def emit_error(exc: BaseException) -> None:
    print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False, separators=(",", ":")))


def load_elk_query_module():
    path = ELK_FETCH_DIR / "elk_query.py"
    if not path.exists():
        raise OnlineMiningError(f"elk-fetch script not found: {path}")
    spec = importlib.util.spec_from_file_location("elk_query", path)
    if spec is None or spec.loader is None:
        raise OnlineMiningError(f"cannot load elk_query.py from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def build_client(source: str):
    elk = load_elk_query_module()
    config = source_config(source)
    profiles = elk.load_profiles()
    return elk.build_client(profiles, str(config["profile"]))


def source_config(source: str) -> dict[str, Any]:
    if source not in SOURCE_CONFIGS:
        raise OnlineMiningError(f"unknown source: {source}")
    return SOURCE_CONFIGS[source]


def date_range_ms(date_str: str | None) -> tuple[int, int]:
    if date_str:
        date_obj = datetime.strptime(str(date_str), "%Y%m%d")
        start_ms = int(date_obj.timestamp() * 1000)
        end_ms = int((date_obj + timedelta(days=1)).timestamp() * 1000) - 1
    else:
        end = datetime.now()
        start_ms = int((end - timedelta(hours=48)).timestamp() * 1000)
        end_ms = int(end.timestamp() * 1000)
    return start_ms, end_ms


def deep_parse(value: Any) -> Any:
    if isinstance(value, str):
        stripped = value.strip()
        if (stripped.startswith("{") and stripped.endswith("}")) or (
            stripped.startswith("[") and stripped.endswith("]")
        ):
            try:
                return deep_parse(json.loads(stripped))
            except Exception:
                return value
        return value
    if isinstance(value, dict):
        return {key: deep_parse(item) for key, item in value.items()}
    if isinstance(value, list):
        return [deep_parse(item) for item in value]
    return value


def get_path(obj: Any, dotted: str, default: Any = None) -> Any:
    cur = obj
    for part in dotted.split("."):
        if not part:
            continue
        if isinstance(cur, dict):
            cur = cur.get(part)
        elif isinstance(cur, list):
            try:
                cur = cur[int(part)]
            except (ValueError, IndexError):
                return default
        else:
            return default
    return default if cur is None else cur


def compact_text(value: Any, limit: int = 500) -> str:
    if value is None:
        return ""
    text = json.dumps(value, ensure_ascii=False) if isinstance(value, (dict, list)) else str(value)
    text = re.sub(r"\s+", " ", text).strip()
    return text if len(text) <= limit else text[: limit - 3] + "..."


def as_list(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    return [value]


def string_values(value: Any) -> list[str]:
    return [str(item) for item in as_list(value) if str(item).strip()]


def match_text_filters(text: str, contains: Any = None, regex: str | None = None) -> bool:
    for item in string_values(contains):
        if item not in text:
            return False
    if regex and not re.search(regex, text):
        return False
    return True


def unique_strings(items: list[Any]) -> list[str]:
    seen: set[str] = set()
    result: list[str] = []
    for item in items:
        text = str(item).strip()
        if not text or text in seen:
            continue
        seen.add(text)
        result.append(text)
    return result


def extract_case(source: str, doc: dict[str, Any]) -> dict[str, Any]:
    parsed = deep_parse(doc)
    if source == "main":
        return extract_main_case(parsed)
    if source == "pre_processing":
        return extract_pre_processing_case(parsed)
    raise OnlineMiningError(f"unknown source: {source}")


def extract_main_case(doc: dict[str, Any]) -> dict[str, Any]:
    device = doc.get("device")
    return {
        "request_id": doc.get("request_id"),
        "timestamp": doc.get("timestamp"),
        "query": doc.get("query"),
        "session_id": doc.get("session_id"),
        "device_id": doc.get("device_id"),
        "device": device,
        "domain": doc.get("domain"),
        "func": doc.get("func") or doc.get("func_name"),
        "text": doc.get("to_speak") or doc.get("text") or doc.get("display_text"),
        "intention": doc.get("intention"),
        "llm_agent_info": get_path(doc, "intention.intent_arbitrator_info.llm_agent_info"),
        "score_domains": get_path(doc, "intention.intent_arbitrator_info.score_domains"),
        "hit_rules": get_path(doc, "intention.intent_arbitrator_info.hit_rules"),
        "raw": doc,
    }


def extract_pre_processing_case(doc: dict[str, Any]) -> dict[str, Any]:
    body = doc.get("responseBoby") or {}
    core = get_path(body, "nodes.0.core", {})
    dispatch = get_path(core, "dispatch_large_model_result", {}) or {}
    dispatch_input = get_path(dispatch, "dispatchLargeModelInput", {}) or {}
    planning_debug = get_path(dispatch, "planningDebugInfo", {}) or {}
    properties0 = get_path(body, "nodes.0.properties.0", {}) or {}
    debug = properties0.get("debug") if isinstance(properties0, dict) else {}
    excellent_domains = get_path(core, "excellent_domains_result", []) or []
    domain_infos = get_path(debug, "domain_infos", []) if isinstance(debug, dict) else []
    return {
        "request_id": doc.get("requestId"),
        "timestamp": doc.get("timestamp"),
        "query": get_path(core, "query.query"),
        "planning_result": planning_debug.get("planningOriginalResult") if isinstance(planning_debug, dict) else None,
        "code": dispatch.get("code") if isinstance(dispatch, dict) else None,
        "prompt_model": dispatch_input.get("promptModel") if isinstance(dispatch_input, dict) else None,
        "hit_rules": dispatch_input.get("hitRuleList") if isinstance(dispatch_input, dict) else None,
        "excellent_domains": excellent_domains,
        "candidate_domains": unique_strings(
            [get_path(item, "domain") for item in as_list(excellent_domains)]
            + [get_path(item, "domain") for item in as_list(domain_infos)]
        ),
        "domain_infos": domain_infos,
        "agent_llm_intent": properties0.get("agentLLMIntent") if isinstance(properties0, dict) else None,
        "prompt_string": get_path(debug, "promptString") if isinstance(debug, dict) else None,
        "planning_prompt": get_path(debug, "planningPrompt") if isinstance(debug, dict) else None,
        "raw": doc,
    }


def case_matches(case: dict[str, Any], filters: dict[str, Any]) -> bool:
    query = str(case.get("query") or "")
    if not match_text_filters(query, filters.get("query_contains"), filters.get("query_regex")):
        return False
    for key in ("domain", "func", "planning_result", "code"):
        expected = string_values(filters.get(key))
        if expected and str(case.get(key) or "") not in expected:
            return False
        contains = filters.get(f"{key}_contains")
        if contains and not match_text_filters(str(case.get(key) or ""), contains):
            return False
    if filters.get("prompt_contains") and not match_text_filters(
        str(case.get("prompt_model") or case.get("prompt_string") or case.get("planning_prompt") or ""),
        filters.get("prompt_contains"),
    ):
        return False
    candidate_domains = set(string_values(case.get("candidate_domains")))
    expected_candidates = set(string_values(filters.get("candidate_domain")))
    if expected_candidates and not expected_candidates.intersection(candidate_domains):
        return False
    return True


def search_docs(
    *,
    source: str,
    date: str | None,
    size: int,
    query_filters: list[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    config = source_config(source)
    client = build_client(source)
    start_ms, end_ms = date_range_ms(date)
    filters = [{"range": {str(config["time_field"]): {"gte": start_ms, "lte": end_ms}}}]
    filters.extend(deepcopy(query_filters or []))
    body = {
        "query": {"bool": {"filter": filters}},
        "size": size,
        "sort": [{str(config["time_field"]): {"order": "desc"}}],
    }
    response = client.search(index=str(config["index"]), body=body)
    return [deep_parse(hit.get("_source", {})) for hit in response.get("hits", {}).get("hits", [])]


def fetch_one_by_request_id(source: str, request_id: str, date: str | None = None) -> dict[str, Any] | None:
    config = source_config(source)
    docs = search_docs(
        source=source,
        date=date,
        size=1,
        query_filters=[{"term": {str(config["id_field"]): request_id}}],
    )
    if not docs:
        # 部分 request id 会带后缀，兜底用 wildcard。
        docs = search_docs(
            source=source,
            date=date,
            size=1,
            query_filters=[{"wildcard": {str(config["id_field"]): {"value": f"{request_id}*"}}}],
        )
    return docs[0] if docs else None


def write_optional_jsonl(path: str | None, rows: list[dict[str, Any]]) -> str | None:
    if not path:
        return None
    output = resolve_portable_path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
    return str(output)


def resolve_portable_path(path: str) -> Path:
    raw = Path(path).expanduser()
    if raw.is_absolute():
        return raw
    return (REPO_ROOT / raw).resolve()