rag/core/llm_utils.py

"""
LLM 调用工具函数

统一封装 LLM 调用模式，减少代码重复。
"""

import json
import re
import logging
from typing import List, Optional, Union, Iterator, Callable

logger = logging.getLogger(__name__)

# MiMo 推理模型关键词（用于识别需要 thinking 参数的模型）
_MIMO_MODEL_KEYWORDS = ('mimo',)

def _is_mimo_model(model_name: str) -> bool:
    """判断是否为小米 MiMo 模型（支持 thinking 参数）"""
    if not model_name:
        return False
    return any(kw in model_name.lower() for kw in _MIMO_MODEL_KEYWORDS)

def _inject_mimo_thinking(kwargs: dict, model: str, disable_thinking: bool = True) -> dict:
    """为 MiMo 模型注入 thinking 参数。返回更新后的 kwargs。"""
    if not _is_mimo_model(model):
        return kwargs
    extra = dict(kwargs.get('extra_body') or {})
    extra['thinking'] = {'type': 'disabled' if disable_thinking else 'enabled'}
    kwargs['extra_body'] = extra
    return kwargs


def call_llm(
    client,
    prompt: str,
    model: str,
    temperature: float = 0.3,
    max_tokens: int = 1000,
    messages: List[dict] = None,
    stream: bool = False,
    **kwargs
) -> Union[str, Iterator, None]:
    """
    统一的 LLM 调用封装

    Args:
        client: OpenAI 客户端实例
        prompt: 用户提示（当 messages 为 None 时使用）
        model: 模型名称
        temperature: 温度参数 (0-1)
        max_tokens: 最大 token 数
        messages: 完整消息列表（优先于 prompt）
        stream: 是否启用流式输出
        **kwargs: 其他参数（如 response_format, tools 等）

    Returns:
        流式模式返回 stream 迭代器，否则返回响应内容字符串
        调用失败返回 None

    Example:
        # 简单调用
        response = call_llm(client, "你好", model, temperature=0.3)

        # 使用消息列表
        response = call_llm(client, "", model, messages=[
            {"role": "system", "content": "你是助手"},
            {"role": "user", "content": "你好"}
        ])

        # 流式调用
        for chunk in call_llm(client, prompt, model, stream=True):
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content
    """
    if messages is None:
        messages = [{"role": "user", "content": prompt}]

    # MiMo 模型自动注入 thinking 参数
    try:
        from config import LLM_DISABLE_THINKING
    except ImportError:
        LLM_DISABLE_THINKING = True
    kwargs = _inject_mimo_thinking(kwargs, model, disable_thinking=LLM_DISABLE_THINKING)

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=stream,
            **kwargs
        )

        if stream:
            return response

        content = response.choices[0].message.content

        # 推理模型兼容（mimo-v2.5 等）：
        # 推理模型思考链消耗大量 token（~1000），max_tokens 不足时 content 为空，
        # 全部输出进入 reasoning_content。此处从思考链中提取有效内容。
        if not content or not content.strip():
            reasoning = getattr(response.choices[0].message, 'reasoning_content', None)
            if reasoning and reasoning.strip():
                # 先去掉 <think>...</think> 标签
                cleaned = re.sub(r'', '', reasoning, flags=re.DOTALL).strip()
                if cleaned:
                    logger.info("LLM: content为空，从reasoning_content提取内容")
                    # 尝试提取 JSON 对象（兼容结构化响应场景）
                    json_match = re.search(r'\{[\s\S]*\}', cleaned)
                    if json_match:
                        try:
                            json.loads(json_match.group())
                            return json_match.group().strip()
                        except (json.JSONDecodeError, ValueError):
                            pass
                    # 尝试提取 JSON 数组
                    bracket_match = re.search(r'\[[\s\S]*\]', cleaned)
                    if bracket_match:
                        try:
                            json.loads(bracket_match.group())
                            return bracket_match.group().strip()
                        except (json.JSONDecodeError, ValueError):
                            pass
                    # 纯文本响应：直接返回清理后的内容
                    return cleaned
            logger.warning("LLM 返回空 content 且 reasoning_content 也无法提取（可能需要增大 max_tokens）")
            return None

        return content.strip()
    except Exception as e:
        logger.warning(f"LLM 调用失败: {e}")
        return None


def call_llm_stream(
    client,
    prompt: str,
    model: str,
    temperature: float = 0.3,
    max_tokens: int = 3000,
    messages: List[dict] = None,
    error_prefix: str = "[错误]",
    **kwargs
) -> Iterator[str]:
    """
    流式 LLM 调用（生成器封装）

    自动处理流式响应，逐块 yield 文本内容。
    兼容推理模型（mimo-v2.5 等）：当 content 为空时回退到 reasoning_content。

    Args:
        client: OpenAI 客户端实例
        prompt: 用户提示
        model: 模型名称
        temperature: 温度参数
        max_tokens: 最大 token 数（推理模型需留足思考链预算）
        messages: 完整消息列表
        error_prefix: 错误时的前缀
        **kwargs: 其他参数

    Yields:
        响应文本片段

    Example:
        for text in call_llm_stream(client, "讲个故事", model):
            print(text, end="", flush=True)
    """
    if messages is None:
        messages = [{"role": "user", "content": prompt}]

    # MiMo 模型自动注入 thinking 参数
    try:
        from config import LLM_DISABLE_THINKING
    except ImportError:
        LLM_DISABLE_THINKING = True
    kwargs = _inject_mimo_thinking(kwargs, model, disable_thinking=LLM_DISABLE_THINKING)

    try:
        stream = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
            **kwargs
        )

        content_yielded = False
        reasoning_buffer = []

        for chunk in stream:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta

            # 正常 content 输出
            if hasattr(delta, 'content') and delta.content:
                content_yielded = True
                yield delta.content
                continue

            # 推理模型：reasoning_content（思考链）
            rc = getattr(delta, 'reasoning_content', None)
            if rc:
                reasoning_buffer.append(rc)

        # 回退：content 为空但 reasoning_content 有内容（推理模型 token 不足时）
        if not content_yielded and reasoning_buffer:
            reasoning_text = ''.join(reasoning_buffer)
            # 去掉 <think>...</think> 标签
            cleaned = re.sub(r'', '', reasoning_text, flags=re.DOTALL).strip()
            if cleaned:
                logger.info("流式 LLM: content为空，从reasoning_content提取内容")
                yield cleaned

    except Exception as e:
        logger.error(f"LLM 流式调用失败: {e}")
        yield f"{error_prefix} 调用大模型失败: {str(e)}"


def call_llm_with_retry(
    client,
    prompt: str,
    model: str,
    max_retries: int = 3,
    retry_delay: float = 1.0,
    **kwargs
) -> Optional[str]:
    """
    带重试的 LLM 调用

    Args:
        client: OpenAI 客户端实例
        prompt: 用户提示
        model: 模型名称
        max_retries: 最大重试次数
        retry_delay: 重试间隔（秒）
        **kwargs: 传递给 call_llm 的其他参数

    Returns:
        响应内容，全部失败返回 None
    """
    import time

    for attempt in range(max_retries):
        result = call_llm(client, prompt, model, **kwargs)
        if result is not None:
            return result

        if attempt < max_retries - 1:
            logger.debug(f"LLM 调用重试 {attempt + 1}/{max_retries}")
            time.sleep(retry_delay)

    logger.warning(f"LLM 调用失败，已重试 {max_retries} 次")
    return None


def parse_json_from_response(content: str) -> Optional[dict]:
    """
    从 LLM 响应中解析 JSON

    自动处理 ```json 或 ``` 代码块格式

    Args:
        content: LLM 返回的原始内容

    Returns:
        解析后的字典，失败返回 None
    """
    if not content:
        return None

    # 提取 JSON 代码块（支持 ```json 或 ```）
    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
    json_str = json_match.group(1) if json_match else content

    try:
        parsed = json.loads(json_str.strip())
        # mimo 等模型可能返回 JSON 数组而非对象，取第一个元素
        if isinstance(parsed, list) and len(parsed) > 0 and isinstance(parsed[0], dict):
            return parsed[0]
        if isinstance(parsed, dict):
            return parsed
        return None
    except (json.JSONDecodeError, TypeError, ValueError):
        return None


def parse_json_list_from_response(content: str) -> Optional[List[dict]]:
    """
    从 LLM 响应中解析 JSON 数组

    Args:
        content: LLM 返回的原始内容

    Returns:
        解析后的列表，失败返回 None
    """
    if not content:
        return None

    # 提取 JSON 代码块
    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
    json_str = json_match.group(1) if json_match else content

    try:
        result = json.loads(json_str.strip())
        if isinstance(result, list):
            return result
        # 如果是 {"items": [...]} 格式，提取 items
        if isinstance(result, dict):
            for key in ['items', 'data', 'results', 'questions']:
                if key in result and isinstance(result[key], list):
                    return result[key]
        return None
    except (json.JSONDecodeError, TypeError, ValueError):
        return None


def extract_json_object(content: str) -> Optional[dict]:
    """
    多策略从 LLM 响应中提取 JSON 对象（增强版）

    在 parse_json_from_response 基础上增加 fallback 策略：
    1. 先调用 parse_json_from_response（markdown 代码块 → 直接解析）
    2. 失败后 fallback 到正则匹配最外层 {...} 块

    Args:
        content: LLM 返回的原始内容

    Returns:
        解析后的字典，全部策略失败返回 None
    """
    if not content:
        return None

    # 策略1+2：markdown 代码块提取 + 直接 json.loads
    result = parse_json_from_response(content)
    if result is not None and isinstance(result, dict):
        return result

    # 策略3（fallback）：正则匹配最外层 JSON 对象 {...}
    brace_match = re.search(r'\{[\s\S]*\}', content)
    if brace_match:
        try:
            parsed = json.loads(brace_match.group(0))
            if isinstance(parsed, dict):
                return parsed
        except (json.JSONDecodeError, TypeError, ValueError):
            pass

    return None


def extract_json_list(content: str) -> Optional[list]:
    """
    多策略从 LLM 响应中提取 JSON 数组（增强版）

    在 parse_json_list_from_response 基础上增加 fallback 策略：
    1. 先调用 parse_json_list_from_response（markdown 代码块 → 直接解析 → 嵌套提取）
    2. 失败后 fallback 到正则匹配最外层 [...] 块

    Args:
        content: LLM 返回的原始内容

    Returns:
        解析后的列表，全部策略失败返回 None
    """
    if not content:
        return None

    # 策略1+2：markdown 代码块提取 + 直接 json.loads + 嵌套 key 提取
    result = parse_json_list_from_response(content)
    if result is not None:
        return result

    # 策略3（fallback）：正则匹配最外层 JSON 数组 [...]
    bracket_match = re.search(r'\[[\s\S]*\]', content)
    if bracket_match:
        try:
            parsed = json.loads(bracket_match.group(0))
            if isinstance(parsed, list):
                return parsed
        except (json.JSONDecodeError, TypeError, ValueError):
            pass

    return None


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

def quick_ask(
    client,
    prompt: str,
    model: str,
    temperature: float = 0.3
) -> str:
    """
    快速提问（简化版）

    Args:
        client: OpenAI 客户端
        prompt: 问题
        model: 模型
        temperature: 温度

    Returns:
        回答内容，失败返回空字符串
    """
    result = call_llm(client, prompt, model, temperature=temperature)
    return result or ""


def quick_yes_no(
    client,
    prompt: str,
    model: str,
    keywords: List[str] = None
) -> bool:
    """
    快速是/否判断

    Args:
        client: OpenAI 客户端
        prompt: 问题（需包含判断标准）
        model: 模型
        keywords: 判断为 True 的关键词（默认 ["是", "需要", "yes", "true"]）

    Returns:
        布尔判断结果
    """
    if keywords is None:
        keywords = ["是", "需要", "yes", "true"]

    result = call_llm(client, prompt, model, temperature=0, max_tokens=128)
    if result is None:
        return False

    result_lower = result.lower()
    return any(kw.lower() in result_lower for kw in keywords)