smartcat 0.5.1

Putting a brain behind `cat`. CLI interface to bring language models in the Unix ecosystem 🐈‍⬛ 
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use log::debug;
use std::io::{self, Read, Result, Write};

use crate::config::{get_api_config, Message, Prompt, PLACEHOLDER_TOKEN};
use crate::request::{make_authenticated_request, OpenAiResponse};

// [tmp] mostly template to write tests
pub fn chunk_process_input<R: Read, W: Write>(
    input: &mut R,
    output: &mut W,
    prefix: &str,
    suffix: &str,
) -> Result<()> {
    let mut first_chunk = true;
    let mut buffer = [0; 1024];
    loop {
        match input.read(&mut buffer) {
            Ok(0) => break, // end of input
            Ok(n) => {
                if first_chunk {
                    output.write_all(prefix.as_bytes())?;
                    first_chunk = false;
                }
                output.write_all(&buffer[..n])?;
            }
            Err(e) => return Err(e),
        }
    }

    if !first_chunk {
        // we actually got some input
        output.write_all(suffix.as_bytes())?;
    }

    Ok(())
}

pub fn process_input_with_request<R: Read, W: Write>(
    mut prompt: Prompt,
    input: &mut R,
    input_string: Option<String>,
    output: &mut W,
    repeat_input: bool,
) -> Result<Prompt> {
    let mut input = match input_string {
        Some(input_string) => input_string,
        None => {
            let mut buffer = Vec::new();
            input.read_to_end(&mut buffer)?;

            String::from_utf8(buffer).unwrap()
        }
    };

    // nothing to do if no input
    if input.is_empty() {
        return Ok(prompt);
    }

    // insert the input in the messages with placeholders
    for message in prompt.messages.iter_mut() {
        message.content = message.content.replace(PLACEHOLDER_TOKEN, &input)
    }
    // fetch the api config tied to the prompt
    let api_config = get_api_config(&prompt.api.to_string());

    // make the request
    let response: OpenAiResponse = make_authenticated_request(api_config, &prompt)
        .map_err(|e| match e {
            ureq::Error::Status(status, response) => {
                let body = match response.into_string() {
                    Ok(body) => body,
                    Err(_) => "(non-UTF-8 response)".to_owned(),
                };
                io::Error::other(format!(
                    "API call failed with status code {status} and body: {body}"
                ))
            }
            ureq::Error::Transport(transport) => io::Error::other(transport),
        })?
        .into_json()?;

    let response_text = response.choices.first().unwrap().message.content.as_str();
    debug!("{}", &response_text);

    prompt.messages.push(Message::assistant(response_text));

    if repeat_input {
        input.push('\n');
        output.write_all(input.as_bytes())?;
    }

    output.write_all(response_text.as_bytes())?;

    Ok(prompt)
}

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

    macro_rules! test_process_input {
        ($test_name:ident, $prefix:expr, $suffix:expr, $input:expr) => {
            #[test]
            fn $test_name() {
                let input = $input.as_bytes();
                let mut output = std::io::Cursor::new(Vec::new());

                let result =
                    chunk_process_input(&mut Cursor::new(input), &mut output, $prefix, $suffix);
                assert!(result.is_ok());

                let expected_output = if !input.is_empty() {
                    format!("{}{}{}", $prefix, $input, $suffix)
                } else {
                    "".into()
                };

                let expected_output_as_bytes = expected_output.as_bytes();
                let output_data: Vec<u8> = output.into_inner();
                assert_eq!(
                    expected_output_as_bytes,
                    output_data,
                    "\nexpected: {}\nGot: {}",
                    String::from_utf8_lossy(expected_output_as_bytes),
                    &expected_output
                );
            }
        };
    }

    test_process_input!(
        test_with_prefix_and_suffix,
        "Prefix: ",
        " Suffix",
        "Input data"
    );
    test_process_input!(
        test_with_custom_prefix_suffix,
        "Start: ",
        " End",
        "Custom input"
    );
    test_process_input!(test_empty_input, "Pre: ", " Post", "");
}