#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : AI.  @by PyCharm
# @File         : message_types
# @Time         : 2024/11/21 19:11
# @Author       : betterme
# @WeChat       : meutils
# @Software     : PyCharm
# @Description  :


from typing import Union, Iterable, Optional, List
from pydantic import BaseModel, Field
from typing_extensions import Literal


class ChatCompletionContentPartTextParam(BaseModel):
    type: Literal["text"] = "text"
    text: str


class ChatCompletionContentPartRefusalParam(BaseModel):
    type: Literal["refusal"] = "refusal"
    refusal: str


class ChatCompletionMessageToolCallParam(BaseModel):
    id: str
    type: str
    function: dict


class FunctionCall(BaseModel):
    arguments: str = Field(...,
                           description="The arguments to call the function with, as generated by the model in JSON format")
    name: str = Field(..., description="The name of the function to call")


class ChatCompletionToolMessageParam(BaseModel):
    content: Union[str, List[ChatCompletionContentPartTextParam]] = Field(...,
                                                                          description="The contents of the tool message")
    role: Literal["tool"] = Field(..., description="The role of the messages author")
    tool_call_id: str = Field(..., description="Tool call that this message is responding to")


class ChatCompletionUserMessageParam(BaseModel):
    content: Union[str, List[ChatCompletionContentPartTextParam]] = Field(...,
                                                                          description="The contents of the user message")
    role: Literal["user"] = Field(..., description="The role of the messages author")
    name: Optional[str] = Field(None, description="An optional name for the participant")


class ChatCompletionSystemMessageParam(BaseModel):
    content: Union[str, List[ChatCompletionContentPartTextParam]] = Field(...,
                                                                          description="The contents of the system message")
    role: Literal["system"] = Field(..., description="The role of the messages author")
    name: Optional[str] = Field(None, description="An optional name for the participant")


class ChatCompletionFunctionMessageParam(BaseModel):
    content: Optional[str] = Field(..., description="The contents of the function message")
    name: str = Field(..., description="The name of the function to call")
    role: Literal["function"] = Field(..., description="The role of the messages author")


class ChatCompletionAssistantMessageParam(BaseModel):
    role: Literal["assistant"] = Field(..., description="The role of the messages author")
    content: Optional[
        Union[str, List[Union[ChatCompletionContentPartTextParam, ChatCompletionContentPartRefusalParam]]]] = Field(
        None, description="The contents of the assistant message"
    )
    function_call: Optional[FunctionCall] = Field(None,
                                                  description="The name and arguments of a function that should be called")
    name: Optional[str] = Field(None, description="An optional name for the participant")
    refusal: Optional[str] = Field(None, description="The refusal message by the assistant")
    tool_calls: Optional[List[ChatCompletionMessageToolCallParam]] = Field(None,
                                                                           description="The tool calls generated by the model")

class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: Union[str, List[Union[ChatCompletionUserMessageParam, FileUrlContent]], FileUrlContent]

class ChatMessages(BaseModel):
    messages: List[Message]

class ChatCompletionMessageParam(
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
    ChatCompletionAssistantMessageParam,
    ChatCompletionToolMessageParam,
    ChatCompletionFunctionMessageParam
):
    pass


if __name__ == '__main__':
    # 创建一个用户消息
    user_msg = ChatCompletionUserMessageParam(
        role="user",
        content="Hello!",
        name="John"
    )

    # 创建一个系统消息
    system_msg = ChatCompletionSystemMessageParam(
        role="system",
        content="You are a helpful assistant."
    )

    # 创建一个助手消息
    assistant_msg = ChatCompletionAssistantMessageParam(
        role="assistant",
        content="How can I help you today?"
    )

    # 验证消息
    # print(user_msg.model_dump())
    # print(system_msg.model_dump())
    # print(assistant_msg.model_dump())
    messages = [
        {
            'role': 'system',
            'content': '你是一个文件问答助手'
        },
        {
            'role': 'user',
            # 'content': {
            #     "type": "file_url",
            #     "file_url": {"url": "https://oss.ffire.cc/files/招标文件备案表（第二次）.pdf", "detai": "auto"}
            # },
            'content': [
                {
                    "type": "text",
                    "text": "这个文件讲了什么？"
                },
                # 多轮的时候要剔除
                {
                    "type": "file_url",
                    "file_url": {"url": "https://oss.ffire.cc/files/招标文件备案表（第二次）.pdf", "detai": "auto"}
                }
            ]
        },
        {'role': 'assistant', 'content': "好的"},
        {
            'role': 'user',
            # 'content': {
            #     "type": "file_url",
            #     "file_url": {"url": "https://oss.ffire.cc/files/招标文件备案表（第二次）.pdf", "detai": "auto"}
            # },
            'content': [
                {
                    "type": "text",
                    "text": "错了 继续回答"
                },
                # {
                #     "type": "file_url",
                #     "file_url": {"url": "https://oss.ffire.cc/files/招标文件备案表（第二次）.pdf", "detai": "auto"}
                # }
            ]
        }
    ]

    for message in messages:
        print(ChatCompletionMessageParam(**message))

