summaryrefslogtreecommitdiffstats
path: root/src/render/mod.rs
blob: efbcdd75d04814ce46f6b3321a87b49cbff04575 (plain) (blame)
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
mod cmd;
mod markdown;
mod repl;

use self::cmd::cmd_render_stream;
pub use self::markdown::MarkdownRender;
use self::repl::repl_render_stream;

use crate::client::ChatGptClient;
use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
use anyhow::Result;
use crossbeam::channel::unbounded;
use crossbeam::sync::WaitGroup;
use std::thread::spawn;

pub fn render_stream(
    input: &str,
    prompt: Option<String>,
    client: &ChatGptClient,
    highlight: bool,
    repl: bool,
    abort: SharedAbortSignal,
    wg: WaitGroup,
) -> Result<String> {
    let mut stream_handler = if highlight {
        let (tx, rx) = unbounded();
        let abort_clone = abort.clone();
        spawn(move || {
            let _ = if repl {
                repl_render_stream(rx, abort)
            } else {
                cmd_render_stream(rx, abort)
            };
            drop(wg);
        });
        ReplyStreamHandler::new(Some(tx), abort_clone)
    } else {
        drop(wg);
        ReplyStreamHandler::new(None, abort)
    };
    client.send_message_streaming(input, prompt, &mut stream_handler)?;
    let buffer = stream_handler.get_buffer();
    Ok(buffer.to_string())
}