rust_ai/utils/
header.rs

1//!
2//! A module for manipulating additional HTTP headers especially when
3//! communicating with proxy servers.
4
5////////////////////////////////////////////////////////////////////////////////
6
7use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
8
9const ADDITIONAL_HEADERS_VAR: &'static str = "RUST_AI_ADDITIONAL_HEADERS";
10
11/// A type alias for (String, String), which denotes HTTP header/value pair.
12type RawHeader<'a> = (&'a str, &'a str);
13type Header = (String, String);
14
15/// This is the only interface to set/get runtime level headers when sending
16/// requests to endpoints such as OpenAI.
17///
18/// # Usage
19/// 1. Call `AdditionalHeaders::default()` to initialize an instance.
20/// 2. Call `set_header()` to set required headers.
21/// 3. Depends on the caller:
22///   - If you are using Rust-AI as an library, call `to_var()` to set related
23///     environment variable.
24///   - If you are calling from Rust-AI itself, you should call `provide()` to
25///     turn current instance into [`HeaderMap`].
26#[derive(serde::Serialize, serde::Deserialize, Debug)]
27pub struct AdditionalHeaders {
28    headers: Vec<Header>,
29}
30
31impl Default for AdditionalHeaders {
32    /// Create a new instance of [`AdditionalHeaders`] that doesn't contain any
33    /// headers.
34    fn default() -> Self {
35        Self { headers: vec![] }
36    }
37}
38
39impl AdditionalHeaders {
40    /// Try to initialize [`AdditionalHeaders`] instance from environment
41    /// variable. If no such environment variables available, then create an
42    /// empty instance.
43    pub fn from_var() -> Self {
44        if let Ok(c) = std::env::var(ADDITIONAL_HEADERS_VAR) {
45            if let Ok(deserialized) = serde_json::from_str::<Self>(&c) {
46                return deserialized;
47            }
48        }
49        Default::default()
50    }
51}
52
53impl AdditionalHeaders {
54    /// Set a new header pair. Header name and values should never include NUL
55    /// characters (`\0`).
56    pub fn set_header<'a>(&mut self, header: RawHeader<'a>) {
57        let header_name = header.0.to_string();
58        let header_value = header.1.to_string();
59        if header_name.contains("\0") || header_value.contains("\0") {
60            panic!("`\0` cannot present in any field of the header");
61        }
62        self.headers.push((header_name, header_value));
63    }
64
65    /// Turn the contained headers into [`HeaderMap`]. Will not consume current
66    /// instance.
67    pub fn provide(&self) -> HeaderMap {
68        let mut hm = HeaderMap::new();
69        self.headers.iter().for_each(|h| {
70            hm.append(
71                HeaderName::from_lowercase(h.0.to_lowercase().as_bytes()).unwrap(),
72                HeaderValue::from_str(&h.1).unwrap(),
73            );
74        });
75
76        hm
77    }
78
79    /// Try to serialize current instance into string form and set to specific
80    /// environment variable.
81    pub fn to_var(&self) {
82        let serialized = serde_json::to_string(self).unwrap();
83        std::env::set_var(ADDITIONAL_HEADERS_VAR, serialized);
84    }
85}
rust_ai/utils/header.rs

rust_ai/utils/
header.rs
