tools/llm_client.py

import json
import re
import ast
from dataclasses import dataclass
from typing import Any

import requests

from config import (
    HF_PLANNER_USE_RESPONSE_FORMAT,
    HF_ROUTER_URL,
    HF_TEXT_MODEL,
    get_hf_token,
)
from tools.common import normalize_answer


@dataclass
class AgentAction:
    action: str
    args: dict[str, Any]
    thought: str = ""
    answer: str | None = None
    confidence: str = "medium"

    @property
    def is_final(self) -> bool:
        return self.action == "final_answer"


@dataclass
class QuestionClassification:
    question_type: str
    confidence: str = "medium"
    reason: str = ""
    query: str = ""
    raw: str = ""


@dataclass
class FinalAnswerFormat:
    answer: str
    confidence: str = "medium"
    raw: str = ""


def call_hf_chat(
    messages: list[dict[str, Any]],
    model: str = HF_TEXT_MODEL,
    max_tokens: int = 512,
    temperature: float = 0.1,
    response_format: dict[str, Any] | None = None,
) -> str:
    token = get_hf_token()
    if not token:
        raise RuntimeError("未配置 HF_TOKEN。")

    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    if response_format:
        payload["response_format"] = response_format

    response = requests.post(
        HF_ROUTER_URL,
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        },
        json=payload,
        timeout=180,
    )
    if not response.ok:
        raise RuntimeError(
            f"HF chat 请求失败：status={response.status_code}, body={response.text[:500]}"
        )
    data = response.json()
    try:
        message = data["choices"][0]["message"]
    except (KeyError, IndexError, TypeError) as exc:
        raise RuntimeError(f"HF chat 响应结构异常：{_compact_json(data)}") from exc
    return extract_message_text(message, data)


def extract_message_text(message: dict[str, Any], raw_response: dict[str, Any]) -> str:
    """兼容不同 HF Router provider 返回的 chat message content 格式。"""
    content = message.get("content")
    if isinstance(content, str) and content.strip():
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict):
                for key in ("text", "content"):
                    value = item.get(key)
                    if isinstance(value, str):
                        parts.append(value)
        joined = "\n".join(part for part in parts if part.strip()).strip()
        if joined:
            return joined

    fallback_parts = []
    for key in ("reasoning_content", "reasoning", "text"):
        value = message.get(key)
        if isinstance(value, str) and value.strip():
            fallback_parts.append(f"{key}={value}")
    if fallback_parts:
        return "\n".join(fallback_parts)

    raise RuntimeError(
        "HF chat 返回空内容。"
        f" message_keys={sorted(message.keys())}; raw={_compact_json(raw_response)}"
    )


def extract_json_object(text: str) -> dict[str, Any]:
    text = text.strip()
    candidates = [text]

    fenced_blocks = re.findall(
        r"```(?:json)?\s*(.*?)```",
        text,
        flags=re.IGNORECASE | re.DOTALL,
    )
    candidates.extend(block.strip() for block in fenced_blocks)
    candidates.extend(_balanced_json_candidates(text))

    errors = []
    for candidate in candidates:
        if not candidate:
            continue
        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError as exc:
            errors.append(str(exc))
        else:
            if isinstance(parsed, dict):
                return parsed

        try:
            parsed_literal = ast.literal_eval(candidate)
        except (SyntaxError, ValueError):
            continue
        if isinstance(parsed_literal, dict):
            return parsed_literal

    error_hint = f"; parse_errors={errors[:2]}" if errors else ""
    raise ValueError(f"模型没有输出 JSON 对象：{text[:300]}{error_hint}")


def _balanced_json_candidates(text: str) -> list[str]:
    candidates = []
    starts = [index for index, char in enumerate(text) if char == "{"]
    for start in starts:
        depth = 0
        in_string = False
        escaped = False
        for index in range(start, len(text)):
            char = text[index]
            if in_string:
                if escaped:
                    escaped = False
                elif char == "\\":
                    escaped = True
                elif char == '"':
                    in_string = False
                continue
            if char == '"':
                in_string = True
            elif char == "{":
                depth += 1
            elif char == "}":
                depth -= 1
                if depth == 0:
                    candidates.append(text[start : index + 1])
                    break
    return candidates


def _compact_json(data: Any, limit: int = 1200) -> str:
    text = json.dumps(data, ensure_ascii=False, default=str)
    if len(text) <= limit:
        return text
    return text[:limit] + "...<truncated>"


def parse_agent_action(text: str) -> AgentAction:
    data = extract_json_object(text)
    action = data.get("action") or data.get("tool")
    if not isinstance(action, str) or not action.strip():
        raise ValueError(f"模型 JSON 缺少 action 字段：{data}")

    args = data.get("args") or {}
    if not isinstance(args, dict):
        raise ValueError(f"action args 必须是对象：{data}")

    answer = data.get("answer")
    if answer is not None:
        answer = str(answer)

    return AgentAction(
        action=action.strip(),
        args=args,
        thought=str(data.get("thought", "")).strip(),
        answer=answer,
        confidence=str(data.get("confidence", "medium")).strip() or "medium",
    )


def classify_question_type(
    question: str,
    task_id: str,
    file_name: str,
    type_specs: list[dict[str, str]],
) -> QuestionClassification:
    allowed_types = {spec["type"] for spec in type_specs}
    type_lines = "\n".join(
        (
            f"- {spec['type']}: tool={spec['tool']}; "
            f"description={spec['description']}"
        )
        for spec in type_specs
    )
    prompt = f"""
Classify this GAIA task into exactly one question_type.

Available question types:
{type_lines}

Return exactly one JSON object:
{{
  "question_type": "one allowed type",
  "confidence": "high|medium|low",
  "reason": "short reason",
  "query": "optional search/query string if useful"
}}

Rules:
1. Do not answer the question.
2. Choose the tool category that should be tried first.
3. If no type clearly fits, use "unknown".
4. Prefer direct_text for self-contained table, string, list, or regex-like tasks.
5. Prefer python_code for .py attachments.
6. Prefer spreadsheet for .xlsx/.xls attachments.

Task:
task_id={task_id}
file_name={file_name}
question={question}
""".strip()

    raw = call_hf_chat(
        [
            {
                "role": "system",
                "content": (
                    "You are a strict GAIA task classifier. "
                    "Return JSON only. Never solve the task."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        max_tokens=350,
        temperature=0.0,
    )
    data = extract_json_object(raw)
    question_type = str(data.get("question_type") or data.get("type") or "unknown").strip()
    if question_type not in allowed_types:
        question_type = "unknown"
    confidence = str(data.get("confidence", "medium")).strip() or "medium"
    reason = str(data.get("reason", "")).strip()
    query = str(data.get("query", "")).strip()
    return QuestionClassification(
        question_type=question_type,
        confidence=confidence,
        reason=reason,
        query=query,
        raw=raw,
    )


def format_final_answer_with_llm(
    question: str,
    question_type: str,
    candidate_answer: str,
    evidence: str,
    source: str,
    confidence: str,
) -> FinalAnswerFormat:
    prompt = f"""
Format the tool result into the exact final answer required by the GAIA question.

Question:
{question}

Question type:
{question_type}

Tool source:
{source}

Tool confidence:
{confidence}

Candidate answer:
{candidate_answer or "<none>"}

Tool evidence:
{evidence or "<none>"}

Return exactly one JSON object:
{{"answer":"final answer only","confidence":"high|medium|low"}}

Rules:
1. Do not explain.
2. If candidate_answer is present, preserve its factual content and only fix formatting.
3. Obey requested capitalization, decimal places, ordering, separators, and units.
4. If evidence is insufficient and no candidate_answer is available, return answer "无法确定" with low confidence.
""".strip()
    raw = call_hf_chat(
        [
            {
                "role": "system",
                "content": (
                    "You are a strict final-answer formatter. "
                    "Return JSON only and never include reasoning."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        max_tokens=300,
        temperature=0.0,
    )
    data = extract_json_object(raw)
    answer = data.get("answer") or data.get("final_answer")
    if answer is None:
        raise ValueError(f"格式化模型 JSON 缺少 answer 字段：{data}")
    return FinalAnswerFormat(
        answer=normalize_answer(str(answer)),
        confidence=str(data.get("confidence", confidence)).strip() or confidence,
        raw=raw,
    )


def call_planner_model(messages: list[dict[str, Any]]) -> str:
    if not HF_PLANNER_USE_RESPONSE_FORMAT:
        return call_hf_chat(
            messages,
            max_tokens=700,
            temperature=0.0,
        )

    try:
        return call_hf_chat(
            messages,
            max_tokens=700,
            temperature=0.0,
            response_format={"type": "json_object"},
        )
    except RuntimeError as exc:
        if not _should_retry_without_response_format(str(exc)):
            raise
        return call_hf_chat(
            messages,
            max_tokens=700,
            temperature=0.0,
        )


def _should_retry_without_response_format(error_text: str) -> bool:
    lower_error = error_text.lower()
    return any(
        marker in lower_error
        for marker in (
            "response_format",
            "json_object",
            "json mode",
            "空内容",
            "empty",
        )
    )


def plan_next_action(messages: list[dict[str, Any]]) -> AgentAction:
    raw = call_planner_model(messages)
    try:
        return parse_agent_action(raw)
    except Exception as first_error:
        repair_messages = messages + [
            {
                "role": "assistant",
                "content": raw or "<empty response>",
            },
            {
                "role": "user",
                "content": (
                    "Your previous response was invalid because it was not one JSON object. "
                    "Return exactly one JSON object now, with no markdown and no prose. "
                    "Valid tool-call schema: "
                    "{\"thought\":\"...\",\"action\":\"tool_name\",\"args\":{}}. "
                    "Valid final-answer schema: "
                    "{\"thought\":\"...\",\"action\":\"final_answer\",\"answer\":\"...\",\"confidence\":\"high|medium|low\"}."
                ),
            },
        ]
        repaired_raw = call_planner_model(repair_messages)
        try:
            return parse_agent_action(repaired_raw)
        except Exception as second_error:
            raise ValueError(
                "planner 连续两次没有输出合法 JSON。"
                f" first_error={first_error}; first_raw={raw[:300]!r};"
                f" second_error={second_error}; second_raw={repaired_raw[:300]!r}"
            ) from second_error


def answer_with_light_model(question: str, evidence: str) -> str:
    prompt = f"""
你是一个 GAIA 问答 Agent。你会收到问题和工具已经整理好的证据。

硬性规则：
1. 只输出最终答案，不输出推理过程。
2. 严格遵守题目要求的格式、大小写、排序、逗号、单位和小数位。
3. 如果证据不足，不要编造；输出你能从证据中最可靠得到的答案。
4. 不要添加 "Answer:"、"最终答案：" 或解释性文字。

问题：
{question}

工具证据：
{evidence}
""".strip()

    raw_answer = call_hf_chat(
        [
            {
                "role": "system",
                "content": "You are a precise GAIA final-answer agent. Return only the final answer.",
            },
            {"role": "user", "content": prompt},
        ],
        max_tokens=500,
    )
    return normalize_answer(raw_answer)