aether-llm 0.9.17

Multi-provider LLM abstraction layer for the Aether AI agent framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use tracing::debug;

use super::auth::MantleAuth;
use crate::catalog::transport::ModelTransport;
use crate::providers::http::{SseData, http_client, post_sse, responses_code};
use crate::providers::openai::responses::mappers::{ResponsesRequestPolicy, build_request};
use crate::providers::response_stream::OpenedStream;
use crate::{Context, LlmError, LlmModel, Result};
use utils::variables::Vars;

/// Transport for Bedrock models served through the `OpenAI` Responses API.
#[derive(Clone)]
pub(crate) struct MantleClient {
    http: reqwest::Client,
    template_vars: Vars,
    auth: MantleAuth,
    base_url_override: Option<String>,
}

impl MantleClient {
    pub(crate) fn new(region: String, auth: MantleAuth, base_url_override: Option<String>) -> Self {
        Self {
            http: http_client(),
            template_vars: Vars::new().with("AWS_REGION", region.clone()).with("AWS_DEFAULT_REGION", region),
            auth,
            base_url_override: base_url_override.map(|url| url.trim_end_matches('/').to_string()),
        }
    }

    pub(crate) fn with_auth(mut self, auth: MantleAuth) -> Self {
        self.auth = auth;
        self
    }

    /// Endpoint this client will POST to for `transport`.
    ///
    /// An explicit base URL — a test server or a private endpoint — wins over the
    /// catalog template, which is why the override is checked first.
    pub(crate) fn endpoint(&self, transport: &ModelTransport) -> Result<String> {
        let ModelTransport::OpenAiResponses { base_url_template } = transport;
        let base = match &self.base_url_override {
            Some(base) => base.clone(),
            None => {
                self.template_vars.expand(base_url_template).map_err(|e| LlmError::InvalidArgument(e.to_string()))?
            }
        };

        Ok(format!("{}/responses", base.trim_end_matches('/')))
    }

    pub(crate) async fn stream(
        &self,
        identity: Option<&LlmModel>,
        model: &str,
        transport: &ModelTransport,
        context: &Context,
    ) -> Result<OpenedStream<SseData>> {
        let url = self.endpoint(transport)?;
        let body = build_request(identity, model, context, &ResponsesRequestPolicy::MANTLE)?;

        debug!(model, url, auth = %self.auth, "Sending Bedrock responses request");

        let encoded = serde_json::to_vec(&body)?;
        let headers = self.auth.headers("POST", &url, &encoded).await?;
        post_sse(&self.http, &url, headers, &body, responses_code).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    fn transport(api: &'static str) -> ModelTransport {
        ModelTransport::OpenAiResponses { base_url_template: api }
    }

    fn client(base_url_override: Option<String>) -> MantleClient {
        MantleClient::new("us-west-2".to_string(), MantleAuth::None, base_url_override)
    }

    #[test]
    fn endpoint_expands_the_region_from_the_model_template() {
        let endpoint =
            client(None).endpoint(&transport("https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1")).unwrap();

        assert_eq!(endpoint, "https://bedrock-mantle.us-west-2.api.aws/openai/v1/responses");
    }

    #[test]
    fn endpoint_preserves_a_models_own_api_path() {
        let endpoint = client(None).endpoint(&transport("https://bedrock-mantle.${AWS_REGION}.api.aws/v1")).unwrap();

        assert_eq!(endpoint, "https://bedrock-mantle.us-west-2.api.aws/v1/responses");
    }

    #[test]
    fn base_url_override_replaces_the_catalog_template() {
        let endpoint = client(Some("http://127.0.0.1:9999/".to_string()))
            .endpoint(&transport("https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"))
            .unwrap();

        assert_eq!(endpoint, "http://127.0.0.1:9999/responses");
    }
}