llm_parser.py

"""
llm_parser.py
利用 LLM 将用户的自然语言描述解析为结构化、经过校验的桥梁设计参数
支持"一次性生成"与"多轮对话式增量修改"两种调用方式
"""
import json
import os
import re
import time
from dataclasses import dataclass
from typing import List, Optional

from dotenv import load_dotenv
from openai import OpenAI
from pydantic import BaseModel, field_validator

import design_search
import llm_logger
import span_layout_search

# 匹配数字（含负号、小数点），用于从可能带单位的字符串中提取数值，
# 例如不同 LLM/服务商对"请返回浮点数"的遵循程度不一，
# 智谱 GLM 实测就会返回 "30m" 这类带单位的跨径而不是纯数字。
_NUMBER_PATTERN = re.compile(r"-?\d+\.?\d*")


def _extract_float(value, default: float) -> float:
    """从数字或"带单位的字符串"中提取出第一个数值；提取失败则返回给定的默认值"""
    if isinstance(value, (int, float)):
        return float(value)
    match = _NUMBER_PATTERN.search(str(value))
    if not match:
        return default
    try:
        return float(match.group())
    except ValueError:
        return default


load_dotenv()

_API_KEY = os.getenv("OPENAI_API_KEY")
_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.deepseek.com")
_MODEL = os.getenv("OPENAI_MODEL", "deepseek-chat")

# 供前端展示默认值 / 判断是否已经在本机 .env 中配置了 Key
DEFAULT_BASE_URL = _BASE_URL
DEFAULT_MODEL = _MODEL


def has_env_key() -> bool:
    """本机 .env 中是否已配置 OPENAI_API_KEY（用于 UI 提示，不代表运行时实际使用的 Key）"""
    return bool(_API_KEY)


SYSTEM_PROMPT = """
你是一个专业的桥梁结构工程助手。请从用户的输入中提取桥梁设计参数，并严格返回 JSON 格式。
需要提取的字段如下：
- spans: 跨径序列（字符串，如 "30.0, 30.0, 30.0"）
- beam_height: 主梁高度（浮点数，单位m）
- pier_height: 桥墩高度（浮点数，单位m）
- bridge_width: 桥面宽度（浮点数，单位m）

如果用户没有提供某个参数，请给出合理的工程默认值（spans: "30.0, 30.0, 30.0", beam_height: 2.5, pier_height: 12.0, bridge_width: 15.0）。
只返回合法 JSON，不要包含任何解释文字。
"""

REFINE_SYSTEM_PROMPT = """
你是一个专业的桥梁结构工程助手，正在与用户进行多轮对话式桥梁参数设计。
你会收到"当前设计参数"（JSON）和用户这一轮提出的需求（可能是全新描述，也可能是局部调整，
例如"把第二跨改成25米"或"墩再高一点"）。

请在当前参数基础上，按用户本轮要求更新对应字段，未提及的字段保持不变，
并严格返回更新后的完整 JSON（不是增量 diff，而是完整的新参数集合）。

字段定义：
- spans: 跨径序列（字符串，如 "30.0, 30.0, 30.0"，逗号分隔；若用户要求增减跨数请相应调整跨的个数）
- beam_height: 主梁高度（浮点数，单位m）
- pier_height: 桥墩高度（浮点数，单位m）
- bridge_width: 桥面宽度（浮点数，单位m）

只返回合法 JSON，不要包含任何解释文字。
"""


class BridgeParams(BaseModel):
    """LLM 解析结果的字段校验与安全范围钳制（范围需与 app.py 中滑块的取值范围保持一致）"""

    spans: str = "30.0, 30.0, 30.0"
    beam_height: float = 2.5
    pier_height: float = 12.0
    bridge_width: float = 15.0

    @field_validator("spans")
    @classmethod
    def validate_spans(cls, v):
        # 对每一段逗号分隔的文本单独提取数字，兼容 "30m" "30.0 米" " 30 " 等各种带单位/空格的写法
        values = []
        for token in str(v).split(","):
            match = _NUMBER_PATTERN.search(token)
            if not match:
                continue
            try:
                num = float(match.group())
            except ValueError:
                continue
            if num > 0:
                values.append(num)

        if not values:
            return "30.0, 30.0, 30.0"
        return ", ".join(f"{num:.1f}" for num in values)

    @field_validator("beam_height", mode="before")
    @classmethod
    def extract_beam_height(cls, v):
        return _extract_float(v, default=2.5)

    @field_validator("beam_height")
    @classmethod
    def clamp_beam_height(cls, v):
        return max(1.0, min(5.0, float(v)))

    @field_validator("pier_height", mode="before")
    @classmethod
    def extract_pier_height(cls, v):
        return _extract_float(v, default=12.0)

    @field_validator("pier_height")
    @classmethod
    def clamp_pier_height(cls, v):
        return max(4.0, min(30.0, float(v)))

    @field_validator("bridge_width", mode="before")
    @classmethod
    def extract_bridge_width(cls, v):
        return _extract_float(v, default=15.0)

    @field_validator("bridge_width")
    @classmethod
    def clamp_bridge_width(cls, v):
        return max(3.0, min(40.0, float(v)))


def _resolve_config(api_key: str | None, base_url: str | None, model: str | None):
    resolved_key = api_key or _API_KEY
    resolved_base_url = base_url or _BASE_URL
    resolved_model = model or _MODEL
    if not resolved_key:
        raise RuntimeError(
            "未检测到 API Key：请在侧边栏手动粘贴你的 Key，"
            "或复制 .env.example 为 .env 并填入后重启应用。"
        )
    return resolved_key, resolved_base_url, resolved_model


def _call_llm_for_params(
    messages: list,
    log_user_input: str,
    api_key: str | None,
    base_url: str | None,
    model: str | None,
) -> tuple[BridgeParams, str]:
    """调用 LLM、解析并校验为 BridgeParams，同时把本次调用记录到本地日志（logs/llm_calls.db）

    返回 (解析结果, 实际服务这次请求的模型名)。第二个值优先取 API 响应里的
    response.model——像 DeepSeek 这类服务商的"deepseek-chat"是会漂移的别名
    （同一个别名两年内先后指向过好几个不同的底层模型版本），只记录请求时传的
    别名字符串不足以在事后确定当时到底是哪个模型在服务；response.model 拿不到
    时才退回请求用的别名字符串。"""
    resolved_key, resolved_base_url, resolved_model = _resolve_config(api_key, base_url, model)

    # 不缓存全局单例：不同用户可能带着各自的 Key 调用同一个服务
    client = OpenAI(api_key=resolved_key, base_url=resolved_base_url)

    start = time.perf_counter()
    raw_response_text = None

    try:
        response = client.chat.completions.create(
            model=resolved_model,
            messages=messages,
            temperature=0.1,
            response_format={"type": "json_object"},
        )
        raw_response_text = response.choices[0].message.content.strip()
        actual_model = getattr(response, "model", None) or resolved_model
    except Exception as exc:
        latency_ms = (time.perf_counter() - start) * 1000
        llm_logger.log_call(
            resolved_base_url, resolved_model, log_user_input, None, False, latency_ms, str(exc)
        )
        raise RuntimeError(f"调用 LLM 接口失败：{exc}") from exc

    try:
        raw = json.loads(raw_response_text)
        result = BridgeParams(**raw)
    except Exception as exc:
        latency_ms = (time.perf_counter() - start) * 1000
        llm_logger.log_call(
            resolved_base_url, resolved_model, log_user_input, raw_response_text, False, latency_ms, str(exc)
        )
        if isinstance(exc, json.JSONDecodeError):
            raise RuntimeError(f"LLM 返回内容不是合法 JSON：{raw_response_text[:200]}") from exc
        raise RuntimeError(f"LLM 返回的参数格式不符合预期：{exc}") from exc

    latency_ms = (time.perf_counter() - start) * 1000
    llm_logger.log_call(
        resolved_base_url, resolved_model, log_user_input, raw_response_text, True, latency_ms
    )
    return result, actual_model


def parse_bridge_params_with_llm(
    user_input: str,
    api_key: str | None = None,
    base_url: str | None = None,
    model: str | None = None,
) -> BridgeParams:
    """利用 LLM 将用户的自然语言一次性转化为结构化、经过校验的桥梁设计参数

    :param api_key: 显式传入的 Key（例如用户在网页上临时输入的），优先于本机 .env
    :param base_url: 显式传入的 API 地址，优先于本机 .env
    :param model: 显式传入的模型名，优先于本机 .env
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_input},
    ]
    result, _actual_model = _call_llm_for_params(messages, user_input, api_key, base_url, model)
    return result


def parse_bridge_params_with_llm_returning_model(
    user_input: str,
    api_key: str | None = None,
    base_url: str | None = None,
    model: str | None = None,
) -> tuple[BridgeParams, str]:
    """与 parse_bridge_params_with_llm 相同，但额外返回实际服务这次请求的模型名
    （API 响应里的 response.model，不是请求时传的别名字符串）。供评测脚本
    （eval/evaluate_llm.py）使用，好在结果 JSON 里记录清楚 Table 5 每一条 case
    到底是哪个具体模型跑出来的——不加进 parse_bridge_params_with_llm 本体，
    是为了不改动真实 app.py 里现有调用方依赖的返回类型。"""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_input},
    ]
    return _call_llm_for_params(messages, user_input, api_key, base_url, model)


def refine_bridge_params_with_llm(
    user_input: str,
    current_params: BridgeParams,
    api_key: str | None = None,
    base_url: str | None = None,
    model: str | None = None,
) -> BridgeParams:
    """在当前参数基础上，根据用户这一轮的自然语言指令做增量修改（多轮对话式设计）

    :param current_params: 当前的设计参数，作为上下文一并发给 LLM
    """
    context = current_params.model_dump_json()
    messages = [
        {"role": "system", "content": REFINE_SYSTEM_PROMPT},
        {"role": "user", "content": f"当前设计参数：{context}\n\n本轮需求：{user_input}"},
    ]
    result, _actual_model = _call_llm_for_params(messages, user_input, api_key, base_url, model)
    return result


# ----------------------------------------------------
# 对话式工具调用（function calling）：
# 让 LLM 判断这一轮消息是"修改设计参数"还是"要求自动调整配筋以满足承载力"，
# 并分别路由到对应的确定性函数——LLM 只负责理解意图、组织回复文字，
# 不自己计算任何数字，实际数值都由本地的确定性代码算出。
# ----------------------------------------------------

UPDATE_PARAMS_TOOL = {
    "type": "function",
    "function": {
        "name": "update_bridge_params",
        "description": (
            "当用户描述或修改桥梁跨径、梁高、墩高、桥宽等基本几何参数时调用。"
            "只需要包含本轮用户明确要修改的字段——没有提到、不需要变化的字段"
            "请不要包含在参数里（系统会自动保持原值不变），绝不要猜测或凭经验"
            "编造未提及字段的数值，也不要把当前值原样重复传一遍。"
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "spans": {
                    "type": "string",
                    "description": "跨径序列，逗号分隔，如 30.0, 30.0, 30.0；仅在用户明确要修改跨径时提供",
                },
                "beam_height": {"type": "number", "description": "主梁高度(m)；仅在用户明确要修改梁高时提供"},
                "pier_height": {"type": "number", "description": "桥墩高度(m)；仅在用户明确要修改墩高时提供"},
                "bridge_width": {"type": "number", "description": "桥面宽度(m)；仅在用户明确要修改桥宽时提供"},
            },
            "required": [],
        },
    },
}

ADJUST_REINFORCEMENT_TOOL = {
    "type": "function",
    "function": {
        "name": "auto_adjust_reinforcement",
        "description": (
            "当用户询问当前方案的截面承载力是否足够、或要求自动调整方案"
            "以满足结构承载力要求时调用。这个工具会在本地做精确数值搜索，"
            "优先只调整梁片数/主筋直径（改动最小）；如果这样不够，还会尝试"
            "加大梁高（改动更大，会明确在结果里说明）。"
            "你不需要、也不应该自己计算或编造任何数字。"
        ),
        "parameters": {"type": "object", "properties": {}, "required": []},
    },
}

SUGGEST_SPAN_LAYOUT_TOOL = {
    "type": "function",
    "function": {
        "name": "suggest_span_layout",
        "description": (
            "当用户给出桥梁总长（以及可选的跨数范围要求），询问应该怎么分跨、"
            "推荐什么跨径布置方案时调用。这个工具会在本地对等跨、边跨缩短两类"
            "常见布置方式做完整的结构可行性搜索和工程量比较，推荐一个满足承载力"
            "要求且混凝土用量最省的方案。你不需要、也不应该自己计算或猜测"
            "任何跨径数值或工程量。"
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "total_length": {"type": "number", "description": "桥梁总长(m)，必须提供"},
                "num_spans_min": {
                    "type": "integer",
                    "description": "允许的最小跨数；用户没有明确要求时不填，使用默认值",
                },
                "num_spans_max": {
                    "type": "integer",
                    "description": "允许的最大跨数；用户没有明确要求时不填，使用默认值",
                },
            },
            "required": ["total_length"],
        },
    },
}

CHAT_SYSTEM_PROMPT = """
你是一个专业的桥梁结构工程助手，正在与用户进行多轮对话式桥梁参数设计。
对话历史中可能包含之前几轮的交流，用于帮你理解上下文（例如用户在追问上一轮
调整的结果是否满足要求）；但每一轮消息里提供的"当前设计参数"JSON 才是当前
真正生效的数值（可能已经被用户手动调整过，请以这个 JSON 为准，不要用对话
历史里更早的数字）。

你有三个工具可以调用：
1. update_bridge_params：当用户在描述/修改跨径、梁高、墩高、桥宽等基本几何参数时调用
2. auto_adjust_reinforcement：当用户询问承载力/配筋是否足够、是否安全、或要求自动
   调整配筋/梁片数以满足承载力要求时调用——包括"现在够了吗""这样安全吗"这类
   延续上一轮承载力话题的追问
3. suggest_span_layout：当用户给出桥梁总长、询问应该怎么分跨/推荐跨径布置方案时调用
   ——例如"总长60米，应该怎么分跨比较好""帮我看看100米分3到5跨哪种方案好"

请结合对话历史判断真实意图，调用其中一个工具来完成请求，不要直接用文字回复、
也不要自己编造任何工程数值。

调用 auto_adjust_reinforcement 后，工具会返回一个带 "applied" 字段的 JSON：
- applied 为 true：调整已经真正生效，如实告诉用户新的梁片数/主筋直径/利用率；
  如果 JSON 里 "beam_h_changed" 为 true，说明光调配筋不够，梁高也被加大了，
  这是比单纯加钢筋更大的改动，必须在回复里明确提到梁高变化（新的梁高数值），
  不能只字不提，让用户误以为只是加了点钢筋
- applied 为 false：**没有做任何改动**，梁高、梁片数、主筋直径都还是原来的值，
  你的回复必须明确说明"未能找到可行方案、方案没有被修改"，绝对不能说成
  "已经调整为...片/...mm"这种暗示改动已生效的表述，即使 JSON 里的字段有
  具体数字，那也只是当前保持不变的原值，不是新方案。

调用 suggest_span_layout 后，工具同样会返回一个带 "found" 字段的 JSON：
- found 为 true：已经找到推荐方案并自动应用到当前设计（跨径、梁高、梁片数、
  主筋直径都会更新），如实告诉用户推荐的跨径列表、方案类型（等跨/边跨缩短）、
  以及对应的混凝土用量，说明这是同类可行方案里最省的一个
- found 为 false：**没有应用任何改动**，当前设计保持不变，你的回复必须明确
  说明在给定的跨数范围内找不到可行方案，并转达 message 里给出的建议
  （比如扩大跨数范围），不能编造一个"可行"的方案
"""


@dataclass
class ChatResult:
    """handle_chat_message 的返回结果：updated_params / reinforcement_result /
    span_layout_result 三者至多一个非空（取决于本轮 LLM 调用了哪个工具）"""

    updated_params: Optional[BridgeParams] = None
    reinforcement_result: Optional[design_search.ReinforcementSearchResult] = None
    span_layout_result: Optional[span_layout_search.SpanLayoutResult] = None
    reply: str = ""
    resolved_model: str = ""  # API 响应里的 response.model，不是请求时传的别名字符串（见 _call_llm_for_params 的说明）


MAX_CHAT_HISTORY_MESSAGES = 10  # 只带最近几轮对话，避免上下文无限增长拖慢/推高成本


def handle_chat_message(
    user_input: str,
    current_params: BridgeParams,
    pier_h: float,
    q_live: float,
    current_num_girders: int,
    current_bar_diameter_mm: int,
    chat_history: Optional[List[dict]] = None,
    api_key: str | None = None,
    base_url: str | None = None,
    model: str | None = None,
) -> ChatResult:
    """
    统一的对话入口：用 function calling 让模型判断这一轮用户消息属于哪种意图，
    执行对应的确定性计算（更新设计参数 / 搜索满足承载力的配筋方案），
    再请模型把执行结果组织成一句自然语言回复。

    :param chat_history: 之前几轮对话（[{"role": "user"/"assistant", "content": str}, ...]），
        不包含本轮消息。用于让模型理解"现在够了吗"这类延续上一轮话题的追问，
        只取最近 MAX_CHAT_HISTORY_MESSAGES 条。
    """
    resolved_key, resolved_base_url, resolved_model = _resolve_config(api_key, base_url, model)
    client = OpenAI(api_key=resolved_key, base_url=resolved_base_url)

    context = current_params.model_dump_json()
    messages: List[dict] = [{"role": "system", "content": CHAT_SYSTEM_PROMPT}]
    if chat_history:
        for msg in chat_history[-MAX_CHAT_HISTORY_MESSAGES:]:
            messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": f"当前设计参数：{context}\n\n本轮需求：{user_input}"})

    start = time.perf_counter()
    try:
        response = client.chat.completions.create(
            model=resolved_model,
            messages=messages,
            tools=[UPDATE_PARAMS_TOOL, ADJUST_REINFORCEMENT_TOOL, SUGGEST_SPAN_LAYOUT_TOOL],
            tool_choice="auto",
            temperature=0.1,
        )
    except Exception as exc:
        latency_ms = (time.perf_counter() - start) * 1000
        llm_logger.log_call(
            resolved_base_url, resolved_model, user_input, None, False, latency_ms, str(exc)
        )
        raise RuntimeError(f"调用 LLM 接口失败：{exc}") from exc

    actual_model = getattr(response, "model", None) or resolved_model
    assistant_message = response.choices[0].message
    tool_calls = assistant_message.tool_calls

    if not tool_calls:
        latency_ms = (time.perf_counter() - start) * 1000
        llm_logger.log_call(
            resolved_base_url, resolved_model, user_input, assistant_message.content, True, latency_ms
        )
        return ChatResult(
            reply=assistant_message.content or "抱歉，我没能理解这个请求，可以换个说法再试一次。",
            resolved_model=actual_model,
        )

    result = ChatResult(resolved_model=actual_model)
    messages.append(assistant_message.model_dump(exclude_none=True))
    called_tool_names = []

    for tool_call in tool_calls:
        name = tool_call.function.name
        called_tool_names.append(name)
        try:
            args = json.loads(tool_call.function.arguments or "{}")
        except json.JSONDecodeError:
            args = {}

        if name == "update_bridge_params":
            try:
                # 只用模型返回的字段覆盖当前参数，其余保持 current_params 真实值——
                # 不能直接 BridgeParams(**args)，模型没返回的字段会被类默认值覆盖
                merged = current_params.model_dump()
                merged.update(args)
                parsed = BridgeParams(**merged)
                result.updated_params = parsed
                tool_output = parsed.model_dump_json()
            except Exception as exc:
                tool_output = f"参数格式有误：{exc}"
        elif name == "auto_adjust_reinforcement":
            try:
                spans = [float(s.strip()) for s in current_params.spans.split(",") if s.strip()]
                search_result = design_search.find_adequate_reinforcement(
                    spans=spans,
                    beam_h=current_params.beam_height,
                    pier_h=pier_h,
                    bridge_w=current_params.bridge_width,
                    q_live=q_live,
                    current_num_girders=current_num_girders,
                    current_bar_diameter_mm=current_bar_diameter_mm,
                )
                result.reinforcement_result = search_result
                # 传结构化 JSON 而不是一句话：显式的 applied/beam_h_changed 字段
                # 让模型难以把"没有生效的调整"转述成"已经调整好了"
                tool_output = json.dumps(
                    {
                        "applied": search_result.found,
                        "beam_h": search_result.beam_h,
                        "beam_h_changed": search_result.found
                        and search_result.beam_h != current_params.beam_height,
                        "num_girders": search_result.num_girders,
                        "bar_diameter_mm": search_result.bar_diameter_mm,
                        "worst_utilization": (
                            None
                            if search_result.worst_utilization == float("inf")
                            else search_result.worst_utilization
                        ),
                        "message": search_result.message,
                    },
                    ensure_ascii=False,
                )
            except Exception as exc:
                tool_output = f"自动调整过程出错：{exc}"
        elif name == "suggest_span_layout":
            try:
                total_length = float(args["total_length"])
                layout_result = span_layout_search.find_span_layout(
                    total_length=total_length,
                    bridge_w=current_params.bridge_width,
                    pier_h=pier_h,
                    q_live=q_live,
                    num_spans_min=int(args.get("num_spans_min", 2)),
                    num_spans_max=int(args.get("num_spans_max", 5)),
                    current_num_girders=current_num_girders,
                    current_bar_diameter_mm=current_bar_diameter_mm,
                )
                result.span_layout_result = layout_result
                # 同上：显式的 found 字段防止模型在没找到方案时编出跨径列表
                best = layout_result.best
                tool_output = json.dumps(
                    {
                        "found": layout_result.found,
                        "spans": best.spans if best else None,
                        "pattern": best.pattern if best else None,
                        "beam_h": best.beam_h if best else None,
                        "num_girders": best.num_girders if best else None,
                        "bar_diameter_mm": best.bar_diameter_mm if best else None,
                        "worst_utilization": best.worst_utilization if best else None,
                        "total_concrete_m3": best.total_concrete_m3 if best else None,
                        "num_feasible_candidates": len(layout_result.candidates),
                        "message": layout_result.message,
                    },
                    ensure_ascii=False,
                )
            except Exception as exc:
                tool_output = f"跨径布置搜索过程出错：{exc}"
        else:
            tool_output = f"未知工具：{name}"

        messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": tool_output})

    try:
        followup = client.chat.completions.create(
            model=resolved_model, messages=messages, temperature=0.1
        )
        reply_text = followup.choices[0].message.content or ""
    except Exception as exc:
        # 工具已经在本地执行成功，只是总结文案生成失败，兜底给用户一句说明
        reply_text = f"（已完成计算，但生成总结文案时出错：{exc}）"

    latency_ms = (time.perf_counter() - start) * 1000
    llm_logger.log_call(
        resolved_base_url,
        resolved_model,
        user_input,
        json.dumps({"tool_calls": called_tool_names}, ensure_ascii=False),
        True,
        latency_ms,
    )

    result.reply = reply_text
    return result