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
|
mod cli;
mod client;
mod config;
mod render;
mod repl;
mod term;
#[macro_use]
mod utils;
use crate::cli::Cli;
use crate::client::ChatGptClient;
use crate::config::{Config, SharedConfig};
use anyhow::{anyhow, Result};
use clap::Parser;
use crossbeam::sync::WaitGroup;
use is_terminal::IsTerminal;
use parking_lot::RwLock;
use render::{render_stream, MarkdownRender};
use repl::{AbortSignal, Repl};
use std::io::{stdin, Read};
use std::sync::Arc;
use std::{io::stdout, process::exit};
use utils::cl100k_base_singleton;
fn main() -> Result<()> {
let cli = Cli::parse();
let text = cli.text();
let config = Arc::new(RwLock::new(Config::init(text.is_none())?));
if cli.list_roles {
config
.read()
.roles
.iter()
.for_each(|v| println!("{}", v.name));
exit(0);
}
if cli.list_models {
for (name, _) in &config::MODELS {
println!("{name}");
}
exit(0);
}
if cli.dry_run {
config.write().dry_run = true;
}
let role = match &cli.role {
Some(name) => Some(
config
.read()
.get_role(name)
.ok_or_else(|| anyhow!("Unknown role '{name}'"))?,
),
None => None,
};
if let Some(model) = &cli.model {
config.write().set_model(model)?;
}
config.write().role = role;
if cli.no_highlight {
config.write().highlight = false;
}
if let Some(prompt) = &cli.prompt {
config.write().add_prompt(prompt)?;
}
if cli.info {
let info = config.read().info()?;
println!("{info}");
exit(0);
}
let no_stream = cli.no_stream;
let client = ChatGptClient::init(config.clone())?;
if atty::isnt(atty::Stream::Stdin) {
let mut input = String::new();
stdin().read_to_string(&mut input)?;
if let Some(text) = text {
input = format!("{text}\n{input}");
}
start_directive(&client, &config, &input, no_stream)
} else {
match text {
Some(text) => start_directive(&client, &config, &text, no_stream),
None => start_interactive(client, config),
}
}
}
fn start_directive(
client: &ChatGptClient,
config: &SharedConfig,
input: &str,
no_stream: bool,
) -> Result<()> {
if !stdout().is_terminal() {
config.write().highlight = false;
}
config.read().maybe_print_send_tokens(input);
let output = if no_stream {
let (highlight, light_theme) = config.read().get_render_options();
let output = client.send_message(input)?;
if highlight {
let mut markdown_render = MarkdownRender::new(light_theme);
println!("{}", markdown_render.render(&output).trim());
} else {
println!("{}", output.trim());
}
output
} else {
let wg = WaitGroup::new();
let abort = AbortSignal::new();
let abort_clone = abort.clone();
ctrlc::set_handler(move || {
abort_clone.set_ctrlc();
})
.expect("Error setting Ctrl-C handler");
let output = render_stream(input, client, config, false, abort, wg.clone())?;
wg.wait();
output
};
config.read().save_message(input, &output)
}
fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> {
cl100k_base_singleton();
config.write().on_repl()?;
let mut repl = Repl::init(config.clone())?;
repl.run(client, config)
}
|
