how/lib/llm_cached.py

#!/usr/bin/env python3
"""
通用 LLM 缓存分析模块

根据 prompt + 模型 + 参数 进行缓存，相同输入直接返回缓存结果。
"""

import asyncio
import hashlib
import json
import re
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from agents import Agent, Runner, ModelSettings
from lib.client import get_model
from lib.config import get_cache_dir
from lib.utils import parse_json_from_text
from lib.my_trace import set_trace_smith as set_trace


# ===== 配置 =====

@dataclass
class LLMConfig:
    """LLM 配置"""
    model_name: str = "google/gemini-3-pro-preview"
    temperature: float = 0.0
    max_tokens: int = 65536

    def to_dict(self) -> Dict:
        return {
            "model_name": self.model_name,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }


@dataclass
class AnalyzeResult:
    """分析结果（包含元数据）"""
    data: Any  # 解析后的数据或原始字符串
    cache_hit: bool  # 是否命中缓存
    model_name: str  # 使用的模型
    cache_key: str  # 缓存键
    log_url: Optional[str] = None  # trace URL（仅当实际调用 LLM 时）
    retries: int = 0  # 重试次数

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "cache_hit": self.cache_hit,
            "model_name": self.model_name,
            "cache_key": self.cache_key,
            "log_url": self.log_url,
            "retries": self.retries,
        }


# 默认模型
DEFAULT_MODEL = "google/gemini-3-pro-preview"

# 预设配置
PRESETS = {
    "default": LLMConfig(
        model_name=DEFAULT_MODEL,
        temperature=0.0,
        max_tokens=65536,
    ),
    "fast": LLMConfig(
        model_name="openai/gpt-4.1-mini",
        temperature=0.0,
        max_tokens=65536,
    ),
    "balanced": LLMConfig(
        model_name="google/gemini-2.5-flash-preview-05-20",
        temperature=0.0,
        max_tokens=65536,
    ),
    "quality": LLMConfig(
        model_name="anthropic/claude-sonnet-4",
        temperature=0.0,
        max_tokens=65536,
    ),
    "best": LLMConfig(
        model_name="google/gemini-2.5-pro-preview-05-06",
        temperature=0.0,
        max_tokens=65536,
    ),
}


# ===== 缓存工具函数 =====

def _get_cache_dir(task_name: str) -> Path:
    """获取缓存目录"""
    return Path(get_cache_dir(f"llm_cached/{task_name}"))


def _generate_cache_key(
    prompt: str,
    config: LLMConfig,
) -> str:
    """生成缓存键（MD5 哈希）"""
    cache_string = f"{prompt}||{config.model_name}||{config.temperature}||{config.max_tokens}"
    return hashlib.md5(cache_string.encode('utf-8')).hexdigest()


def _sanitize_filename(text: str, max_length: int = 30) -> str:
    """将文本转换为安全的文件名"""
    sanitized = re.sub(r'[^\w\u4e00-\u9fff]', '_', text)
    sanitized = re.sub(r'_+', '_', sanitized)
    if len(sanitized) > max_length:
        sanitized = sanitized[:max_length]
    return sanitized.strip('_')


def _get_cache_filepath(
    task_name: str,
    cache_key: str,
    prompt_preview: str,
    config: LLMConfig,
) -> Path:
    """
    获取缓存文件路径

    文件名格式: {prompt_preview}_{model}_{hash[:8]}.json
    """
    cache_dir = _get_cache_dir(task_name)

    # 清理 prompt 预览
    clean_preview = _sanitize_filename(prompt_preview, max_length=40)

    # 简化模型名
    model_short = config.model_name.split('/')[-1]
    model_short = _sanitize_filename(model_short, max_length=20)

    # 哈希前8位
    hash_short = cache_key[:8]

    filename = f"{clean_preview}_{model_short}_{hash_short}.json"
    return cache_dir / filename


def _load_from_cache(
    task_name: str,
    cache_key: str,
    prompt_preview: str,
    config: LLMConfig,
) -> Optional[Dict]:
    """从缓存加载，返回 {raw: str, log_url: str}"""
    cache_file = _get_cache_filepath(task_name, cache_key, prompt_preview, config)

    # 如果文件不存在，尝试通过哈希匹配
    if not cache_file.exists():
        cache_dir = _get_cache_dir(task_name)
        if cache_dir.exists():
            hash_short = cache_key[:8]
            matching_files = list(cache_dir.glob(f"*_{hash_short}.json"))
            if matching_files:
                cache_file = matching_files[0]
            else:
                return None
        else:
            return None

    try:
        with open(cache_file, 'r', encoding='utf-8') as f:
            cached_data = json.load(f)
            return {
                "raw": cached_data['output']['raw'],
                "log_url": cached_data.get('metadata', {}).get('log_url'),
            }
    except (json.JSONDecodeError, IOError, KeyError):
        return None


def _save_to_cache(
    task_name: str,
    cache_key: str,
    prompt_preview: str,
    prompt: str,
    config: LLMConfig,
    result: str,
    log_url: Optional[str] = None,
) -> None:
    """保存到缓存（包含 log_url）"""
    cache_file = _get_cache_filepath(task_name, cache_key, prompt_preview, config)
    cache_file.parent.mkdir(parents=True, exist_ok=True)

    # 尝试解析 JSON
    parsed_result = parse_json_from_text(result)

    cache_data = {
        "input": {
            "prompt": prompt,
            "prompt_preview": prompt_preview,
            **config.to_dict(),
        },
        "output": {
            "raw": result,
            "parsed": parsed_result,
        },
        "metadata": {
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "cache_key": cache_key,
            "cache_file": str(cache_file.name),
            "log_url": log_url,
        }
    }

    try:
        with open(cache_file, 'w', encoding='utf-8') as f:
            json.dump(cache_data, f, ensure_ascii=False, indent=2)
    except IOError:
        pass


# ===== 核心 API =====

async def analyze(
    prompt: str,
    task_name: str = "default",
    config: Optional[LLMConfig] = None,
    preset: Optional[str] = None,
    force: bool = False,
    parse_json: bool = True,
    max_retries: int = 3,
    log_url: Optional[str] = None,
) -> AnalyzeResult:
    """
    通用 LLM 分析（带缓存）

    Args:
        prompt: 完整的 prompt
        task_name: 任务名称（用于缓存目录分类）
        config: LLM 配置，如果为 None 则使用 preset 或默认配置
        preset: 预设配置名称 ("default", "fast", "balanced", "quality", "best")
        force: 强制重新分析（跳过缓存），默认 False
        parse_json: 是否解析为 JSON
        max_retries: 最大重试次数（默认3次）
        log_url: 外部传入的 trace URL（如果为 None 且缓存未命中，则自动生成）

    Returns:
        AnalyzeResult 对象，包含：
        - data: 解析后的数据或原始字符串
        - cache_hit: 是否命中缓存
        - model_name: 使用的模型
        - cache_key: 缓存键
        - log_url: trace URL（仅当实际调用 LLM 时）
        - retries: 实际重试次数

    Examples:
        >>> # 使用缓存（默认）
        >>> result = await analyze("分析...", task_name="origin")

        >>> # 强制重新分析
        >>> result = await analyze("分析...", task_name="origin", force=True)

        >>> # 外部控制 trace（多个分析共享同一个 trace）
        >>> _, log_url = set_trace()
        >>> result1 = await analyze("分析1...", log_url=log_url)
        >>> result2 = await analyze("分析2...", log_url=log_url)
    """
    # 确定配置
    if config is None:
        if preset and preset in PRESETS:
            config = PRESETS[preset]
        else:
            config = PRESETS["default"]

    # 生成缓存键
    cache_key = _generate_cache_key(prompt, config)

    # prompt 预览（用于文件名）
    prompt_preview = prompt[:50].replace('\n', ' ')

    # 尝试从缓存加载（除非 force=True）
    if not force:
        cached_data = _load_from_cache(task_name, cache_key, prompt_preview, config)
        if cached_data is not None:
            cached_raw = cached_data["raw"]
            cached_log_url = cached_data.get("log_url")
            if parse_json:
                parsed = parse_json_from_text(cached_raw)
                if parsed:
                    return AnalyzeResult(
                        data=parsed,
                        cache_hit=True,
                        model_name=config.model_name,
                        cache_key=cache_key,
                        log_url=cached_log_url,  # 返回缓存时的 log_url
                        retries=0,
                    )
            else:
                return AnalyzeResult(
                    data=cached_raw,
                    cache_hit=True,
                    model_name=config.model_name,
                    cache_key=cache_key,
                    log_url=cached_log_url,  # 返回缓存时的 log_url
                    retries=0,
                )

    # 设置 trace（仅当实际调用 LLM 且未传入 log_url 时）
    if log_url is None:
        _, log_url = set_trace()

    # 创建 Agent
    agent = Agent(
        name=f"LLM-{task_name}",
        model=get_model(config.model_name),
        model_settings=ModelSettings(
            temperature=config.temperature,
            max_tokens=config.max_tokens,
        ),
        tools=[],
    )

    last_error = None
    retries = 0

    for attempt in range(max_retries):
        try:
            result = await Runner.run(agent, input=prompt)
            raw_output = result.final_output

            if parse_json:
                parsed = parse_json_from_text(raw_output)
                if parsed:
                    # 解析成功，缓存并返回
                    _save_to_cache(task_name, cache_key, prompt_preview, prompt, config, raw_output, log_url)
                    return AnalyzeResult(
                        data=parsed,
                        cache_hit=False,
                        model_name=config.model_name,
                        cache_key=cache_key,
                        log_url=log_url,
                        retries=retries,
                    )
                else:
                    # 解析失败，重试
                    retries += 1
                    last_error = f"JSON 解析失败 (尝试 {attempt + 1}/{max_retries})\n响应: {raw_output[:500]}..."
                    print(f"  ⚠️ {last_error}")
                    if attempt < max_retries - 1:
                        await asyncio.sleep(1)
            else:
                # 不需要解析 JSON
                _save_to_cache(task_name, cache_key, prompt_preview, prompt, config, raw_output, log_url)
                return AnalyzeResult(
                    data=raw_output,
                    cache_hit=False,
                    model_name=config.model_name,
                    cache_key=cache_key,
                    log_url=log_url,
                    retries=retries,
                )

        except Exception as e:
            retries += 1
            last_error = f"API 调用失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}"
            print(f"  ⚠️ {last_error}")
            if attempt < max_retries - 1:
                await asyncio.sleep(1)

    raise ValueError(f"所有 {max_retries} 次重试均失败: {last_error}")


async def analyze_batch(
    prompts: list[str],
    task_name: str = "default",
    config: Optional[LLMConfig] = None,
    preset: Optional[str] = None,
    force: bool = False,
    parse_json: bool = True,
    max_concurrent: int = 10,
    log_url: Optional[str] = None,
    progress_callback: Optional[Callable] = None,
) -> list[AnalyzeResult]:
    """
    批量 LLM 分析（带并发控制）

    Args:
        prompts: prompt 列表
        task_name: 任务名称
        config: LLM 配置
        preset: 预设配置名称
        force: 强制重新分析（跳过缓存），默认 False
        parse_json: 是否解析为 JSON
        max_concurrent: 最大并发数
        log_url: 外部传入的 trace URL（所有分析共享同一个 trace）
        progress_callback: 进度回调

    Returns:
        AnalyzeResult 列表
    """
    # 如果没有传入 log_url，生成一个共享的
    if log_url is None:
        _, log_url = set_trace()

    semaphore = asyncio.Semaphore(max_concurrent)

    async def limited_analyze(prompt: str):
        async with semaphore:
            result = await analyze(
                prompt=prompt,
                task_name=task_name,
                config=config,
                preset=preset,
                force=force,
                parse_json=parse_json,
                log_url=log_url,
            )
            if progress_callback:
                progress_callback(1)
            return result

    tasks = [limited_analyze(p) for p in prompts]
    return await asyncio.gather(*tasks)


# ===== 便捷函数 =====

async def analyze_fast(prompt: str, task_name: str = "default", **kwargs) -> AnalyzeResult:
    """快速分析（使用 fast 预设）"""
    return await analyze(prompt, task_name=task_name, preset="fast", **kwargs)


async def analyze_quality(prompt: str, task_name: str = "default", **kwargs) -> AnalyzeResult:
    """高质量分析（使用 quality 预设）"""
    return await analyze(prompt, task_name=task_name, preset="quality", **kwargs)


# ===== 测试 =====

if __name__ == "__main__":
    async def main():
        prompt = """
分析"猫咪"和"宠物"的关系，输出 JSON：
```json
{
  "关系": "...",
  "说明": "..."
}
```
"""
        print("测试 1: 基本用法（自动生成 trace）")
        result = await analyze(prompt, task_name="test", preset="fast")
        print(f"  cache_hit: {result.cache_hit}")
        print(f"  model: {result.model_name}")
        print(f"  log_url: {result.log_url}")
        print(f"  data: {result.data}")

        print("\n测试 2: 缓存命中")
        result = await analyze(prompt, task_name="test", preset="fast")
        print(f"  cache_hit: {result.cache_hit}")
        print(f"  log_url: {result.log_url}")  # 应该是 None

        print("\n测试 3: 强制重新分析 (force=True)")
        _, shared_log_url = set_trace()
        print(f"  共享 trace: {shared_log_url}")
        result = await analyze(
            "输出 JSON: {\"test\": 123}",
            task_name="test",
            log_url=shared_log_url,
            force=True,
        )
        print(f"  cache_hit: {result.cache_hit}")
        print(f"  log_url: {result.log_url}")

    asyncio.run(main())