long_articles_agentic_supply/supply/infra/shared/tools.py

"""通用工具函数 —— 只包含与业务无关的基础设施工具"""

import hashlib
import json
import random
import re
import string
from datetime import datetime, timezone, date, timedelta
from typing import Dict, List, Optional


def safe_json_parse(text: str) -> Optional[Dict | List]:
    """多层降级解析 JSON：直接解析 → 提取代码块 → 提取 JSON 对象/数组

    LLM 可能返回 ```json ... ``` 包裹的文本，或文本中夹杂 markdown 前缀/后缀。
    先尝试直接解析（最常见路径），失败后逐层降级提取。
    """
    if not text:
        return None

    # 降级 1：直接解析
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        pass

    clean = text.strip()

    # 降级 2：提取最外层 json 代码块
    m = re.search(r"```json\s*(.*?)\s*```", clean, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(1))
        except (json.JSONDecodeError, TypeError):
            pass
    else:
        m = re.search(r"```\s*(.*?)\s*```", clean, re.DOTALL)
        if m:
            try:
                return json.loads(m.group(1))
            except (json.JSONDecodeError, TypeError):
                pass

    # 降级 3：在文本中查找第一个完整 JSON 对象 { ... } 或数组 [ ... ]
    for bracket_pair in [("{}", "{", "}"), ("[]", "[", "]")]:
        opener, closer = bracket_pair[1], bracket_pair[2]
        start = clean.find(opener)
        if start == -1:
            continue
        depth = 0
        in_string = False
        escape_next = False
        for i in range(start, len(clean)):
            ch = clean[i]
            if escape_next:
                escape_next = False
                continue
            if ch == "\\":
                escape_next = True
                continue
            if ch == '"' and not escape_next:
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch == opener:
                depth += 1
            elif ch == closer:
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(clean[start : i + 1])
                    except (json.JSONDecodeError, TypeError):
                        return None
        try:
            return json.loads(clean[start:])
        except (json.JSONDecodeError, TypeError):
            pass

    return None


def str_to_md5(strings: str) -> str:
    """字符串转 MD5"""
    original_bytes = strings.encode("utf-8")
    md5_hash = hashlib.md5()
    md5_hash.update(original_bytes)
    return md5_hash.hexdigest()


def generate_task_trace_id() -> str:
    """生成任务追踪 ID，格式: Task-YYYYMMDDHHmmss-16位随机"""
    random_str = "".join(random.choices(string.ascii_lowercase + string.digits, k=16))
    return f"Task-{datetime.now().strftime('%Y%m%d%H%M%S')}-{random_str}"


def yield_batch(data: list, batch_size: int):
    """生成批次数据"""
    for i in range(0, len(data), batch_size):
        yield data[i : i + batch_size]


def timestamp_to_str(timestamp: int, string_format: str = "%Y-%m-%d %H:%M:%S") -> str:
    """时间戳转字符串"""
    dt_object = (
        datetime.utcfromtimestamp(timestamp).replace(tzinfo=timezone.utc).astimezone()
    )
    return dt_object.strftime(string_format)


def days_remaining_in_month() -> int:
    """计算本月剩余天数"""
    today = date.today()
    if today.month == 12:
        next_month = today.replace(year=today.year + 1, month=1, day=1)
    else:
        next_month = today.replace(month=today.month + 1, day=1)
    last_day_of_month = next_month - timedelta(days=1)
    return (last_day_of_month - today).days