nah_chat 0.8.0

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

//!
//! # Introduction
//! This crate exposes an async stream API for the widely-used OpenAI
//! [chat completion API](https://platform.openai.com/docs/api-reference/chat) and the
//! [Responses API](https://platform.openai.com/docs/api-reference/responses).
//!
//! Supported features:
//! * Stream generation
//! * Tool calls
//! * Reasoning content (Qwen3, Deepseek R1, etc)
//! * Token usage in the chat completion stream (via `stream_options.include_usage`)
//! * Responses API (stream + non-stream, tool calls, reasoning)
//!
//! This crate is built on top of `tokio`, `reqwest` and `serde_json`.
//!
//! ```rust
//! use nah_chat::{ChatClient, ChatCompletionStreamEvent, ChatMessage};
//! use futures_util::{pin_mut, StreamExt};
//!
//! # async fn make_request_example() {
//! # let base_url = "http://localhost:8080".to_string();
//! # let auth_token = None;
//! # let model_name = "deepseek-r1";
//! # let messages = vec![];
//! # let params = std::collections::HashMap::new();
//!
//! let chat_client = ChatClient::init(base_url, auth_token);
//!
//! // create and pin the stream
//! let stream = chat_client
//!        .chat_completion_stream(model_name, &messages, &params)
//!        .await
//!        .unwrap();
//! pin_mut!(stream);
//!
//! // buffer for the new message
//! let mut message = ChatMessage::new();
//!
//! // consume the stream
//! while let Some(event_result) = stream.next().await {
//!   match event_result {
//!     Ok(ChatCompletionStreamEvent::Delta(delta)) => {
//!       message.apply_model_response_chunk(delta);
//!     }
//!     Ok(ChatCompletionStreamEvent::Usage(usage)) => {
//!       // Optional: the final chunk carries the authoritative token usage.
//!       eprintln!("Usage: {} prompt + {} completion tokens",
//!                 usage.prompt_tokens.unwrap_or(0), usage.completion_tokens.unwrap_or(0));
//!     }
//!     Err(e) => {
//!       eprintln!("Error occurred while processing the chat completion: {}", e);
//!     }
//!   }
//! }
//! # }
//! ```
//! # Notice
//! Copyright 2025, [Mengxiao Lin](linmx0130@gmail.com).
//! This is a part of [nah](https://github.com/linmx0130/nah) project. `nah` means "*N*ot *A*
//! *H*uman". Source code is available under [MPL-2.0](https://mozilla.org/MPL/2.0/).
//!
mod error;
mod message;
mod responses;
pub use error::{Error, ErrorKind, Result};
pub use message::*;
pub use responses::*;

use async_stream::stream;
use futures_core::stream::Stream;
use serde_json::{Number, Value, json};

use responses::SseBuffer;
/**
 * A builder for creating parameters for chat completion requests.
 */
#[derive(Debug, Clone)]
pub struct ChatCompletionParamsBuilder {
  data: std::collections::HashMap<String, Value>,
}

impl ChatCompletionParamsBuilder {
  /**
   * Initialize a [ChatCompletionParamsBuilder] object.
   */
  pub fn new() -> Self {
    ChatCompletionParamsBuilder {
      data: std::collections::HashMap::new(),
    }
  }

  /**
   * Consume the data builder to get a hash map of the parameters for chat completion requests.
   */
  pub fn build(self) -> std::collections::HashMap<String, Value> {
    self.data
  }

  /**
   * Set max token parameter.
   */
  pub fn max_tokens(&mut self, n: usize) -> &mut Self {
    self.data.insert(
      "max_tokens".to_owned(),
      Value::Number(Number::from_u128(n as u128).unwrap()),
    );
    self
  }

  /**
   * Set temperature parameter.
   */
  pub fn temperature(&mut self, t: f64) -> &mut Self {
    self.data.insert(
      "temperature".to_owned(),
      Value::Number(Number::from_f64(t).unwrap()),
    );
    self
  }

  /**
   * Set top_p parameter.
   */
  pub fn top_p(&mut self, p: f64) -> &mut Self {
    self.data.insert(
      "top_p".to_owned(),
      Value::Number(Number::from_f64(p).unwrap()),
    );
    self
  }

  /**
   * Set frequency_penalty parameter.
   */
  pub fn frequency_penalty(&mut self, p: f64) -> &mut Self {
    self.data.insert(
      "frequency_penalty".to_owned(),
      Value::Number(Number::from_f64(p).unwrap()),
    );
    self
  }

  /**
   * Set a parameter with key of `name` and value of `value`.
   */
  pub fn insert(&mut self, name: &str, value: Value) -> &mut Self {
    self.data.insert(name.to_owned(), value);
    self
  }

  /**
   * Ask the server to include the token usage of the call in the final stream chunk
   * (OpenAI / DeepSeek / vLLM ...). The usage is reported as a
   * [ChatCompletionStreamEvent::Usage] right before `[DONE]`.
   */
  pub fn include_usage(&mut self) -> &mut Self {
    self
      .data
      .insert("stream_options".to_owned(), json!({"include_usage": true}));
    self
  }
}

impl<'a> std::iter::IntoIterator for &'a ChatCompletionParamsBuilder {
  type Item = (&'a String, &'a Value);
  type IntoIter = std::collections::hash_map::Iter<'a, String, Value>;

  fn into_iter(self) -> Self::IntoIter {
    (&self.data).into_iter()
  }
}

/**
 * A parsed event from the chat completion streaming interface.
 *
 * `Delta` carries the incremental assistant content (apply it with
 * [ChatMessage::apply_model_response_chunk]); `Usage` carries the token usage of
 * the whole call and is yielded after the final delta, right before `[DONE]`
 * (only when the request sets `stream_options.include_usage`).
 */
#[derive(Debug, Clone)]
pub enum ChatCompletionStreamEvent {
  Delta(ChatResponseChunkDelta),
  Usage(ChatCompletionUsage),
}

/**
 * The object to hold information about the model server and `reqwest` HTTP client.
 */
#[derive(Debug)]
pub struct ChatClient {
  pub base_url: String,
  pub auth_token: Option<String>,
  pub http_client: reqwest::Client,
}

impl ChatClient {
  /**
   * Create a new ChatClient instance, which hosts the basic information and reqwest client
   * for making the requests
   *
   * Args:
   * * `base_url` Base url of the API server. This URL should NOT end with '/'.
   * * `auth_token` Bearer authentication token. It is often called "API Key".
   */
  pub fn init(base_url: String, auth_token: Option<String>) -> Self {
    let client = reqwest::Client::new();
    ChatClient {
      base_url,
      auth_token,
      http_client: client,
    }
  }

  /**
   * Create a chat completion request.
   *
   * Args:
   * * `model` Name of the model to be called.
   * * `messages` A list of [ChatMessage] as the context.
   * * `is_stream` Whether the request is stream-based.
   * * `params` Other parameters to be sent.
   */
  pub fn create_chat_completion_request<'a, 'b, P, M>(
    &self,
    model: &str,
    messages: M,
    is_stream: bool,
    params: P,
  ) -> reqwest::RequestBuilder
  where
    P: IntoIterator<Item = (&'a String, &'a Value)>,
    M: IntoIterator<Item = &'b ChatMessage>,
  {
    let mut data = json!({
        "model": model.to_owned(),
        "messages": messages.into_iter().collect::<Vec<_>>(),
        "stream": is_stream,
        "n": 1,
    });
    params.into_iter().for_each(|(key, value)| {
      data
        .as_object_mut()
        .and_then(|o| o.insert(key.to_owned(), value.to_owned()));
    });

    let endpoint = format!("{}/chat/completions", self.base_url);

    let mut req = self
      .http_client
      .post(&endpoint)
      .header(reqwest::header::CONTENT_TYPE, "application/json")
      .body(serde_json::to_string(&data).unwrap());
    if self.auth_token.is_some() {
      req = req.bearer_auth(self.auth_token.as_ref().unwrap().as_str());
    }

    req
  }
  /**
   * Request chat completion in the non-stream approach.
   *
   * Args:
   * * `model` Name of the model to be called.
   * * `messages` An list of [ChatMessage] as the context.
   * * `params` Other parameters to be sent.
   */
  pub async fn chat_completion<'a, 'b, P, M>(
    &self,
    model: &str,
    messages: M,
    params: P,
  ) -> Result<ChatMessage>
  where
    P: IntoIterator<Item = (&'a String, &'a Value)>,
    M: IntoIterator<Item = &'b ChatMessage>,
  {
    let req = self.create_chat_completion_request(model, messages, false, params);
    let res = req.send().await?;
    if !res.status().is_success() {
      let code = res.status().as_u16();
      let error_content = res.text().await.unwrap();
      return Err(Error {
        kind: ErrorKind::ModelServerError,
        message: Some(format!(
          "Model server responded with error: HTTP status {}, error message = {}",
          code, error_content
        )),
        cause: None,
      });
    }
    let res_text = res.text().await?;
    let mut response_data: Value = serde_json::from_str(&res_text).map_err(|e| Error {
      kind: ErrorKind::ModelServerError,
      message: Some("Failed to parse model server response".to_string()),
      cause: Some(Box::new(e)),
    })?;
    let mut choices = response_data
      .as_object_mut()
      .and_then(|obj| obj.remove("choices"))
      .ok_or(Error {
        kind: ErrorKind::ModelServerError,
        message: Some("Invalid response format: missing choices field".to_string()),
        cause: None,
      })?;
    let choice: &mut Value = choices
      .as_array_mut()
      .and_then(|choices| choices.first_mut())
      .ok_or(Error {
        kind: ErrorKind::ModelServerError,
        message: Some("Invalid response format: missing choices item".to_string()),
        cause: None,
      })?;

    let message_value: Value = choice
      .as_object_mut()
      .and_then(|choice| choice.remove("message"))
      .ok_or(Error {
        kind: ErrorKind::ModelServerError,
        message: Some("Invalid response format: missing message".to_string()),
        cause: None,
      })?;

    let message: ChatMessage = match serde_json::from_value(message_value) {
      Ok(m) => m,
      Err(e) => {
        return Err(Error {
          kind: ErrorKind::ModelServerError,
          message: Some("Failed to parse message from response".to_string()),
          cause: Some(Box::new(e)),
        });
      }
    };
    Ok(message)
  }

  /**
   * Request chat completion in the async stream approach.
   *
   * The stream yields [ChatCompletionStreamEvent]s. `Delta` events carry the
   * incremental assistant content; when the request sets
   * `stream_options.include_usage` (see [ChatCompletionParamsBuilder::include_usage]),
   * a `Usage` event is yielded after the final delta, right before `[DONE]`.
   *
   * Args:
   * * `model` Name of the model to be called.
   * * `messages` An list of [ChatMessage] as the context.
   * * `params` Other parameters to be sent.
   */
  pub async fn chat_completion_stream<'a, 'b, P, M>(
    &self,
    model: &str,
    messages: M,
    params: P,
  ) -> Result<impl Stream<Item = Result<ChatCompletionStreamEvent>>>
  where
    P: IntoIterator<Item = (&'a String, &'a Value)>,
    M: IntoIterator<Item = &'b ChatMessage>,
  {
    let req = self.create_chat_completion_request(model, messages, true, params);
    let mut res = req.send().await?;

    if !res.status().is_success() {
      let code = res.status().as_u16();
      let error_content = res.text().await.unwrap();
      return Err(Error {
        kind: ErrorKind::ModelServerError,
        message: Some(format!(
          "Model server responded with error: HTTP status {}, error message = {}",
          code, error_content
        )),
        cause: None,
      });
    }

    let stream = stream! {
      let mut buffer = SseBuffer::new();
      let mut reach_done = false;
      while !reach_done {
        let Some(chunk_data) = res.chunk().await? else {
            break;
        };
        for message in buffer.push(&chunk_data) {
          for chunk in self.get_model_response_chunks(&message) {
            match chunk {
              ChatResponseChunk::Delta(d) => {
                yield Ok(ChatCompletionStreamEvent::Delta(d));
              }
              ChatResponseChunk::Usage(u) => {
                yield Ok(ChatCompletionStreamEvent::Usage(u));
              }
              ChatResponseChunk::Done => {
                reach_done = true;
              }
            }
          }
        }
      }
      for message in buffer.finish() {
        for chunk in self.get_model_response_chunks(&message) {
          match chunk {
            ChatResponseChunk::Delta(d) => {
              yield Ok(ChatCompletionStreamEvent::Delta(d));
            }
            ChatResponseChunk::Usage(u) => {
              yield Ok(ChatCompletionStreamEvent::Usage(u));
            }
            _ => {}
          }
        }
      }
    };
    Ok(stream)
  }

  /**
   * Parse one complete SSE message (payload between blank lines) from the stream
   * chat completion API to obtain the chunks it carries.
   *
   * A normal chunk carries `choices[0].delta` and is parsed as a `Delta`; a
   * chunk with a top-level `usage` object (only sent when the request sets
   * `stream_options.include_usage`) is parsed as `Usage`.
   *
   * Some providers (OpenAI) send usage in a dedicated final chunk with empty
   * `choices`; others (DeepSeek) attach it to the *same* chunk as the final
   * delta. Both fields are parsed independently, so neither event is lost
   * regardless of provider. `usage: null` on ordinary chunks never parses.
   */
  fn get_model_response_chunks(&self, message: &str) -> Vec<ChatResponseChunk> {
    let data_str = message
      .lines()
      .find_map(|line| line.strip_prefix("data:"))
      .map(str::trim);
    let Some(data_str) = data_str else {
      return Vec::new();
    };
    if data_str == "[DONE]" {
      return vec![ChatResponseChunk::Done];
    }
    let chunk_value: Value = match serde_json::from_str(data_str) {
      Ok(v) => v,
      Err(_) => return Vec::new(),
    };
    let mut chunks = Vec::new();
    let delta_value = chunk_value
      .as_object()
      .and_then(|chunk| chunk.get("choices"))
      .and_then(|choices_value| choices_value.as_array())
      .and_then(|choices_arr| choices_arr.first())
      .and_then(|choice_value| choice_value.as_object())
      .and_then(|choice_obj| choice_obj.get("delta"));
    if let Some(delta_value) = delta_value
      && let Ok(delta) = serde_json::from_value::<ChatResponseChunkDelta>(delta_value.to_owned())
    {
      chunks.push(ChatResponseChunk::Delta(delta));
    }
    // `usage` is `null` on ordinary chunks; deserializing null into the
    // struct fails, so only a real usage object produces a `Usage` chunk.
    if let Ok(usage) = serde_json::from_value::<ChatCompletionUsage>(chunk_value["usage"].clone()) {
      chunks.push(ChatResponseChunk::Usage(usage));
    }
    chunks
  }
}

#[cfg(test)]
mod tests;