feat: add `-c/--code` to generate only code (#327)

2 files changed, 21 insertions, 2 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index f8d97c2..6b99483 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -288,6 +288,11 @@ impl Config {
         self.set_role_obj(role)
     }
 
+    pub fn set_code_role(&mut self) -> Result<()> {
+        let role = Role::for_code();
+        self.set_role_obj(role)
+    }
+
     pub fn set_role_obj(&mut self, role: Role) -> Result<()> {
         if let Some(session) = self.session.as_mut() {
             session.update_role(Some(role.clone()))?;
diff --git a/src/config/role.rs b/src/config/role.rs
index a63877b..39f0f17 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -29,7 +29,7 @@ impl Role {
             _ => "&&",
         };
         Self {
-            name: "__for_execute__".into(),
+            name: "__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.
@@ -44,7 +44,7 @@ Do not provide markdown formatting such as ```"#
 
     pub fn for_describe() -> Self {
         Self {
-            name: "__for_describe__".into(),
+            name: "__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.
@@ -54,6 +54,20 @@ APPLY MARKDOWN formatting when possible."#
         }
     }
 
+    pub fn for_code() -> Self {
+        Self {
+            name: "__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 info(&self) -> Result<String> {
         let output = serde_yaml::to_string(&self)
             .with_context(|| format!("Unable to show info about role {}", &self.name))?;