feat: add role-specific config (#42)

Now we can set role temperature. For example:

```
- name: shell
  prompt: >
    I want you to act as a linux shell expert.
    I want you to answer only with bash code.
    Do not write explanations.
  temperature: 0.4
```
This commit is contained in:
sigoden 2023-03-08 19:43:11 +08:00 committed by GitHub
parent 2539e24fe9
commit 56483e04f0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 32 additions and 5 deletions

View File

@ -148,7 +148,7 @@ impl ChatGptClient {
"messages": messages,
});
if let Some(v) = self.config.borrow().temperature {
if let Some(v) = self.config.borrow().get_temperature() {
body.as_object_mut()
.and_then(|m| m.insert("temperature".into(), json!(v)));
}

View File

@ -10,9 +10,9 @@ use std::{
use anyhow::{anyhow, Context, Result};
use inquire::{Confirm, Text};
use serde::Deserialize;
use serde::{Deserialize, Serialize};
use crate::utils::now;
use crate::utils::{emphasis, now};
const CONFIG_FILE_NAME: &str = "config.yaml";
const ROLES_FILE_NAME: &str = "roles.yaml";
@ -153,7 +153,19 @@ impl Config {
pub fn change_role(&mut self, name: &str) -> String {
match self.find_role(name) {
Some(role) => {
let output = format!("{}>> {}", role.name, role.prompt.trim());
let temperature = match role.temperature {
Some(v) => format!("{v}"),
None => "null".into(),
};
let output = format!(
"{}: {}\n{}: {}\n{}: {}",
emphasis("name"),
role.name,
emphasis("prompt"),
role.prompt.trim(),
emphasis("temperature"),
temperature
);
self.role = Some(role);
output
}
@ -165,6 +177,7 @@ impl Config {
self.role = Some(Role {
name: TEMP_ROLE_NAME.into(),
prompt: prompt.into(),
temperature: self.temperature,
});
}
@ -178,6 +191,13 @@ impl Config {
})
}
pub fn get_temperature(&self) -> Option<f64> {
self.role
.as_ref()
.and_then(|v| v.temperature)
.or(self.temperature)
}
pub fn merge_prompt(&self, content: &str) -> String {
match self.get_prompt() {
Some(prompt) => format!("{}\n{content}", prompt.trim()),
@ -305,12 +325,14 @@ impl Config {
}
}
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Role {
/// Role name
pub name: String,
/// Prompt text send to ai for setting up a role
pub prompt: String,
/// What sampling temperature to use, between 0 and 2
pub temperature: Option<f64>,
}
fn create_config_file(config_path: &Path) -> Result<()> {

View File

@ -1,4 +1,5 @@
use chrono::prelude::*;
use crossterm::style::{Color, Stylize};
use std::io::{stdout, Write};
#[macro_export]
@ -17,3 +18,7 @@ pub fn now() -> String {
let now = Local::now();
now.to_rfc3339_opts(SecondsFormat::Secs, false)
}
pub fn emphasis(text: &str) -> String {
text.stylize().with(Color::White).to_string()
}