linger-openai-sdk 0.1.1

use crate::error::LingerError;
use crate::stream::{SseEvent, SseStream};
use crate::transport::BodyStream;
use crate::RequestId;
use futures_core::Stream;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
use std::pin::Pin;
use std::task::{Context, Poll};

/// EN: Request body for `POST /v1/completions`.
/// 中文：`POST /v1/completions` 的请求体。
#[derive(Clone, Debug, Serialize, PartialEq)]
#[non_exhaustive]
pub struct CreateCompletionRequest {
    /// EN: Model id used to create the completion.
    /// 中文：用于创建 completion 的模型 ID。
    pub model: String,
    /// EN: Prompt input for the legacy completions endpoint.
    /// 中文：legacy completions 端点的提示输入。
    pub prompt: CompletionPrompt,
    /// EN: Optional server-side candidate count.
    /// 中文：可选的服务端候选数量。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub best_of: Option<u32>,
    /// EN: Whether to echo the prompt with the generated completion.
    /// 中文：是否在生成的 completion 中回显提示。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub echo: Option<bool>,
    /// EN: Optional frequency penalty.
    /// 中文：可选的频率惩罚。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    /// EN: Optional token logit bias map.
    /// 中文：可选的 token logit 偏置映射。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<BTreeMap<String, f32>>,
    /// EN: Optional log probability count.
    /// 中文：可选的 log probability 数量。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logprobs: Option<u32>,
    /// EN: Optional maximum generated token count.
    /// 中文：可选的最大生成 token 数量。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
    /// EN: Optional number of completions to generate.
    /// 中文：可选的 completion 生成数量。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u32>,
    /// EN: Optional presence penalty.
    /// 中文：可选的存在惩罚。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    /// EN: Optional deterministic sampling seed.
    /// 中文：可选的确定性采样种子。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<i64>,
    /// EN: Optional stop sequence or sequences.
    /// 中文：可选的停止序列。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<CompletionStop>,
    /// EN: Optional stream flag set by streaming convenience methods.
    /// 中文：由流式便捷方法设置的可选 stream 标志。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    /// EN: Optional suffix inserted after generated text.
    /// 中文：可选的生成文本后缀。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub suffix: Option<String>,
    /// EN: Optional sampling temperature.
    /// 中文：可选的采样温度。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    /// EN: Optional nucleus sampling value.
    /// 中文：可选的 nucleus sampling 值。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,
    /// EN: Optional end-user identifier.
    /// 中文：可选的终端用户标识。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
    /// EN: Forward-compatible optional fields not yet covered by handwritten types.
    /// 中文：手写类型尚未覆盖的前向兼容可选字段。
    #[serde(flatten)]
    pub extra: BTreeMap<String, Value>,
}

impl CreateCompletionRequest {
    /// EN: Starts building a legacy completion request.
    /// 中文：开始构建 legacy completion 请求。
    pub fn builder() -> CreateCompletionRequestBuilder {
        CreateCompletionRequestBuilder::default()
    }

    pub(crate) fn into_streaming(mut self) -> Self {
        self.stream = Some(true);
        self
    }
}

/// EN: Builder for legacy completion requests.
/// 中文：legacy completion 请求的构建器。
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct CreateCompletionRequestBuilder {
    model: Option<String>,
    prompt: Option<CompletionPrompt>,
    best_of: Option<u32>,
    echo: Option<bool>,
    frequency_penalty: Option<f32>,
    logit_bias: Option<BTreeMap<String, f32>>,
    logprobs: Option<u32>,
    max_tokens: Option<u32>,
    n: Option<u32>,
    presence_penalty: Option<f32>,
    seed: Option<i64>,
    stop: Option<CompletionStop>,
    suffix: Option<String>,
    temperature: Option<f32>,
    top_p: Option<f32>,
    user: Option<String>,
    extra: BTreeMap<String, Value>,
}

impl CreateCompletionRequestBuilder {
    /// EN: Sets the model id.
    /// 中文：设置模型 ID。
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// EN: Sets the prompt.
    /// 中文：设置提示。
    pub fn prompt(mut self, prompt: impl Into<CompletionPrompt>) -> Self {
        self.prompt = Some(prompt.into());
        self
    }

    /// EN: Sets the optional server-side candidate count.
    /// 中文：设置可选的服务端候选数量。
    pub fn best_of(mut self, best_of: u32) -> Self {
        self.best_of = Some(best_of);
        self
    }

    /// EN: Sets whether to echo the prompt.
    /// 中文：设置是否回显提示。
    pub fn echo(mut self, echo: bool) -> Self {
        self.echo = Some(echo);
        self
    }

    /// EN: Sets the frequency penalty.
    /// 中文：设置频率惩罚。
    pub fn frequency_penalty(mut self, frequency_penalty: f32) -> Self {
        self.frequency_penalty = Some(frequency_penalty);
        self
    }

    /// EN: Sets the token logit bias map.
    /// 中文：设置 token logit 偏置映射。
    pub fn logit_bias(mut self, logit_bias: BTreeMap<String, f32>) -> Self {
        self.logit_bias = Some(logit_bias);
        self
    }

    /// EN: Sets the log probability count.
    /// 中文：设置 log probability 数量。
    pub fn logprobs(mut self, logprobs: u32) -> Self {
        self.logprobs = Some(logprobs);
        self
    }

    /// EN: Sets the maximum generated token count.
    /// 中文：设置最大生成 token 数量。
    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    /// EN: Sets the number of completions to generate.
    /// 中文：设置 completion 生成数量。
    pub fn n(mut self, n: u32) -> Self {
        self.n = Some(n);
        self
    }

    /// EN: Sets the presence penalty.
    /// 中文：设置存在惩罚。
    pub fn presence_penalty(mut self, presence_penalty: f32) -> Self {
        self.presence_penalty = Some(presence_penalty);
        self
    }

    /// EN: Sets the deterministic sampling seed.
    /// 中文：设置确定性采样种子。
    pub fn seed(mut self, seed: i64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// EN: Sets stop sequence configuration.
    /// 中文：设置停止序列配置。
    pub fn stop(mut self, stop: impl Into<CompletionStop>) -> Self {
        self.stop = Some(stop.into());
        self
    }

    /// EN: Sets the suffix inserted after generated text.
    /// 中文：设置生成文本后的后缀。
    pub fn suffix(mut self, suffix: impl Into<String>) -> Self {
        self.suffix = Some(suffix.into());
        self
    }

    /// EN: Sets the sampling temperature.
    /// 中文：设置采样温度。
    pub fn temperature(mut self, temperature: f32) -> Self {
        self.temperature = Some(temperature);
        self
    }

    /// EN: Sets the nucleus sampling value.
    /// 中文：设置 nucleus sampling 值。
    pub fn top_p(mut self, top_p: f32) -> Self {
        self.top_p = Some(top_p);
        self
    }

    /// EN: Sets the optional end-user identifier.
    /// 中文：设置可选的终端用户标识。
    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    /// EN: Adds a forward-compatible JSON field.
    /// 中文：添加前向兼容的 JSON 字段。
    pub fn extra(mut self, name: impl Into<String>, value: Value) -> Self {
        self.extra.insert(name.into(), value);
        self
    }

    /// EN: Builds and validates the request.
    /// 中文：构建并校验请求。
    pub fn build(self) -> Result<CreateCompletionRequest, LingerError> {
        let model = self
            .model
            .filter(|value| !value.trim().is_empty())
            .ok_or_else(|| LingerError::invalid_config("model is required"))?;
        let prompt = self
            .prompt
            .ok_or_else(|| LingerError::invalid_config("prompt is required"))?;
        if self.logprobs.is_some_and(|logprobs| logprobs > 5) {
            return Err(LingerError::invalid_config(
                "logprobs must be between 0 and 5",
            ));
        }
        if self
            .frequency_penalty
            .is_some_and(|penalty| !(-2.0..=2.0).contains(&penalty))
        {
            return Err(LingerError::invalid_config(
                "frequency_penalty must be between -2.0 and 2.0",
            ));
        }
        if self
            .presence_penalty
            .is_some_and(|penalty| !(-2.0..=2.0).contains(&penalty))
        {
            return Err(LingerError::invalid_config(
                "presence_penalty must be between -2.0 and 2.0",
            ));
        }
        if self
            .temperature
            .is_some_and(|temperature| !(0.0..=2.0).contains(&temperature))
        {
            return Err(LingerError::invalid_config(
                "temperature must be between 0.0 and 2.0",
            ));
        }
        if self
            .top_p
            .is_some_and(|top_p| !(0.0..=1.0).contains(&top_p))
        {
            return Err(LingerError::invalid_config(
                "top_p must be between 0.0 and 1.0",
            ));
        }
        if self.n.is_some_and(|n| !(1..=128).contains(&n)) {
            return Err(LingerError::invalid_config("n must be between 1 and 128"));
        }
        validate_extra_fields(&self.extra)?;
        Ok(CreateCompletionRequest {
            model,
            prompt,
            best_of: self.best_of,
            echo: self.echo,
            frequency_penalty: self.frequency_penalty,
            logit_bias: self.logit_bias,
            logprobs: self.logprobs,
            max_tokens: self.max_tokens,
            n: self.n,
            presence_penalty: self.presence_penalty,
            seed: self.seed,
            stop: self.stop,
            stream: None,
            suffix: self.suffix,
            temperature: self.temperature,
            top_p: self.top_p,
            user: self.user,
            extra: self.extra,
        })
    }
}

/// EN: Legacy completions prompt value.
/// 中文：legacy completions 的提示值。
#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
#[serde(untagged)]
#[non_exhaustive]
pub enum CompletionPrompt {
    /// EN: A single text prompt.
    /// 中文：单个文本提示。
    Text(String),
    /// EN: Multiple text prompts.
    /// 中文：多个文本提示。
    Texts(Vec<String>),
    /// EN: A single token prompt.
    /// 中文：单个 token 提示。
    Tokens(Vec<i64>),
    /// EN: Multiple token prompts.
    /// 中文：多个 token 提示。
    TokenArrays(Vec<Vec<i64>>),
    /// EN: Explicit JSON null prompt, matching the official API type.
    /// 中文：显式 JSON null 提示，与官方 API 类型一致。
    Null,
}

impl CompletionPrompt {
    /// EN: Creates a token prompt.
    /// 中文：创建 token 提示。
    pub fn tokens<I, T>(tokens: I) -> Self
    where
        I: IntoIterator<Item = T>,
        T: Into<i64>,
    {
        Self::Tokens(tokens.into_iter().map(Into::into).collect())
    }

    /// EN: Creates multiple token prompts.
    /// 中文：创建多个 token 提示。
    pub fn token_arrays<I, J, T>(token_arrays: I) -> Self
    where
        I: IntoIterator<Item = J>,
        J: IntoIterator<Item = T>,
        T: Into<i64>,
    {
        Self::TokenArrays(
            token_arrays
                .into_iter()
                .map(|tokens| tokens.into_iter().map(Into::into).collect())
                .collect(),
        )
    }

    /// EN: Creates multiple text prompts.
    /// 中文：创建多个文本提示。
    pub fn texts<I, T>(texts: I) -> Self
    where
        I: IntoIterator<Item = T>,
        T: Into<String>,
    {
        Self::Texts(texts.into_iter().map(Into::into).collect())
    }
}

impl From<&str> for CompletionPrompt {
    fn from(value: &str) -> Self {
        Self::Text(value.to_string())
    }
}

impl From<String> for CompletionPrompt {
    fn from(value: String) -> Self {
        Self::Text(value)
    }
}

impl From<Vec<String>> for CompletionPrompt {
    fn from(value: Vec<String>) -> Self {
        Self::Texts(value)
    }
}

/// EN: Legacy completions stop sequence value.
/// 中文：legacy completions 的停止序列值。
#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
#[serde(untagged)]
#[non_exhaustive]
pub enum CompletionStop {
    /// EN: A single stop sequence.
    /// 中文：单个停止序列。
    Sequence(String),
    /// EN: Multiple stop sequences.
    /// 中文：多个停止序列。
    Sequences(Vec<String>),
}

impl From<&str> for CompletionStop {
    fn from(value: &str) -> Self {
        Self::Sequence(value.to_string())
    }
}

impl From<String> for CompletionStop {
    fn from(value: String) -> Self {
        Self::Sequence(value)
    }
}

impl From<Vec<String>> for CompletionStop {
    fn from(value: Vec<String>) -> Self {
        Self::Sequences(value)
    }
}

/// EN: Completion response object.
/// 中文：Completion 响应对象。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[non_exhaustive]
pub struct Completion {
    /// EN: Completion id.
    /// 中文：Completion ID。
    pub id: String,
    /// EN: Completion choices returned by the model.
    /// 中文：模型返回的 completion 候选项。
    #[serde(default)]
    pub choices: Vec<CompletionChoice>,
    /// EN: Unix timestamp for creation.
    /// 中文：创建时间的 Unix 时间戳。
    pub created: u64,
    /// EN: Model used for completion.
    /// 中文：用于 completion 的模型。
    pub model: String,
    /// EN: API object type, normally `text_completion`.
    /// 中文：API 对象类型，通常为 `text_completion`。
    pub object: String,
    /// EN: Backend configuration fingerprint, when returned.
    /// 中文：后端配置指纹，如响应中存在。
    #[serde(default)]
    pub system_fingerprint: Option<String>,
    /// EN: Token usage, when returned.
    /// 中文：Token 用量，如响应中存在。
    #[serde(default)]
    pub usage: Option<CompletionUsage>,
    /// EN: Additional fields preserved for forward compatibility.
    /// 中文：为前向兼容保留的额外字段。
    #[serde(flatten)]
    pub extra: BTreeMap<String, Value>,
    /// EN: OpenAI request id from response headers.
    /// 中文：响应头中的 OpenAI 请求 ID。
    #[serde(skip)]
    request_id: Option<RequestId>,
}

impl Completion {
    pub(crate) fn with_request_id(mut self, request_id: Option<RequestId>) -> Self {
        self.request_id = request_id;
        self
    }

    /// EN: Returns the OpenAI request id, when present.
    /// 中文：返回 OpenAI 请求 ID，如存在。
    pub fn request_id(&self) -> Option<&RequestId> {
        self.request_id.as_ref()
    }
}

/// EN: Single completion choice.
/// 中文：单个 completion 候选项。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[non_exhaustive]
pub struct CompletionChoice {
    /// EN: Generated text.
    /// 中文：生成的文本。
    pub text: String,
    /// EN: Choice index.
    /// 中文：候选项索引。
    pub index: u32,
    /// EN: Token log probability details, when requested.
    /// 中文：请求时返回的 token log probability 详情。
    #[serde(default)]
    pub logprobs: Option<CompletionLogprobs>,
    /// EN: Finish reason, when returned.
    /// 中文：结束原因，如响应中存在。
    #[serde(default)]
    pub finish_reason: Option<String>,
    /// EN: Additional fields preserved for forward compatibility.
    /// 中文：为前向兼容保留的额外字段。
    #[serde(flatten)]
    pub extra: BTreeMap<String, Value>,
}

/// EN: Token log probability details for a completion choice.
/// 中文：completion 候选项的 token log probability 详情。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[non_exhaustive]
pub struct CompletionLogprobs {
    /// EN: Text offsets for returned tokens.
    /// 中文：返回 token 的文本偏移量。
    #[serde(default)]
    pub text_offset: Option<Vec<u32>>,
    /// EN: Log probabilities for chosen tokens.
    /// 中文：已选择 token 的 log probability。
    #[serde(default)]
    pub token_logprobs: Option<Vec<f64>>,
    /// EN: Returned tokens.
    /// 中文：返回的 token。
    #[serde(default)]
    pub tokens: Option<Vec<String>>,
    /// EN: Top token log probabilities.
    /// 中文：最高 token log probability。
    #[serde(default)]
    pub top_logprobs: Option<Vec<BTreeMap<String, f64>>>,
    /// EN: Additional fields preserved for forward compatibility.
    /// 中文：为前向兼容保留的额外字段。
    #[serde(flatten)]
    pub extra: BTreeMap<String, Value>,
}

/// EN: Token usage for a completion request.
/// 中文：completion 请求的 token 用量。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[non_exhaustive]
pub struct CompletionUsage {
    /// EN: Generated completion token count.
    /// 中文：生成 completion token 数量。
    pub completion_tokens: u64,
    /// EN: Prompt token count.
    /// 中文：Prompt token 数量。
    pub prompt_tokens: u64,
    /// EN: Total token count.
    /// 中文：总 token 数量。
    pub total_tokens: u64,
    /// EN: Optional completion token detail counts.
    /// 中文：可选的 completion token 明细计数。
    #[serde(default)]
    pub completion_tokens_details: Option<CompletionTokenDetails>,
    /// EN: Optional prompt token detail counts.
    /// 中文：可选的 prompt token 明细计数。
    #[serde(default)]
    pub prompt_tokens_details: Option<CompletionPromptTokenDetails>,
}

/// EN: Completion-side token usage details.
/// 中文：completion 侧 token 用量明细。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[non_exhaustive]
pub struct CompletionTokenDetails {
    /// EN: Accepted prediction token count.
    /// 中文：已接受的预测 token 数量。
    #[serde(default)]
    pub accepted_prediction_tokens: Option<u64>,
    /// EN: Audio token count.
    /// 中文：音频 token 数量。
    #[serde(default)]
    pub audio_tokens: Option<u64>,
    /// EN: Reasoning token count.
    /// 中文：推理 token 数量。
    #[serde(default)]
    pub reasoning_tokens: Option<u64>,
    /// EN: Rejected prediction token count.
    /// 中文：已拒绝的预测 token 数量。
    #[serde(default)]
    pub rejected_prediction_tokens: Option<u64>,
}

/// EN: Prompt-side token usage details.
/// 中文：prompt 侧 token 用量明细。
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[non_exhaustive]
pub struct CompletionPromptTokenDetails {
    /// EN: Audio token count.
    /// 中文：音频 token 数量。
    #[serde(default)]
    pub audio_tokens: Option<u64>,
    /// EN: Cached token count.
    /// 中文：缓存 token 数量。
    #[serde(default)]
    pub cached_tokens: Option<u64>,
}

/// EN: Incremental legacy completions stream.
/// 中文：增量 legacy completions 流。
pub struct CompletionStream {
    inner: SseStream,
    done: bool,
}

impl CompletionStream {
    /// EN: Creates a completion stream from an HTTP body stream.
    /// 中文：通过 HTTP 响应体流创建 completion 流。
    pub fn new(body: BodyStream) -> Self {
        Self {
            inner: SseStream::new(body),
            done: false,
        }
    }
}

impl Stream for CompletionStream {
    type Item = Result<CompletionStreamItem, LingerError>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        if this.done {
            return Poll::Ready(None);
        }
        match Pin::new(&mut this.inner).poll_next(cx) {
            Poll::Ready(Some(Ok(raw))) => match parse_completion_event(raw) {
                Ok(Some(item)) => Poll::Ready(Some(Ok(item))),
                Ok(None) => {
                    this.done = true;
                    Poll::Ready(None)
                }
                Err(error) => {
                    this.done = true;
                    Poll::Ready(Some(Err(error)))
                }
            },
            Poll::Ready(Some(Err(error))) => {
                this.done = true;
                Poll::Ready(Some(Err(error)))
            }
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// EN: Parsed legacy completion stream item.
/// 中文：解析后的 legacy completion 流项目。
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct CompletionStreamItem {
    /// EN: Stream event category.
    /// 中文：流事件类别。
    pub event: CompletionStreamEvent,
    /// EN: Completion payload carried by this SSE frame.
    /// 中文：此 SSE 帧携带的 completion 载荷。
    pub completion: Completion,
}

/// EN: Legacy completion stream event category.
/// 中文：legacy completion 流事件类别。
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum CompletionStreamEvent {
    /// EN: A completion payload was received.
    /// 中文：收到 completion 载荷。
    Completion,
}

fn parse_completion_event(raw: SseEvent) -> Result<Option<CompletionStreamItem>, LingerError> {
    if raw.data.trim() == "[DONE]" {
        return Ok(None);
    }
    let completion = serde_json::from_str::<Completion>(&raw.data).map_err(|error| {
        LingerError::streaming(format!("failed to parse completion stream event: {error}"))
    })?;
    Ok(Some(CompletionStreamItem {
        event: CompletionStreamEvent::Completion,
        completion,
    }))
}

fn validate_extra_fields(extra: &BTreeMap<String, Value>) -> Result<(), LingerError> {
    if extra.keys().any(|key| key.trim().is_empty()) {
        return Err(LingerError::invalid_config(
            "extra field names must not be empty",
        ));
    }
    if extra.values().any(Value::is_null) {
        return Err(LingerError::invalid_config(
            "extra field values must not be null",
        ));
    }
    Ok(())
}