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
|
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 fn for_execute() -> Self {
let os = detect_os();
let shell = detect_shell();
let shell = match shell.rsplit_once('/') {
Some((_, v)) => v,
None => &shell,
};
Self {
name: "__for_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 &&.
Provide only plain text without Markdown formatting.
Do not provide markdown formatting such as ```"#
),
temperature: None,
}
}
pub fn for_describe() -> Self {
Self {
name: "__for_describe__".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 info(&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"
);
}
}
|
