agent.py

import os
import json
from typing import List, Dict, Any
from litellm import completion
from dotenv import load_dotenv
from logger import get_logger

logger = get_logger("agent")

load_dotenv()

def build_review_messages(template: str, data: str, data_type: str = "text") -> List[Dict[str, Any]]:
    """构建发给大模型的多模态 Messages"""
    system_msg = {"role": "system", "content": "你是一个专业的数据标注预审专家。请根据给定的评测标准对提交数据进行打分。"}
    
    # 如果是纯文本
    if data_type == "text":
        if "{data}" in template:
            user_content = template.replace("{data}", data)
        else:
            user_content = f"{template}\n\n待审核数据：\n{data}"
        return [system_msg, {"role": "user", "content": user_content}]
        
    # 如果是图片 (Vision 模式)
    elif data_type == "image_url":
        user_content = [
            {"type": "text", "text": template},
            {
                "type": "image_url",
                "image_url": {
                    "url": data # 假设前端传过来的是图片的公网 URL 或者 base64 字符串
                }
            }
        ]
        return [system_msg, {"role": "user", "content": user_content}]
    
    else:
        raise ValueError(f"Unsupported data_type: {data_type}")

def generate_review(messages: List[Dict[str, Any]], dimensions: List[str], ai_model: str = "gpt-4o-mini") -> Dict[str, Any]:
    """调用大模型进行结构化打分"""
    
    properties = {
        dim: {
            "type": "integer",
            "description": f"评分维度: {dim} (1-10分)"
        } for dim in dimensions
    }
    
    tools = [
        {
            "type": "function",
            "function": {
                "name": "submit_review_result",
                "description": "提交审核评分与结论",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "scores": {
                            "type": "object",
                            "properties": properties,
                            "required": dimensions,
                            "description": "各维度的具体评分"
                        },
                        "ai_comment": {
                            "type": "string",
                            "description": "AI 综合评语与质检建议"
                        },
                        "conclusion": {
                            "type": "string",
                            "enum": ["pass", "return"],
                            "description": "最终结论：pass (通过) 或 return (打回人工复核)"
                        }
                    },
                    "required": ["scores", "ai_comment", "conclusion"]
                }
            }
        }
    ]

    try:
        response = completion(
            model=ai_model,
            messages=messages,
            tools=tools,
            tool_choice={"type": "function", "function": {"name": "submit_review_result"}},
            temperature=0.1
        )
        
        tool_call = response.choices[0].message.tool_calls[0]
        arguments = json.loads(tool_call.function.arguments)
        return arguments
    except Exception as e:
        logger.error(f"大模型 {ai_model} 调用失败: {str(e)}")
        raise RuntimeError(f"大模型 {ai_model} 调用失败: {str(e)}")