Skip to main content

hyperinfer_providers/
lib.rs

1pub mod provider_trait;
2pub mod registry;
3
4#[cfg(feature = "anthropic")]
5pub mod anthropic;
6#[cfg(feature = "openai")]
7pub mod openai;
8
9pub use provider_trait::LlmProvider;
10pub use registry::ProviderRegistry;
11
12pub fn drain_lines(raw_buf: &mut Vec<u8>, lines: &mut Vec<String>) {
13    if raw_buf.is_empty() {
14        return;
15    }
16
17    match std::str::from_utf8(raw_buf) {
18        Ok(s) => {
19            let mut rest = s;
20            while let Some(pos) = rest.find('\n') {
21                let (line, next_rest) = rest.split_at(pos);
22                let line = line.strip_suffix('\r').unwrap_or(line);
23                lines.push(line.to_string());
24                rest = &next_rest[1..];
25            }
26            if rest.is_empty() {
27                raw_buf.clear();
28            } else {
29                *raw_buf = rest.as_bytes().to_vec();
30            }
31        }
32        Err(e) => {
33            let valid_up_to = e.valid_up_to();
34            let valid_prefix = &raw_buf[..valid_up_to];
35
36            if let Some(pos) = valid_prefix.iter().rposition(|&b| b == b'\n') {
37                let line_bytes = &valid_prefix[..pos];
38                let line_bytes = line_bytes.strip_suffix(b"\r").unwrap_or(line_bytes);
39                lines.push(String::from_utf8_lossy(line_bytes).into_owned());
40                raw_buf.drain(..=pos + 1);
41            }
42        }
43    }
44}
45
46pub fn init_default_registry(registry: &ProviderRegistry) {
47    #[cfg(feature = "openai")]
48    {
49        if let Ok(provider) = openai::OpenAiProvider::new() {
50            if !registry.contains(provider.name()) {
51                registry.register(provider);
52            }
53        }
54    }
55
56    #[cfg(feature = "anthropic")]
57    {
58        if let Ok(provider) = anthropic::AnthropicProvider::new() {
59            if !registry.contains(provider.name()) {
60                registry.register(provider);
61            }
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    fn feed_chunks(chunks: &[&[u8]]) -> (Vec<String>, Vec<u8>) {
71        let mut raw_buf: Vec<u8> = Vec::new();
72        let mut all_lines: Vec<String> = Vec::new();
73        for chunk in chunks {
74            raw_buf.extend_from_slice(chunk);
75            drain_lines(&mut raw_buf, &mut all_lines);
76        }
77        (all_lines, raw_buf)
78    }
79
80    #[test]
81    fn test_drain_lines_single_chunk() {
82        let (lines, remainder) = feed_chunks(&[b"data: hello\ndata: world\n"]);
83        assert_eq!(lines, vec!["data: hello", "data: world"]);
84        assert!(remainder.is_empty());
85    }
86
87    #[test]
88    fn test_drain_lines_crlf_endings() {
89        let (lines, remainder) = feed_chunks(&[b"data: hello\r\ndata: world\r\n"]);
90        assert_eq!(lines, vec!["data: hello", "data: world"]);
91        assert!(remainder.is_empty());
92    }
93
94    #[test]
95    fn test_drain_lines_incomplete_line_buffered() {
96        let (lines, remainder) = feed_chunks(&[b"data: hello\n", b"data: partial"]);
97        assert_eq!(lines, vec!["data: hello"]);
98        assert_eq!(remainder, b"data: partial");
99    }
100
101    #[test]
102    fn test_drain_lines_multibyte_split_across_chunks() {
103        let chunk1: &[u8] = b"data: caf\xc3";
104        let chunk2: &[u8] = b"\xa9\ndata: done\n";
105        let (lines, remainder) = feed_chunks(&[chunk1, chunk2]);
106        assert_eq!(lines[0], "data: café");
107        assert_eq!(lines[1], "data: done");
108        assert!(remainder.is_empty());
109    }
110
111    #[test]
112    fn test_drain_lines_three_byte_split_across_three_chunks() {
113        let chunk1: &[u8] = b"data: \xe4";
114        let chunk2: &[u8] = b"\xb8";
115        let chunk3: &[u8] = b"\xad\n";
116        let (lines, remainder) = feed_chunks(&[chunk1, chunk2, chunk3]);
117        assert_eq!(lines, vec!["data: 中"]);
118        assert!(remainder.is_empty());
119    }
120
121    #[test]
122    fn test_drain_lines_empty_lines_preserved() {
123        let (lines, _) = feed_chunks(&[b"data: hello\n\ndata: world\n"]);
124        assert_eq!(lines, vec!["data: hello", "", "data: world"]);
125    }
126
127    #[test]
128    fn test_drain_lines_no_newline_nothing_emitted() {
129        let (lines, remainder) = feed_chunks(&[b"data: no newline yet"]);
130        assert!(lines.is_empty());
131        assert_eq!(remainder, b"data: no newline yet");
132    }
133
134    #[test]
135    fn test_drain_lines_utf8_invalid_bytes_preserved() {
136        let chunk1: &[u8] = b"data: \xc3";
137        let chunk2: &[u8] = b"\xa9\n";
138        let (lines, remainder) = feed_chunks(&[chunk1, chunk2]);
139        assert_eq!(lines, vec!["data: é"]);
140        assert!(remainder.is_empty());
141    }
142
143    #[test]
144    fn test_drain_lines_multiple_incomplete_chunks() {
145        let chunk1: &[u8] = b"data: \xe4\xb8";
146        let chunk2: &[u8] = b"\xad";
147        let chunk3: &[u8] = b"\ndata: done\n";
148        let (lines, remainder) = feed_chunks(&[chunk1, chunk2, chunk3]);
149        assert_eq!(lines, vec!["data: 中", "data: done"]);
150        assert!(remainder.is_empty());
151    }
152
153    #[test]
154    fn test_drain_lines_mixed_crlf_and_lf() {
155        let (lines, _) = feed_chunks(&[b"line1\r\nline2\nline3\r\n"]);
156        assert_eq!(lines, vec!["line1", "line2", "line3"]);
157    }
158}