import openai

# 定义函数工具
tools = [
    {
        "type": "function",
        "function": {
            "name": "search_weather",
            "description": "查询城市天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称",
                        "enum": ["北京", "上海", "广州", "深圳"]
                    },
                    "date": {
                        "type": "string",
                        "description": "查询日期"
                    }
                },
                "required": ["city", "date"]
            }
        }
    }
]

# LLM会严格按照Schema输出
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "北京明天天气怎么样？"}],
    tools=tools,
    tool_choice="auto"
)

# response中的tool_calls已经是结构化的
print(response.choices[0].message.tool_calls)

from pydantic import BaseModel, validator
from typing import Literal

class WeatherQuery(BaseModel):
    city: Literal["北京", "上海", "广州", "深圳"]
    date: str

    @validator("date")
    def validate_date(cls, v):
        import re
        if not re.match(r"^\d{4}-\d{2}-\d{2}$", v):
            raise ValueError("日期格式必须为YYYY-MM-DD")
        return v

def handle_llm_output(raw_output):
    try:
        # 尝试解析
        data = json.loads(raw_output)

        # Pydantic验证
        validated = WeatherQuery(**data)
        return validated.dict()

    except json.JSONDecodeError:
        # JSON格式错误，尝试修复
        return fix_and_retry(raw_output)

    except ValidationError as e:
        # 验证失败，尝试修复
        return fix_validation_error(raw_output, e)

def robust_tool_call(prompt, max_retries=3):
    for attempt in range(max_retries):
        try:
            response = llm.generate(prompt)

            # 解析验证
            parsed = parse_and_validate(response)

            # 执行工具调用
            result = execute_tool(parsed)

            return {"success": True, "result": result}

        except (JSONDecodeError, ValidationError) as e:
            if attempt == max_retries - 1:
                # 最后一次尝试失败，使用降级策略
                return fallback_strategy(prompt)

            # 反馈错误，让LLM重新生成
            correction_prompt = f"""
                上一轮输出存在问题：
                {str(e)}

                请重新生成正确的JSON输出。
                用户原始需求：{prompt}
            """
            prompt = correction_prompt

    return {"success": False, "error": "max retries exceeded"}

correction_prompt = """
上一轮输出存在问题，请修正：

错误类型：{error_type}
错误详情：{error_message}

请严格按照以下Schema重新生成：
{schema_definition}

只输出JSON，不要其他内容。
"""

import re
import json

def fix_json_string(raw_string):
    """修复常见的JSON格式错误"""

    # 1. 移除多余逗号
    fixed = re.sub(r',(\s*[}\]])', r'\1', raw_string)

    # 2. 修复单引号为双引号
    # 注意：这是一个简化版本，复杂场景需要更智能的处理

    # 3. 移除注释
    fixed = re.sub(r'//.*', '', fixed)
    fixed = re.sub(r'/\*.*?\*/', '', fixed, flags=re.DOTALL)

    return fixed

def smart_json_parse(text):
    """智能JSON解析，带自动修复"""

    # 尝试直接解析
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 尝试修复后解析
    fixed = fix_json_string(text)
    try:
        return json.loads(fixed)
    except json.JSONDecodeError:
        pass

    # 尝试提取JSON部分
    json_match = re.search(r'\{.*\}', text, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group())
        except:
            pass

    raise ValueError("无法解析为有效JSON")

from pydantic import BaseModel, Field
from typing import Optional

class WeatherQuery(BaseModel):
    city: str = Field(..., min_length=2, max_length=10)
    date: str = Field(..., pattern=r"^\d{4}-\d{2}-\d{2}$")

class WeatherOptions(BaseModel):
    include_aqi: bool = True
    include_forecast: bool = False

class WeatherRequest(BaseModel):
    query: WeatherQuery
    options: Optional[WeatherOptions] = None

from typing import List, Literal

class WeatherBatchRequest(BaseModel):
    cities: List[Literal["北京", "上海", "广州", "深圳"]] = Field(
        ...,
        min_items=1,
        max_items=10,
        description="城市列表，最多10个城市"
    )
    metrics: List[Literal["temperature", "humidity", "wind"]] = Field(
        ...,
        min_items=1
    )

from pydantic import BaseModel, root_validator

class OrderRequest(BaseModel):
    order_type: Literal["buy", "sell"]

    # 根据订单类型，必填字段不同
    buy_quantity: Optional[float] = None
    sell_quantity: Optional[float] = None

    @root_validator
    def validate_quantities(cls, values):
        if values.get("order_type") == "buy" and not values.get("buy_quantity"):
            raise ValueError("买入订单必须指定buy_quantity")
        if values.get("order_type") == "sell" and not values.get("sell_quantity"):
            raise ValueError("卖出订单必须指定sell_quantity")
        return values

def sanitize_tool_input(data):
    """清理可能的安全威胁"""

    # 移除可疑字段
    blocked_fields = ["system_prompt", "instructions", "admin"]
    for field in blocked_fields:
        data.pop(field, None)

    # 类型检查
    for key, value in data.items():
        if isinstance(value, str) and len(value) > 1000:
            raise ValueError(f"字段 {key} 值过长，可能存在注入风险")

    return data

class SafeWeatherQuery(BaseModel):
    city: str = Field(..., min_length=1, max_length=50)
    date: str = Field(..., regex=r"^\d{4}-\d{2}-\d{2}$")

    @field_validator("city")
    @classmethod
    def validate_city_safe(cls, v):
        # 白名单校验
        allowed = {"北京", "上海", "广州", "深圳", "杭州", "成都"}
        if v not in allowed:
            raise ValueError(f"不支持的城市：{v}")
        return v

from functools import lru_cache
import hashlib

@lru_cache(maxsize=1000)
def cached_parse(template, raw_output):
    """缓存解析结果"""
    cache_key = hashlib.md5(f"{template}:{raw_output}".encode()).hexdigest()
    return parse_json(raw_output)

from concurrent.futures import ThreadPoolExecutor

def parallel_validate(requests):
    with ThreadPoolExecutor(max_workers=10) as executor:
        futures = [executor.submit(validate_single, req) for req in requests]
        return [f.result() for f in futures]

import pytest

def test_weather_query_valid():
    data = {"city": "北京", "date": "2026-05-10"}
    result = WeatherQuery(**data)
    assert result.city == "北京"
    assert result.date == "2026-05-10"

def test_weather_query_invalid_city():
    with pytest.raises(ValidationError):
        WeatherQuery(city="东京", date="2026-05-10")

def test_weather_query_invalid_date():
    with pytest.raises(ValidationError):
        WeatherQuery(city="北京", date="05-10-2026")

import fuzzing

@fuzzing.produces(WeatherQuery)
def generate_weather_query():
    return {
        "city": random.choice(["北京", "上海", "INVALID"]),
        "date": f"{random.randint(2020, 2030)}-{random.randint(1,12):02d}-{random.randint(1,28):02d}"
    }

# 测试边界情况和随机输入