llmrix-rust-sdk 1.0.0

use reqwest::{Client, Response, StatusCode};
use serde::{de::DeserializeOwned, Serialize};
use std::time::Duration;

use crate::error::{LlmrixError, Result};
use crate::streaming::{event::StreamEvent, parser::parse_sse_stream};

const API_V1: &str = "/api/open/v1";

/// Low-level HTTP + SSE transport shared across all resources.
pub(crate) struct Transport {
    pub(crate) base_url: String,
    pub(crate) api_key:  String,
    /// Regular client (has timeout).
    client:     Client,
    /// SSE client (no timeout — context controls stream lifetime).
    sse_client: Client,
}

impl Transport {
    pub(crate) fn new(base_url: String, api_key: String, timeout: Duration) -> Result<Self> {
        let client = Client::builder()
            .timeout(timeout)
            .build()
            .map_err(LlmrixError::Transport)?;
        let sse_client = Client::builder()
            // No overall timeout for long-lived streams.
            .build()
            .map_err(LlmrixError::Transport)?;
        Ok(Self { base_url, api_key, client, sse_client })
    }

    fn url(&self, path: &str) -> String {
        format!("{}{}", self.base_url, path)
    }

    fn auth_header(&self) -> Option<(&'static str, String)> {
        if self.api_key.is_empty() {
            None
        } else {
            Some(("Authorization", format!("Bearer {}", self.api_key)))
        }
    }

    // -----------------------------------------------------------------------
    // Regular JSON requests
    // -----------------------------------------------------------------------

    pub(crate) async fn get<T: DeserializeOwned>(&self, path: &str) -> Result<T> {
        let mut req = self.client.get(self.url(path));
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_json(req.send().await?).await
    }

    pub(crate) async fn post<B: Serialize, T: DeserializeOwned>(
        &self, path: &str, body: &B,
    ) -> Result<T> {
        let mut req = self.client.post(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_json(req.send().await?).await
    }

    pub(crate) async fn post_no_body(&self, path: &str) -> Result<()> {
        let mut req = self.client.post(self.url(path)).json(&serde_json::Value::Object(Default::default()));
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_empty(req.send().await?).await
    }

    /// POST with a body; ignore the response body (only check status).
    pub(crate) async fn post_void<B: Serialize>(&self, path: &str, body: &B) -> Result<()> {
        let mut req = self.client.post(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_empty(req.send().await?).await
    }

    pub(crate) async fn patch<B: Serialize, T: DeserializeOwned>(
        &self, path: &str, body: &B,
    ) -> Result<T> {
        let mut req = self.client.patch(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_json(req.send().await?).await
    }

    pub(crate) async fn put<B: Serialize, T: DeserializeOwned>(
        &self, path: &str, body: &B,
    ) -> Result<T> {
        let mut req = self.client.put(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_json(req.send().await?).await
    }

    pub(crate) async fn delete(&self, path: &str) -> Result<()> {
        let mut req = self.client.delete(self.url(path));
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        self.execute_empty(req.send().await?).await
    }

    /// POST to `path`, unwrap the JSON response envelope by extracting `key`.
    pub(crate) async fn post_unwrap<B: Serialize, T: DeserializeOwned>(
        &self, path: &str, body: &B, key: &str,
    ) -> Result<T> {
        let mut req = self.client.post(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        let raw = self.execute_bytes(req.send().await?).await?;
        self.unwrap_key(&raw, key)
    }

    pub(crate) async fn get_unwrap<T: DeserializeOwned>(
        &self, path: &str, key: &str,
    ) -> Result<T> {
        let mut req = self.client.get(self.url(path));
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        let raw = self.execute_bytes(req.send().await?).await?;
        self.unwrap_key(&raw, key)
    }

    pub(crate) async fn patch_unwrap<B: Serialize, T: DeserializeOwned>(
        &self, path: &str, body: &B, key: &str,
    ) -> Result<T> {
        let mut req = self.client.patch(self.url(path)).json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        let raw = self.execute_bytes(req.send().await?).await?;
        self.unwrap_key(&raw, key)
    }

    fn unwrap_key<T: DeserializeOwned>(&self, raw: &[u8], key: &str) -> Result<T> {
        let envelope: serde_json::Map<String, serde_json::Value> =
            serde_json::from_slice(raw)?;
        let sub = envelope
            .get(key)
            .ok_or_else(|| LlmrixError::Other(format!("response missing key \"{}\"", key)))?;
        Ok(serde_json::from_value(sub.clone())?)
    }

    // -----------------------------------------------------------------------
    // SSE streaming
    // -----------------------------------------------------------------------

    pub(crate) async fn stream<B, F>(
        &self, path: &str, body: &B, handler: F,
    ) -> Result<()>
    where
        B: Serialize,
        F: FnMut(&StreamEvent) -> Result<()>,
    {
        let mut req = self.sse_client
            .post(self.url(path))
            .header("Accept", "text/event-stream")
            .header("Cache-Control", "no-cache")
            .json(body);
        if let Some((k, v)) = self.auth_header() {
            req = req.header(k, v);
        }
        let resp = req.send().await.map_err(LlmrixError::Transport)?;
        self.check_status(&resp)?;
        parse_sse_stream(resp.bytes_stream(), handler).await
    }

    // -----------------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------------

    async fn execute_json<T: DeserializeOwned>(&self, resp: Response) -> Result<T> {
        self.check_status(&resp)?;
        Ok(resp.json::<T>().await?)
    }

    async fn execute_bytes(&self, resp: Response) -> Result<Vec<u8>> {
        self.check_status(&resp)?;
        Ok(resp.bytes().await?.to_vec())
    }

    async fn execute_empty(&self, resp: Response) -> Result<()> {
        self.check_status(&resp)?;
        Ok(())
    }

    fn check_status(&self, resp: &Response) -> Result<()> {
        let status = resp.status();
        if status.is_success() {
            return Ok(());
        }
        let code = status.as_u16();
        let msg = format!("server returned HTTP {}", code);
        if status == StatusCode::UNAUTHORIZED || status == StatusCode::SERVICE_UNAVAILABLE {
            Err(LlmrixError::Auth { status: code, message: msg })
        } else {
            Err(LlmrixError::Api { status: code, message: msg, body: String::new() })
        }
    }
}

// ---------------------------------------------------------------------------
// Path helpers
// ---------------------------------------------------------------------------

pub(crate) fn path_conversations() -> String           { format!("{}/conversations", API_V1) }
pub(crate) fn path_conversation(id: &str) -> String    { format!("{}/{}", path_conversations(), id) }
pub(crate) fn path_messages(id: &str) -> String        { format!("{}/messages", path_conversation(id)) }
pub(crate) fn path_chat(id: &str) -> String            { format!("{}/chat", path_conversation(id)) }
pub(crate) fn path_chat_stop(id: &str) -> String       { format!("{}/stop", path_chat(id)) }
pub(crate) fn path_chat_decide(id: &str) -> String     { format!("{}/hitl/decide", path_chat(id)) }
pub(crate) fn path_cron_tasks() -> String              { format!("{}/cron/tasks", API_V1) }
pub(crate) fn path_cron_task(id: &str) -> String       { format!("{}/{}", path_cron_tasks(), id) }
pub(crate) fn path_cron_pause(id: &str) -> String      { format!("{}/pause", path_cron_task(id)) }
pub(crate) fn path_cron_resume(id: &str) -> String     { format!("{}/resume", path_cron_task(id)) }
pub(crate) fn path_agents() -> String                  { format!("{}/agent", API_V1) }
pub(crate) fn path_agent(id: i64) -> String            { format!("{}/{}", path_agents(), id) }
pub(crate) fn path_agent_mates(id: i64) -> String      { format!("{}/mates", path_agent(id)) }