AkurAI-Framework / crates/llm/src/request.rs

Latest change c6179dcf3abbf53272e4279e6740397ad2e75726 - AKURAI-FRAMEWORK-1..9 consolidate framework APIs by Ólafur Búi Ólafsson

//! Endpoint parsing and HTTP/1.1 request building for a streaming
//! `POST /v1/chat/completions`.
//!
//! The pure parts here — [`parse_endpoint`], [`build_body`], [`build_request`]
//! — are unit-tested against hand-written strings; no network is touched.

use akurai_json::Value;

use crate::error::LlmError;
use crate::types::Message;

/// A parsed plain-HTTP endpoint: `host` and `port`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Endpoint {
    pub host: String,
    pub port: u16,
}

/// Parse an endpoint like `http://host:port` (or `http://host`, defaulting to
/// port 80). A trailing path is ignored — we always POST to
/// `/v1/chat/completions`.
///
/// `https://` is rejected with [`LlmError::TlsUnsupported`].
pub fn parse_endpoint(endpoint: &str) -> Result<Endpoint, LlmError> {
    let trimmed = endpoint.trim();
    if akurai_http::endpoint::is_https_endpoint(trimmed) {
        return Err(LlmError::TlsUnsupported(trimmed.to_string()));
    }
    let (host, port) =
        akurai_http::endpoint::parse_plain_endpoint(trimmed).map_err(LlmError::InvalidEndpoint)?;
    Ok(Endpoint { host, port })
}

/// Serialize the chat-completions request body using `akurai_json` so strings
/// are escaped correctly:
/// `{"model": ..., "messages": [{"role","content"}, ...], "stream": true,
///  "tools": [...]?}`.
///
/// `stream` is always `true`. `tools` is omitted entirely when `None` or empty.
/// Keys are emitted in insertion order.
pub fn build_body(model: &str, messages: &[Message], tools: Option<&[Value]>) -> String {
    let msgs = Value::Array(
        messages
            .iter()
            .map(|m| {
                Value::Object(vec![
                    ("role".into(), Value::Str(m.role.clone())),
                    ("content".into(), Value::Str(m.content.clone())),
                ])
            })
            .collect(),
    );
    let mut fields = vec![
        ("model".into(), Value::Str(model.to_string())),
        ("messages".into(), msgs),
        ("stream".into(), Value::Bool(true)),
    ];
    if let Some(tools) = tools {
        if !tools.is_empty() {
            fields.push(("tools".into(), Value::Array(tools.to_vec())));
        }
    }
    Value::Object(fields).to_json()
}

/// Build a complete HTTP/1.1 request string for the streaming
/// `POST /v1/chat/completions`.
///
/// `Content-Length` is the exact UTF-8 byte length of the body. `Accept:
/// text/event-stream` requests the SSE response, and `Connection: close` lets
/// the client read the stream until EOF.
pub fn build_request(
    host: &str,
    port: u16,
    model: &str,
    messages: &[Message],
    tools: Option<&[Value]>,
) -> String {
    let body = build_body(model, messages, tools);
    let host_header = if port == 80 {
        host.to_string()
    } else {
        format!("{host}:{port}")
    };
    format!(
        "POST /v1/chat/completions HTTP/1.1\r\n\
         Host: {host_header}\r\n\
         Content-Type: application/json\r\n\
         Accept: text/event-stream\r\n\
         Content-Length: {len}\r\n\
         Connection: close\r\n\
         \r\n\
         {body}",
        host_header = host_header,
        len = body.len(),
        body = body,
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_host_and_port() {
        let e = parse_endpoint("http://192.168.8.14:8080").unwrap();
        assert_eq!(e.host, "192.168.8.14");
        assert_eq!(e.port, 8080);
    }

    #[test]
    fn defaults_port_80() {
        let e = parse_endpoint("http://localhost").unwrap();
        assert_eq!(e.host, "localhost");
        assert_eq!(e.port, 80);
    }

    #[test]
    fn ignores_trailing_path() {
        let e = parse_endpoint("http://host:9000/v1/chat/completions").unwrap();
        assert_eq!(e.host, "host");
        assert_eq!(e.port, 9000);
    }

    #[test]
    fn rejects_https() {
        assert_eq!(
            parse_endpoint("https://host:443"),
            Err(LlmError::TlsUnsupported("https://host:443".into()))
        );
    }

    #[test]
    fn rejects_unknown_scheme_and_empty_host() {
        assert!(matches!(
            parse_endpoint("ftp://host"),
            Err(LlmError::InvalidEndpoint(_))
        ));
        assert!(matches!(
            parse_endpoint("http://"),
            Err(LlmError::InvalidEndpoint(_))
        ));
    }

    #[test]
    fn body_serializes_messages_in_order_with_stream_true() {
        let msgs = [
            Message::new("system", "be brief"),
            Message::new("user", "hi"),
        ];
        let body = build_body("gemma", &msgs, None);
        assert_eq!(
            body,
            r#"{"model":"gemma","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"hi"}],"stream":true}"#
        );
    }

    #[test]
    fn body_escapes_content() {
        let msgs = [Message::new("user", "say \"hæ\"")];
        let body = build_body("m", &msgs, None);
        assert!(body.contains(r#""content":"say \"hæ\"""#));
    }

    #[test]
    fn body_includes_tools_when_present() {
        let tool = akurai_json::parse(r#"{"type":"function","function":{"name":"now"}}"#).unwrap();
        let body = build_body("m", &[Message::new("user", "x")], Some(&[tool]));
        assert!(body
            .contains(r#""stream":true,"tools":[{"type":"function","function":{"name":"now"}}]"#));
    }

    #[test]
    fn body_omits_empty_tools() {
        let body = build_body("m", &[Message::new("user", "x")], Some(&[]));
        assert!(!body.contains("tools"));
    }

    #[test]
    fn request_is_well_formed() {
        let msgs = [Message::new("user", "hi")];
        let req = build_request("titan", 8080, "gemma", &msgs, None);
        let body = build_body("gemma", &msgs, None);
        assert!(req.starts_with("POST /v1/chat/completions HTTP/1.1\r\n"));
        assert!(req.contains("Host: titan:8080\r\n"));
        assert!(req.contains("Content-Type: application/json\r\n"));
        assert!(req.contains("Accept: text/event-stream\r\n"));
        assert!(req.contains("Connection: close\r\n"));
        assert!(req.contains(&format!("Content-Length: {}\r\n", body.len())));
        assert!(req.ends_with(&format!("\r\n\r\n{body}")));
    }

    #[test]
    fn request_host_header_omits_port_80() {
        let req = build_request("localhost", 80, "m", &[Message::new("user", "x")], None);
        assert!(req.contains("Host: localhost\r\n"));
    }
}