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
147
148
149
150
151
152
153
154
155
156
|
use crate::{
client::{Message, MessageContent, MessageRole},
utils::{detect_os, detect_shell},
};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use super::Input;
const INPUT_PLACEHOLDER: &str = "__INPUT__";
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Role {
/// Role name
pub name: String,
/// Prompt text
pub prompt: String,
/// What sampling temperature to use, between 0 and 2
pub temperature: Option<f64>,
}
impl Role {
pub const EXECUTE: &'static str = "__execute__";
pub const DESCRIBE_COMMAND: &'static str = "__describe_command__";
pub const CODE: &'static str = "__code__";
pub fn for_execute() -> Self {
let os = detect_os();
let (shell, _, _) = detect_shell();
let combine = match shell.as_str() {
"nushell" | "powershell" => ";",
_ => "&&",
};
Self {
name: Self::EXECUTE.into(),
prompt: format!(
r#"Provide only {shell} commands for {os} without any description.
If there is a lack of details, provide most logical solution.
Ensure the output is a valid {shell} command.
If multiple steps required try to combine them together using {combine}.
Provide only plain text without Markdown formatting.
Do not provide markdown formatting such as ```"#
),
temperature: None,
}
}
pub fn for_describe_command() -> Self {
Self {
name: Self::DESCRIBE_COMMAND.into(),
prompt: r#"Provide a terse, single sentence description of the given shell command.
Describe each argument and option of the command.
Provide short responses in about 80 words.
APPLY MARKDOWN formatting when possible."#
.into(),
temperature: None,
}
}
pub fn for_code() -> Self {
Self {
name: Self::CODE.into(),
prompt: r#"Provide only code as output without any description.
Provide only code in plain text format without Markdown formatting.
Do not include symbols such as ``` or ```python.
If there is a lack of details, provide most logical solution.
You are not allowed to ask for more details.
For example if the prompt is "Hello world Python", you should return "print('Hello world')"."#
.into(),
temperature: None,
}
}
pub fn export(&self) -> Result<String> {
let output = serde_yaml::to_string(&self)
.with_context(|| format!("Unable to show info about role {}", &self.name))?;
Ok(output.trim_end().to_string())
}
pub fn embedded(&self) -> bool {
self.prompt.contains(INPUT_PLACEHOLDER)
}
pub fn complete_prompt_args(&mut self, name: &str) {
self.name = name.to_string();
self.prompt = complete_prompt_args(&self.prompt, &self.name);
}
pub fn match_name(&self, name: &str) -> bool {
if self.name.contains(':') {
let role_name_parts: Vec<&str> = self.name.split(':').collect();
let name_parts: Vec<&str> = name.split(':').collect();
role_name_parts[0] == name_parts[0] && role_name_parts.len() == name_parts.len()
} else {
self.name == name
}
}
pub fn echo_messages(&self, input: &Input) -> String {
let input_markdown = input.render();
if self.embedded() {
self.prompt.replace(INPUT_PLACEHOLDER, &input_markdown)
} else {
format!("{}\n\n{}", self.prompt, input.render())
}
}
pub fn build_messages(&self, input: &Input) -> Vec<Message> {
let mut content = input.to_message_content();
if self.embedded() {
content.merge_prompt(|v: &str| self.prompt.replace(INPUT_PLACEHOLDER, v));
vec![Message {
role: MessageRole::User,
content,
}]
} else {
vec![
Message {
role: MessageRole::System,
content: MessageContent::Text(self.prompt.clone()),
},
Message {
role: MessageRole::User,
content,
},
]
}
}
}
fn complete_prompt_args(prompt: &str, name: &str) -> String {
let mut prompt = prompt.trim().to_string();
for (i, arg) in name.split(':').skip(1).enumerate() {
prompt = prompt.replace(&format!("__ARG{}__", i + 1), arg);
}
prompt
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_merge_prompt_name() {
assert_eq!(
complete_prompt_args("convert __ARG1__", "convert:foo"),
"convert foo"
);
assert_eq!(
complete_prompt_args("convert __ARG1__ to __ARG2__", "convert:foo:bar"),
"convert foo to bar"
);
}
}
|
