xai-grok 0.1.0

//! The async [`GrokClient`] and its streaming machinery.

use std::pin::Pin;
use std::task::{Context, Poll};

use futures_core::Stream;
use futures_util::StreamExt;

use stream_rs::accumulators::openai::OpenAiAccumulator;
use stream_rs::sse::SseEvent;
use stream_rs::stream::SseStream;

use crate::error::{Error, Result};
use crate::types::{ChatChunk, ChatRequest, ChatResponse};

/// The default xAI API base URL.
pub const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1";

/// A high-level async client for the xAI Grok API.
///
/// The client is `OpenAI`-wire-compatible (xAI hosts an `OpenAI`-compatible API at
/// `https://api.x.ai/v1`) and uses [`stream-rs`](stream_rs) for byte-accurate
/// SSE parsing of streamed completions — so chunk boundaries, multi-line
/// `data:` fields, and interleaved tool-call deltas are handled correctly
/// regardless of how the bytes arrive off the socket.
///
/// Clone is cheap (`reqwest::Client` is internally reference-counted).
#[derive(Debug, Clone)]
pub struct GrokClient {
    http: reqwest::Client,
    base_url: String,
    api_key: String,
}

impl GrokClient {
    /// Build a client from an explicit API key, using the default base URL.
    pub fn new(api_key: impl Into<String>) -> Result<Self> {
        Self::builder().api_key(api_key).build()
    }

    /// Build a client, reading the key from the `XAI_API_KEY` environment
    /// variable.
    ///
    /// # Errors
    /// Returns [`Error::MissingApiKey`] if the variable is unset or empty.
    pub fn from_env() -> Result<Self> {
        let key = std::env::var("XAI_API_KEY").map_err(|_| Error::MissingApiKey)?;
        Self::new(key)
    }

    /// Start configuring a client.
    #[must_use]
    pub fn builder() -> GrokClientBuilder {
        GrokClientBuilder::default()
    }

    /// Send a non-streaming chat completion request.
    ///
    /// The `stream` flag on the request is forced to `false`.
    ///
    /// # Errors
    /// Fails on transport errors, non-2xx responses, or JSON decode errors.
    pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatResponse> {
        request.stream = Some(false);
        let resp = self.post("/chat/completions", &request).await?;
        let resp = error_for_status(resp).await?;
        Ok(resp.json::<ChatResponse>().await?)
    }

    /// Send a streaming chat completion request.
    ///
    /// Returns a [`ChatStream`] that yields one [`StreamEvent`] per parsed
    /// `chat.completion.chunk`. The `[DONE]` sentinel terminates the stream and
    /// is not surfaced as an event.
    ///
    /// # Errors
    /// Fails up-front on transport errors or a non-2xx status; per-chunk decode
    /// errors surface as `Err` items while polling the stream.
    pub async fn chat_stream(&self, mut request: ChatRequest) -> Result<ChatStream> {
        request.stream = Some(true);
        let resp = self.post("/chat/completions", &request).await?;
        let resp = error_for_status(resp).await?;

        // reqwest's byte stream yields `Result<Bytes, reqwest::Error>`, which is
        // exactly the shape `stream-rs`'s `SseStream` adapter consumes. It
        // drives an `SseParser` across chunk boundaries with no per-chunk
        // allocation beyond the events themselves.
        let byte_stream: ByteStream = resp.bytes_stream().boxed();
        let sse = SseStream::new(byte_stream);
        Ok(ChatStream {
            inner: Box::pin(sse),
            done: false,
        })
    }

    /// Convenience: run a streaming request but fold every delta back into the
    /// final [`ChatResponse`]-like content using `stream-rs`'s
    /// [`OpenAiAccumulator`]. The closure `on_token` is called with each text
    /// fragment as it arrives, so you can render live while still getting the
    /// assembled result.
    ///
    /// # Errors
    /// Propagates stream and decode errors.
    pub async fn chat_stream_collect<F>(
        &self,
        request: ChatRequest,
        mut on_token: F,
    ) -> Result<OpenAiAccumulator>
    where
        F: FnMut(&str),
    {
        let mut stream = self.chat_stream(request).await?;
        let mut acc = OpenAiAccumulator::new();

        while let Some(event) = stream.next().await {
            let chunk = event?;
            for choice in &chunk.choices {
                let i = choice.index as usize;
                if let Some(role) = &choice.delta.role {
                    acc.push_role(i, role);
                }
                if let Some(content) = &choice.delta.content {
                    on_token(content);
                    acc.push_content(i, content);
                }
                if let Some(tool_calls) = &choice.delta.tool_calls {
                    for tc in tool_calls {
                        let (name, args) = tc.function.as_ref().map_or((None, None), |f| {
                            (f.name.as_deref(), f.arguments.as_deref())
                        });
                        acc.push_tool_call(i, tc.index as usize, tc.id.as_deref(), name, args);
                    }
                }
                if let Some(reason) = &choice.finish_reason {
                    acc.set_finish_reason(i, reason);
                }
            }
        }
        Ok(acc)
    }

    async fn post<T: serde::Serialize>(&self, path: &str, body: &T) -> Result<reqwest::Response> {
        let url = format!("{}{}", self.base_url, path);
        Ok(self
            .http
            .post(url)
            .bearer_auth(&self.api_key)
            .json(body)
            .send()
            .await?)
    }
}

async fn error_for_status(resp: reqwest::Response) -> Result<reqwest::Response> {
    let status = resp.status();
    if status.is_success() {
        return Ok(resp);
    }
    let body = resp.text().await.unwrap_or_default();
    Err(Error::Api {
        status: status.as_u16(),
        body,
    })
}

/// Builder for [`GrokClient`].
#[derive(Debug, Default)]
pub struct GrokClientBuilder {
    api_key: Option<String>,
    base_url: Option<String>,
    http: Option<reqwest::Client>,
}

impl GrokClientBuilder {
    /// Set the API key.
    #[must_use]
    pub fn api_key(mut self, key: impl Into<String>) -> Self {
        self.api_key = Some(key.into());
        self
    }

    /// Override the base URL (e.g. for a proxy or compatible gateway).
    #[must_use]
    pub fn base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = Some(url.into());
        self
    }

    /// Provide a preconfigured `reqwest::Client` (timeouts, proxy, etc.).
    #[must_use]
    pub fn http_client(mut self, client: reqwest::Client) -> Self {
        self.http = Some(client);
        self
    }

    /// Finish building.
    ///
    /// # Errors
    /// Returns [`Error::MissingApiKey`] if no non-empty key was set.
    pub fn build(self) -> Result<GrokClient> {
        let api_key = self
            .api_key
            .filter(|k| !k.is_empty())
            .ok_or(Error::MissingApiKey)?;
        Ok(GrokClient {
            http: self.http.unwrap_or_default(),
            base_url: self
                .base_url
                .unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
            api_key,
        })
    }
}

/// The boxed byte stream feeding the SSE parser.
type ByteStream =
    futures_util::stream::BoxStream<'static, std::result::Result<bytes::Bytes, reqwest::Error>>;

/// A stream of parsed Grok completion chunks.
///
/// Wraps `stream-rs`'s SSE adapter, decoding each event's `data` payload into a
/// [`ChatChunk`] and stopping at the `[DONE]` sentinel.
pub struct ChatStream {
    inner: Pin<Box<SseStream<ByteStream, bytes::Bytes, reqwest::Error>>>,
    done: bool,
}

impl Stream for ChatStream {
    type Item = Result<ChatChunk>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            if self.done {
                return Poll::Ready(None);
            }
            match self.inner.as_mut().poll_next(cx) {
                Poll::Ready(Some(Ok(event))) => {
                    if let Some(item) = decode_event(&event) {
                        match item {
                            DecodedEvent::Done => {
                                self.done = true;
                                return Poll::Ready(None);
                            }
                            DecodedEvent::Chunk(result) => return Poll::Ready(Some(result)),
                        }
                    }
                    // Empty / comment-only event: keep polling.
                }
                Poll::Ready(Some(Err(e))) => {
                    self.done = true;
                    return Poll::Ready(Some(Err(Error::Transport(e))));
                }
                Poll::Ready(None) => {
                    self.done = true;
                    return Poll::Ready(None);
                }
                Poll::Pending => return Poll::Pending,
            }
        }
    }
}

enum DecodedEvent {
    Chunk(Result<ChatChunk>),
    Done,
}

fn decode_event(event: &SseEvent) -> Option<DecodedEvent> {
    let data = event.data.trim();
    if data.is_empty() {
        return None;
    }
    if data == "[DONE]" {
        return Some(DecodedEvent::Done);
    }
    Some(DecodedEvent::Chunk(
        serde_json::from_str::<ChatChunk>(data).map_err(Error::Decode),
    ))
}