feat: Support Harmony format responses

Cargo.toml

@@ -6,7 +6,7 @@ edition = "2024"
 [dependencies]
 xdg = "2.5"
 serde = { version = "1.0", features = ["derive"] }
-serde_json = "1.0"
+serde_json = { version = "1.0", features = ["raw_value"] }
 json5 = "0.4"
 reqwest = { version = "0.12", features = ["json"] }
 tokio = { version = "1.0", features = ["full"] }
@@ -15,3 +15,4 @@ anyhow = "1.0"
 shellexpand = "3.0"
 colored = "3.0.0"
 termios = "0.3.3"
+regex = "1.12.2"

README.md

@@ -62,12 +62,20 @@ place a [json5](https://json5.org/) file in `$XDG_CONFIG_HOME/please/config.json
     // Endpoint URL
     "endpoint": "https://openrouter.ai/api/v1",
 
+    // Response format of the model
+    // accepted values are "harmony" | "json_schema"
+    // if not specified, defaults to "json_schema"
+    // "response-format": "json_schema",
+
     "prompts": {
         // Prompt template for generating shell commands
         "command": "You are an expert in the Linux shell. The user would like to perform a task in the shell. \
  Please return ONLY a single shell command compatible with the user's shell (it will be ran with `$SHELL`). \
  Prefer single-line solutions. Do not include any markdown formatting, explanations, or multiple options. \
- Just return the raw command that can be executed directly.",
+ Your answer should just be the raw command that can be executed directly. \
+ Do not include $SHELL at the start of the command the user will take care of inserting that. \
+ The command should be broken into segments (e.g `echo foo` -> [\"echo\", \"foo\"]). \
+ Respond with a JSON object as follows { \"command\":  [\"YOUR\", \"COMMAND\"] }",
     }
 }
 ```

src/api.rs

@@ -1,6 +1,9 @@
 use anyhow::{Context, Result};
+use regex::Regex;
 use serde::{Deserialize, Serialize};
 
+use crate::config::{Config, ResponseFormat};
+
 #[derive(Debug, Serialize)]
 struct ChatRequest {
     model: String,
@@ -23,6 +26,11 @@ struct Choice {
     message: Message,
 }
 
+#[derive(Debug, Deserialize)]
+struct CommandResponse {
+    command: Vec<String>,
+}
+
 pub struct ApiClient {
     client: reqwest::Client,
     endpoint: String,
@@ -45,6 +53,7 @@ impl ApiClient {
         &self,
         system_prompt: &str,
         user_message: &str,
+        config: &Config,
     ) -> Result<String> {
         let url = format!("{}/chat/completions", self.endpoint);
 
@@ -109,9 +118,11 @@ impl ApiClient {
             .context("No response from API")?;
 
         // Clean up the response (remove markdown code blocks)
-        let command = clean_command_response(&raw_response);
-
-        Ok(command)
+        let command = match config.response_format {
+            ResponseFormat::Harmony => cleave_start_parse_json(&raw_response),
+            _ => Ok(clean_command_response(&raw_response)),
+        };
+        command
     }
 
     /// Continue a conversation with feedback from the user
@@ -121,6 +132,7 @@ impl ApiClient {
         original_request: &str,
         previous_command: &str,
         feedback: &str,
+        config: &Config,
     ) -> Result<String> {
         let url = format!("{}/chat/completions", self.endpoint);
 
@@ -178,12 +190,30 @@ impl ApiClient {
             .context("No response from API")?;
 
         // Clean up the response (remove markdown code blocks)
-        let command = clean_command_response(&raw_response);
-
-        Ok(command)
+        let command = match config.response_format {
+            ResponseFormat::Harmony => cleave_start_parse_json(&raw_response),
+            _ => Ok(clean_command_response(&raw_response)),
+        };
+        command
     }
 }
 
+fn cleave_start_parse_json(response: &str) -> Result<String> {
+    let regex = Regex::new(r"(?m)<\|end\|>(\{.*\}$)").unwrap();
+    let captures = regex
+        .captures(response)
+        .context("Harmony parse - Failed to match regex")?;
+
+    let json_str = captures
+        .get(1)
+        .context("Harmony parse - Empty json section")?;
+
+    let json_str = json_str.as_str();
+    let command_response: CommandResponse = serde_json::from_str(json_str)
+        .expect(format!("Failed to parse JSON: {}", json_str).as_str());
+    return Ok(command_response.command.join(" "));
+}
+
 /// Clean up the command response by removing markdown code blocks and extra text
 fn clean_command_response(response: &str) -> String {
     let response = response.trim();

src/config.rs

@@ -21,22 +21,53 @@ pub const DEFAULT_CONFIG_FILE: &str = r#"// please cli configuration
     // Endpoint URL
     "endpoint": "https://openrouter.ai/api/v1",
 
+    // Response format of the model
+    // accepted values are "harmony" | "json_schema"
+    // if not specified, defaults to "json_schema"
+    // "response-format": "json_schema"
+
     "prompts": {
         // Prompt template for generating shell commands
         "command": "You are an expert in the Linux shell. The user would like to perform a task in the shell. \
  Please return ONLY a single shell command compatible with the user's shell (it will be ran with `$SHELL`). \
  Prefer single-line solutions. Do not include any markdown formatting, explanations, or multiple options. \
- Just return the raw command that can be executed directly.",
+ Your answer should just be the raw command that can be executed directly. \
+ Do not include $SHELL at the start of the command the user will take care of inserting that. \
+ The command should be broken into segments (e.g `echo foo` -> [\"echo\", \"foo\"]). \
+ Respond with a JSON object as follows { \"command\":  [\"YOUR\", \"COMMAND\"] }",
     }
 }
 
 /* -*- mode: json5 -*- */
 /* vim: set ft=json5: */
 "#;
 
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ResponseFormat {
+    Harmony,
+    JsonSchema
+}
+impl Default for ResponseFormat {
+    fn default() -> Self {
+        ResponseFormat::JsonSchema
+    }
+} 
+
+impl TryFrom::<&String> for ResponseFormat {
+    type Error = String;
+    fn try_from(value: &String) -> anyhow::Result<Self, Self::Error> {
+        match value.as_str() {
+            "json_schema" => Ok(Self::JsonSchema),
+            "harmony" => Ok(ResponseFormat::Harmony),
+            _ => Err(format!("Unrecognized response format specified {}", value))
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Config {
-    #[serde(rename = "api-key")]
+    #[serde(default, rename = "api-key")]
     pub api_key: String,
 
     #[serde(default = "default_model")]
@@ -51,6 +82,9 @@ pub struct Config {
     #[serde(default = "default_endpoint")]
     pub endpoint: String,
 
+    #[serde(default, rename = "response-format")]
+    pub response_format: ResponseFormat,
+
     #[serde(default)]
     pub prompts: Prompts,
 }
@@ -81,11 +115,18 @@ fn default_endpoint() -> String {
     "https://openrouter.ai/api/v1".to_string()
 }
 
+fn default_response_format() -> ResponseFormat {
+    ResponseFormat::JsonSchema
+}
+
 fn default_command_prompt() -> String {
     r#"You are an expert in the Linux shell. The user would like to perform a task in the shell.
 Please return ONLY a single shell command compatible with the user's shell (it will be ran with `$SHELL`).
 Prefer single-line solutions. Do not include any markdown formatting, explanations, or multiple options.
-Just return the raw command that can be executed directly."#.to_string()
+Your answer should just be the raw command that can be executed directly.
+Do not include $SHELL at the start of the command the user will take care of inserting that.
+The command should be broken into segments (e.g `echo foo` -> ["echo", "foo"]).
+Respond with a JSON object as follows { "command":  ["YOUR", "COMMAND"] }"#.to_string()
 }
 
 impl Config {
@@ -113,6 +154,7 @@ impl Config {
                 quiet: false,
                 shell: default_shell(),
                 endpoint: default_endpoint(),
+                response_format: default_response_format(),
                 prompts: Prompts::default(),
             }
         };
@@ -134,14 +176,22 @@ impl Config {
             config.endpoint = endpoint;
         }
 
+        if let Ok(response_format) = env::var("PLEASE_RESPONSE_FORMAT") {
+            if let Ok(rf) = ResponseFormat::try_from(&response_format) {
+                config.response_format = rf;
+            } else {
+                let msg = format!("Unrecognized response format specified in PLEASE_RESPONSE_FORMAT ({})", response_format);
+                anyhow::bail!(msg);
+            }
+        }
+
         if let Ok(quiet) = env::var("PLEASE_QUIET") {
             config.quiet = quiet == "1" || quiet.to_lowercase() == "true";
         }
 
         if let Ok(command_prompt) = env::var("PLEASE_PROMPTS_COMMAND") {
             config.prompts.command = command_prompt;
         }
-
         Ok(config)
     }
 

src/main.rs

@@ -79,7 +79,7 @@ async fn main() -> Result<()> {
         ui.show_message("Thinking...");
     }
     let mut current_command = match api_client
-        .request_command(&system_prompt, &user_request)
+        .request_command(&system_prompt, &user_request, &config)
         .await
     {
         Ok(cmd) => cmd,
@@ -111,6 +111,7 @@ async fn main() -> Result<()> {
                         &user_request,
                         &current_command,
                         &feedback,
+                        &config,
                     )
                     .await
                 {