"protocol for ark"

from __future__ import annotations

from enum import Enum
from typing import (
    Any,
    Coroutine,
    Dict,
    List,
    Optional,
    Protocol,
    Tuple,
    Union,
)

import volcenginesdkarkruntime.types.chat.chat_completion_chunk as completion_chunk
from pydantic import BaseModel, Field, field_validator, model_validator
from typing_extensions import Annotated, Literal
from volcenginesdkarkruntime.types.chat.chat_completion import ChatCompletion, Choice
from volcenginesdkarkruntime.types.chat.chat_completion_stream_options_param import (
    ChatCompletionStreamOptionsParam,
)
from volcenginesdkarkruntime.types.completion_usage import CompletionUsage

from ark.core.idl.common_protocol import (
    ActionDetails,
    ActionUsage,
    BotUsage,
    ChatUsage,
    Error,
    LabelLogprobosValue,
    ModelDesc,
    Reference,
    Request,
    Response,
)
from ark.core.utils.errorsv3 import InvalidParameter, MissingParameter


class CallableFunction(Protocol):
    def __call__(
        self, request: Dict[str, Any], **kwargs: Any
    ) -> Coroutine[
        Any,
        Any,
        Union[Union[str, BaseModel], Tuple[Union[str, BaseModel], ActionDetails]],
    ]:
        ...


EndpointId = str


class FunctionCallMode(str, Enum):
    SEQUENTIAL = "sequential"
    """
    The function will be called sequentially, one after the other.
    """
    PARALLEL = "parallel"
    """
    The function will be called in parallel, and the results will be combined at the
    end.
    """
    ONCE = "ONCE"
    """
    The function will be called once and returns the result.
    """


class FunctionDefinition(BaseModel):
    name: str
    """The name of the function to be called.

    Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length
    of 64.
    """

    description: str
    """
    A description of what the function does, used by the model to choose when and
    how to call the function.
    """

    parameters: Dict[str, Any]
    """The parameters the functions accepts, described as a JSON Schema object.

    See the
    [JSON Schema reference](https://json-schema.org/understanding-json-schema/) for
    documentation about the format.

    Omitting `parameters` defines a function with an empty parameter list.
    """


class ChatCompletionTool(BaseModel):
    function: FunctionDefinition

    type: Literal["function"]
    """The type of the tool. Currently, only `function` is supported."""


class ArkChatParameters(BaseModel):
    frequency_penalty: Optional[float] = None
    """
    Number between -2.0 and 2.0. Positive values penalize new tokens based on their
    existing frequency in the text so far, decreasing the model's likelihood to
    repeat the same line verbatim.
    """
    logit_bias: Optional[Dict[str, int]] = None
    """
    Modify the likelihood of specified tokens appearing in the completion.
    """
    logprobs: Optional[bool] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = None
    stop: Optional[Union[Optional[str], List[str]]] = None
    stream_options: Optional[ChatCompletionStreamOptionsParam] = None
    temperature: Optional[float] = None
    tools: Optional[List[ChatCompletionTool]] = None
    top_logprobs: Optional[int] = None
    top_p: Optional[float] = None

    n: Optional[int] = Field(default=1, ge=1, le=5)
    """
    How many chat completion choices to generate for each input message. 
    Note that you will be charged based on the number of generated tokens 
    across all of the choices. 

    Keep n as 1 to minimize costs.
    """

    def merge_from(
        self, other: Union[Dict[str, Any], ArkChatParameters]
    ) -> ArkChatParameters:
        other_dict = other if isinstance(other, dict) else other.dict()
        merged_dict = self.dict()

        for key in merged_dict.keys():
            value = other_dict.get(key, None)
            if value is not None:
                merged_dict[key] = value

        return self.__class__(**merged_dict)

    def merge_to(
        self, other: Union[Dict[str, Any], ArkChatParameters]
    ) -> ArkChatParameters:
        self_dict = self.dict()
        merged_dict = other if isinstance(other, dict) else other.dict()

        for key in self_dict.keys():
            value = self_dict.get(key, None)
            if value is not None:
                merged_dict[key] = value

        return self.__class__(**merged_dict)


class Function(BaseModel):
    arguments: str
    """
    The arguments to call the function with, as generated by the model in JSON
    format. Note that the model does not always generate valid JSON, and may
    hallucinate parameters not defined by your function schema. Validate the
    arguments in your code before calling your function.
    """

    name: str
    """The name of the function to call."""


class ChatCompletionMessageToolCallParam(BaseModel):
    id: str
    """The ID of the tool call."""

    function: Function
    """The function that the model called."""

    type: Literal["function"]
    """The type of the tool. Currently, only `function` is supported."""


class ChatCompletionMessageImageUrlPartImageUrl(BaseModel):
    url: str
    detail: Optional[str] = None


class ChatCompletionMessageImageUrlPart(BaseModel):
    type: Literal["image_url"]
    image_url: ChatCompletionMessageImageUrlPartImageUrl


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


ChatCompletionMessagePart = Annotated[
    Union[ChatCompletionMessageImageUrlPart, ChatCompletionMessageTextPart],
    Field(discriminator="type"),
]


class ArkMessage(BaseModel):
    role: Literal["user", "system", "assistant", "tool"]
    content: Union[str, List[ChatCompletionMessagePart]]
    name: Optional[str] = None
    tool_call_id: Optional[str] = None
    tool_calls: Optional[List[ChatCompletionMessageToolCallParam]] = None

    @classmethod
    @model_validator(mode="before")
    def validate_content(cls, v: Dict[str, Any]) -> Dict[str, Any]:
        role, content = v.get("role"), v.get("content")
        if not isinstance(content, str) and role != "user":
            raise InvalidParameter(
                parameter="content",
                cause=f"content must be type of str when role is {role}",
            )

        tool_call_id = v.get("tool_call_id")
        if tool_call_id is not None and role != "tool":
            raise InvalidParameter(
                parameter="tool_call_id",
                cause=f"tool_call_id must be None when role is {role}",
            )

        tool_calls = v.get("tool_calls")
        if tool_calls is not None and role != "assistant":
            raise InvalidParameter(
                parameter="tool_calls",
                cause=f"tool_calls must be None when role is {role}",
            )
        return v


class ArkChatRequest(Request):
    messages: List[ArkMessage]
    """
    chat messages 
    """
    model: str
    """
    bot id
    """
    frequency_penalty: Optional[float] = None
    """
    Number between -2.0 and 2.0. Positive values penalize new tokens based on their
    existing frequency in the text so far, decreasing the model's likelihood to
    repeat the same line verbatim.
    """
    logit_bias: Optional[Dict[str, int]] = None
    """
    Modify the likelihood of specified tokens appearing in the completion.
    """
    logprobs: Optional[bool] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = None
    stop: Optional[Union[Optional[str], List[str]]] = None
    stream: bool = False
    stream_options: Optional[ChatCompletionStreamOptionsParam] = None
    temperature: Optional[float] = None
    tools: Optional[List[ChatCompletionTool]] = None
    top_logprobs: Optional[int] = None
    top_p: Optional[float] = None
    user: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)

    n: Optional[int] = Field(default=1, ge=1, le=5)
    """
    How many chat completion choices to generate for each input message. 
    Note that you will be charged based on the number of generated tokens 
    across all of the choices. 

    Keep n as 1 to minimize costs.
    """

    def get_user_info_extra(self) -> Optional[UserInfoExtra]:
        if not self.metadata or not self.metadata.get("user_info"):
            return None
        user_info = self.metadata.get("user_info", "")
        if isinstance(user_info, str):
            return UserInfoExtra.model_validate_json(user_info)
        else:
            return UserInfoExtra.model_validate(user_info)

    def is_emit_intention_signal_extra(self) -> bool:
        if not self.metadata:
            return False
        intention_signal = self.metadata.get("emit_intention_signal_extra", "false")
        if isinstance(intention_signal, bool):
            return intention_signal
        else:
            return intention_signal == "true"

    def get_chat_request(
        self, extra_body: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        dumped_dict = self.model_dump(exclude_unset=True, exclude_none=True)

        extra_body = extra_body or {}
        if "metadata" in dumped_dict:
            extra_body["metadata"] = dumped_dict.pop("metadata")
        dumped_dict["extra_body"] = extra_body
        return dumped_dict

    @classmethod
    @field_validator("messages")
    def validate_messages(
        cls, v: Optional[List[ArkMessage]]
    ) -> Optional[List[ArkMessage]]:
        if not v:
            raise MissingParameter("messages")

        for msg in v:
            if not msg.role:
                raise MissingParameter("role")
            if msg.content is None:
                msg.content = ""

            if msg.role == "tool" and not msg.tool_call_id:
                raise MissingParameter("tool_call_id")

        return v


class ArkChatResponse(Response):
    id: str
    """A unique identifier for the chat completion."""

    choices: List[Choice]
    """A list of chat completion choices.

    Can be more than one if `n` is greater than 1.
    """

    created: int
    """The Unix timestamp (in seconds) of when the chat completion was created."""

    model: str
    """The model used for the chat completion."""

    object: Literal["chat.completion"]
    """The object type, which is always `chat.completion`."""

    usage: Optional[CompletionUsage] = None
    """Usage statistics for the completion request."""

    bot_usage: Optional[BotUsage] = None

    metadata: Optional[Dict[str, Any]] = None

    references: Optional[List[Reference]] = None

    @staticmethod
    def merge(
        responses: List[Union[ArkChatResponse, ChatCompletion]],
    ) -> ArkChatResponse:
        assert len(responses) > 0, "empty responses"

        """
        to ensure the `merged` have attributes `usage`
        handle the responses in reversed order
        """
        merged = ArkChatResponse(**responses[-1].__dict__)
        for resp in reversed(responses[:-1]):
            for i, j in zip(merged.choices, resp.choices):
                if isinstance(i.message.content, str) and isinstance(
                    j.message.content, str
                ):
                    i.message.content = j.message.content + i.message.content
                elif isinstance(i.message.content, list) and isinstance(
                    j.message.content, list
                ):
                    i.message.content = j.message.content + i.message.content
                else:
                    raise TypeError("no supported merge type")

            if merged.usage and resp.usage:
                merged.usage.prompt_tokens += resp.usage.prompt_tokens
                merged.usage.completion_tokens += resp.usage.completion_tokens
                merged.usage.total_tokens += resp.usage.total_tokens
            if merged.bot_usage and resp.bot_usage:
                merged.bot_usage += resp.bot_usage
        return merged

    def transform_bot_usage(self) -> None:
        """
        used for transforming bot usage without actions
        """

        if not self.usage or self.bot_usage:
            return

        self.bot_usage = BotUsage(
            model_usage=[
                ChatUsage(
                    prompt_tokens=self.usage.prompt_tokens,
                    completion_tokens=self.usage.completion_tokens,
                    total_tokens=self.usage.total_tokens,
                )
            ]
        )
        return

    def merge_usages(
        self, others: Union[CompletionUsage, List[CompletionUsage]]
    ) -> CompletionUsage:
        if not others:
            return self.usage

        if not isinstance(others, list):
            others = [others]

        total_usage = CompletionUsage(
            prompt_tokens=self.usage.prompt_tokens if self.usage else 0,
            completion_tokens=self.usage.completion_tokens if self.usage else 0,
            total_tokens=self.usage.total_tokens if self.usage else 0,
        )
        for usage in others:
            total_usage.prompt_tokens += usage.prompt_tokens
            total_usage.completion_tokens += usage.completion_tokens
            total_usage.total_tokens += usage.total_tokens

        self.usage = total_usage
        return total_usage

    def merge_bot_usage(self, others: Union[Any, List[Any]]) -> Optional[BotUsage]:
        if not others:
            return self.bot_usage

        if not self.bot_usage:
            self.bot_usage = BotUsage()

        if not isinstance(others, list):
            others = [others]

        for other in others:
            if isinstance(other, ActionDetails):
                if self.bot_usage.action_details is None:
                    self.bot_usage.action_details = []
                self.bot_usage.action_details.append(other)

                if self.bot_usage.action_usage is None:
                    self.bot_usage.action_usage = []
                self.bot_usage.action_usage.append(
                    ActionUsage(action_name=other.name, count=other.count)
                )
            elif isinstance(other, (ChatUsage, CompletionUsage)):
                if self.bot_usage.model_usage is None:
                    self.bot_usage.model_usage = []
                self.bot_usage.model_usage.append(other)
            elif isinstance(other, BotUsage):
                self.bot_usage += other
            elif (
                isinstance(other, tuple)
                and len(other) == 2
                and isinstance(other[1], ActionDetails)
            ):
                response, action_detail = other
                # deprecated
                if isinstance(response, ArkSearchResponse) and response.search_usage:
                    if self.bot_usage.action_usage is None:
                        self.bot_usage.action_usage = []
                    self.bot_usage.action_usage.extend(response.search_usage)

                self.merge_bot_usage(action_detail)
        return self.bot_usage


class ArkChatCompletionChunk(Response):
    id: str
    """A unique identifier for the chat completion. Each chunk has the same ID."""

    choices: List[completion_chunk.Choice]
    """A list of chat completion choices.

    Can be more than one if `n` is greater than 1.
    """

    created: int
    """The Unix timestamp (in seconds) of when the chat completion was created.

    Each chunk has the same timestamp.
    """

    model: str
    """The model to generate the completion."""

    object: Literal["chat.completion.chunk"]
    """The object type, which is always `chat.completion.chunk`."""

    usage: Optional[CompletionUsage] = None
    """
    An optional field that will only be present when you set
    `stream_options: {"include_usage": true}` in your request. When present, it
    contains a null value except for the last chunk which contains the token usage
    statistics for the entire request.
    """

    bot_usage: Optional[BotUsage] = None

    metadata: Optional[Dict[str, Any]] = None

    references: Optional[List[Reference]] = None

    def transform_bot_usage(self) -> None:
        """
        used for transforming bot usage without actions
        """

        if not self.usage or (
            self.bot_usage
            and self.bot_usage.model_usage
            and len(self.bot_usage.model_usage) > 0
            and self.bot_usage.model_usage[0].total_tokens
        ):
            return
        if self.bot_usage is None:
            self.bot_usage = BotUsage()
        self.bot_usage.model_usage = [
            ChatUsage(
                prompt_tokens=self.usage.prompt_tokens,
                completion_tokens=self.usage.completion_tokens,
                total_tokens=self.usage.total_tokens,
            )
        ]
        return

    @staticmethod
    def merge(
        responses: List[
            Union[ArkChatCompletionChunk, completion_chunk.ChatCompletionChunk]
        ],
    ) -> Optional[ArkChatCompletionChunk]:
        if len(responses) == 0:
            return None

        """
        to ensure the `merged` have attributes `usage`
        handle the responses in reversed order
        """
        merged = ArkChatCompletionChunk(**responses[-1].__dict__)
        if not merged.bot_usage:
            merged.bot_usage = BotUsage(
                model_usage=[], action_usage=[], action_details=[]
            )
        for resp in reversed(responses[:-1]):
            for i, j in zip(merged.choices, resp.choices):
                if isinstance(i.delta.content, str) and isinstance(
                    j.delta.content, str
                ):
                    i.delta.content = j.delta.content + i.delta.content
                elif isinstance(i.delta.content, list) and isinstance(
                    j.delta.content, list
                ):
                    i.delta.content = j.delta.content + i.delta.content
                else:
                    raise TypeError("no supported merge type")

            if merged.usage and resp.usage:
                merged.usage.prompt_tokens += resp.usage.prompt_tokens
                merged.usage.completion_tokens += resp.usage.completion_tokens
                merged.usage.total_tokens += resp.usage.total_tokens
            if (
                merged.bot_usage
                and isinstance(resp, ArkChatCompletionChunk)
                and resp.bot_usage
            ):
                merged.bot_usage += resp.bot_usage
        return merged

    def merge_usages(
        self, others: Union[CompletionUsage, List[CompletionUsage]]
    ) -> CompletionUsage:
        if not others:
            return self.usage

        if not isinstance(others, list):
            others = [others]

        total_usage = CompletionUsage(
            prompt_tokens=self.usage.prompt_tokens if self.usage else 0,
            completion_tokens=self.usage.completion_tokens if self.usage else 0,
            total_tokens=self.usage.total_tokens if self.usage else 0,
        )
        for usage in others:
            total_usage.prompt_tokens += usage.prompt_tokens
            total_usage.completion_tokens += usage.completion_tokens
            total_usage.total_tokens += usage.total_tokens

        self.usage = total_usage
        return total_usage

    def merge_bot_usage(self, others: Union[Any, List[Any]]) -> Optional[BotUsage]:
        if not others:
            return self.bot_usage

        if not self.bot_usage:
            self.bot_usage = BotUsage()

        if not isinstance(others, list):
            others = [others]

        for other in others:
            if isinstance(other, ActionDetails):
                if self.bot_usage.action_details is None:
                    self.bot_usage.action_details = []
                self.bot_usage.action_details.append(other)

                if self.bot_usage.action_usage is None:
                    self.bot_usage.action_usage = []
                self.bot_usage.action_usage.append(
                    ActionUsage(action_name=other.name, count=other.count)
                )
            elif isinstance(other, (ChatUsage, CompletionUsage)):
                if self.bot_usage.model_usage is None:
                    self.bot_usage.model_usage = []
                self.bot_usage.model_usage.append(other)
            elif isinstance(other, BotUsage):
                self.bot_usage += other
            elif (
                isinstance(other, tuple)
                and len(other) == 2
                and isinstance(other[1], ActionDetails)
            ):
                response, action_detail = other
                # deprecated
                if isinstance(response, ArkSearchResponse) and response.search_usage:
                    if self.bot_usage.action_usage is None:
                        self.bot_usage.action_usage = []
                    self.bot_usage.action_usage.extend(response.search_usage)

                self.merge_bot_usage(action_detail)
        return self.bot_usage


class ArkClassificationRequest(Request):
    query: Optional[str] = None
    labels: List[str]


class ArkClassificationResponse(Response):
    label: Optional[str] = None
    label_logprobos: Optional[Dict[str, LabelLogprobosValue]] = None
    usage: Optional[ChatUsage] = None
    error: Optional[Error] = None


class ArkTokenizeRequest(Request):
    text: Optional[str] = None


class ArkTokenizeResponse(Response):
    total_tokens: Optional[int] = None
    tokens: Optional[List[str]] = None
    error: Optional[Error] = None


class ArkEmbeddingsRequest(Request):
    input: Union[str, List[str]]
    encoding_format: Literal["float", "base64"] = "float"
    user: Optional[str] = None


class UserInfoExtra(BaseModel):
    city: str = ""
    district: str = ""


class ArkActionMeta(Request):
    action_name: str = ""
    model: Optional[str] = None
    function_definition: Optional[FunctionDefinition] = None


class ArkToolMeta(Request):
    action_name: str = ""
    tool_name: str = ""
    model: Optional[str] = None


SourceType = Literal[
    "weather",
    "douyin_baike",
    "toutiao_article",
    "douyin_short_video",
    "search_engine",
    "xigua_feed_video",
    "toutiao_short_content",
]

SOURCE_TAG = {
    "search_engine": "搜索引擎",
    "weather": "墨迹天气",
    "toutiao_article": "头条图文",
    "douyin_short_video": "抖音视频",
    "douyin_baike": "抖音百科",
    "xigua_feed_video": "西瓜视频",
    "toutiao_short_content": "微头条",
}

SearchAction = Literal[
    "WebBrowsing", "WebBrowsingPlus", "WebBrowsingPro", "content_plugin"
]


class ArkSearchSourceType(str, Enum):
    DEFAULT_SEARCH = "search_engine"
    WEATHER = "weather"
    TOUTIAO_ARTICLE = "toutiao_article"
    DOUYIN_SHORT_VIDEO = "douyin_short_video"
    DOUYIN_BAIKE = "douyin_baike"
    XIGUA_FEED_VIDEO = "xigua_feed_video"
    TOUTIAO_SHORT_CONTENT = "toutiao_short_content"


class ArkSearchRequest(Request):
    question: str = ""
    action_name: Literal["WebBrowsingPlus", "WebBrowsingPro", "content_plugin"] = Field(
        default="WebBrowsingPlus"
    )
    source_type: Optional[List[SourceType]] = None
    summary_top_k: int = 10
    utm_source: Optional[str] = None
    keywords: Optional[List[str]] = None
    user_info: Optional[UserInfoExtra] = None
    enhance_param: Optional[Dict[str, Any]] = None


class ArkKnowledgeBaseRequest(Request):
    collection_name: str
    project: str = "default"
    question: str = ""
    retrieve_count: int = 5
    keywords: Optional[List[str]] = None
    rerank_retrieve_count: Optional[int] = None
    query_param: Optional[Dict[str, Any]] = None
    dense_weight: Optional[float] = 0.5
    rerank_switch: bool = False
    show_reference: bool = True


class ArkSearchResponse(Response):
    search_results: List[str] = Field(default_factory=list)
    references: List[Reference] = Field(default_factory=list)
    search_usage: Optional[List[ActionUsage]] = None


class ArkKnowledgeBaseResponse(Response):
    texts: List[str] = Field(default_factory=list)
    references: List[Reference] = Field(default_factory=list)


class ArkSearchIntentionMetadata(ArkActionMeta):
    action_name: SearchAction = Field(default="WebBrowsing")
    result_mapping: Optional[Dict[str, bool]] = None
    keywords: Optional[List[str]] = None


class ArkSearchSummaryMetadata(ArkActionMeta):
    action_name: SearchAction = Field(default="WebBrowsing")
    keywords: Optional[List[str]] = None
    summary_top_k: int = 10
    summary_max_total_len: Optional[int] = None

    biz_id: Literal["seed", "volc"] = "seed"
    extra_params: Optional[Dict[str, Any]] = None
    ab_params: Optional[Dict[str, Any]] = None

    browsing_doc_list: Optional[str] = Field(default="")
    summary_with_system_prompt: Optional[bool] = Field(default=False)
    user_info: Optional[UserInfoExtra] = None
    site: Optional[str] = Field(default="")
    allow_host: Optional[List[str]] = None
    block_host: Optional[List[str]] = None
    publish_time: Dict[str, str] = Field(default_factory=dict)


class ArkCalculatorRequest(Request):
    input: str = ""
    query: Optional[str] = None


class ArkCalculatorResponse(Response):
    result: str = ""
    message: Optional[str] = None


class ArkWebData(BaseModel):
    url: str = ""
    title: str = ""
    content: str = ""
    error_message: Optional[str] = None
    error_code: Optional[str] = None


class ArkLinkReaderRequest(Request):
    url_list: List[str] = Field(default_factory=list)


class ArkLinkReaderResponse(Response):
    ark_web_data_list: List[ArkWebData] = Field(default_factory=list)


class ArkCharacterConfig(BaseModel):
    # character
    name: str = Field(
        default="",
        description="character name, should be unique in single group chat config",
    )
    system_prompt: str = Field(default="", description="character description")
    model_desc: Optional[ModelDesc] = None


class ArkGroupChatConfig(BaseModel):
    # group chat
    description: Optional[str] = Field(
        default="", description="group chat scene description"
    )
    characters: List[ArkCharacterConfig] = Field(
        default=[], description="group chat character configs"
    )
    user_name: Optional[str] = Field(
        default="", description="user character name, optional"
    )


class ArkToolRequest(ArkToolMeta):
    dry_run: Optional[bool] = False
    timeout: Optional[int] = 60
    parameters: Optional[Dict[str, Any]] = None


class ArkToolResponse(Response):
    status_code: Optional[int] = None
    data: Optional[Any] = None
