summaryrefslogtreecommitdiffstats
path: root/src/config/conversation.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2023-03-09 19:10:02 +0800
committerGitHub <noreply@github.com>2023-03-09 19:10:02 +0800
commit9767c07eeeb20cb2f8386faf66cb88452987988b (patch)
treefd6eec9a9ccad49798ab9de7329db507987d8a4d /src/config/conversation.rs
parentc45d71cdeabbced7a55381e3b431c3eae3799ac0 (diff)
downloadaichat-9767c07eeeb20cb2f8386faf66cb88452987988b.tar.gz
feat: support two types of role prompts (#52)
1. embeded prompt use __INPUT__ placeholder will generate one user message when send to gpt ``` - name: shell prompt: > I want you to act as a linux shell expert. Q: How to unzip a file A: unzip file.zip Q: __INPUT__ A: ``` 2. system prompt no __INPUT__ placeholder will generate on system message and one user message when send to gpt ``` - 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. ```
Diffstat (limited to 'src/config/conversation.rs')
-rw-r--r--src/config/conversation.rs108
1 files changed, 48 insertions, 60 deletions
diff --git a/src/config/conversation.rs b/src/config/conversation.rs
index 23eacb5..c3e9bed 100644
--- a/src/config/conversation.rs
+++ b/src/config/conversation.rs
@@ -1,93 +1,81 @@
-use anyhow::Result;
-use serde::{Deserialize, Serialize};
-use serde_json::{json, Value};
+use super::message::{Message, MessageRole, MESSAGE_EXTRA_TOKENS};
+use super::role::Role;
+use super::MAX_TOKENS;
use crate::utils::count_tokens;
-use super::{MAX_TOKENS, MESSAGE_EXTRA_TOKENS};
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Conversation {
pub tokens: usize,
+ pub role: Option<Role>,
pub messages: Vec<Message>,
}
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct Message {
- pub role: MessageRole,
- pub content: String,
-}
-
impl Conversation {
- pub fn new() -> Self {
+ pub fn new(role: Option<Role>) -> Self {
+ let tokens = if let Some(role) = role.as_ref() {
+ role.consume_tokens()
+ } else {
+ 0
+ };
Self {
- tokens: 0,
+ tokens,
+ role,
messages: vec![],
}
}
- pub fn add_chat(&mut self, input: &str, output: &str) -> Result<()> {
- self.messages.push(Message {
- role: MessageRole::User,
- content: input.to_string(),
- });
+ pub fn add_message(&mut self, input: &str, output: &str) -> Result<()> {
+ let mut need_add_msg = true;
+ let mut input_tokens = count_tokens(input);
+ if self.messages.is_empty() {
+ if let Some(role) = self.role.as_ref() {
+ self.messages.extend(role.build_emssages(input));
+ need_add_msg = false;
+ }
+ }
+ if need_add_msg {
+ self.messages.push(Message {
+ role: MessageRole::User,
+ content: input.to_string(),
+ });
+ input_tokens += MESSAGE_EXTRA_TOKENS;
+ }
self.messages.push(Message {
role: MessageRole::Assistant,
content: output.to_string(),
});
- self.tokens += count_tokens(input) + count_tokens(output) + 2 * MESSAGE_EXTRA_TOKENS;
+ self.tokens += input_tokens + count_tokens(output) + MESSAGE_EXTRA_TOKENS;
Ok(())
}
- /// Readline prompt
- pub fn add_prompt(&mut self, prompt: &str) {
- self.messages.push(Message {
- role: MessageRole::System,
- content: prompt.into(),
- });
- self.tokens += count_tokens(prompt) + MESSAGE_EXTRA_TOKENS;
- }
-
pub fn echo_messages(&self, content: &str) -> String {
- let mut messages = self.messages.to_vec();
- messages.push(Message {
- role: MessageRole::User,
- content: content.into(),
- });
+ let messages = self.build_emssages(content);
serde_yaml::to_string(&messages).unwrap_or("Unable to echo message".into())
}
- pub fn build_emssages(&self, content: &str) -> Value {
- let mut messages: Vec<Value> = self.messages.iter().map(msg_to_value).collect();
- messages.push(msg_to_value(&Message {
- role: MessageRole::User,
- content: content.into(),
- }));
- json!(messages)
+ pub fn build_emssages(&self, content: &str) -> Vec<Message> {
+ let mut messages = self.messages.to_vec();
+ let mut need_add_msg = true;
+ if messages.is_empty() {
+ if let Some(role) = self.role.as_ref() {
+ messages = role.build_emssages(content);
+ need_add_msg = false;
+ }
+ };
+ if need_add_msg {
+ messages.push(Message {
+ role: MessageRole::User,
+ content: content.into(),
+ });
+ }
+ messages
}
pub fn reamind_tokens(&self) -> usize {
MAX_TOKENS.saturating_sub(self.tokens)
}
}
-
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub enum MessageRole {
- System,
- Assistant,
- User,
-}
-
-impl MessageRole {
- pub fn name(&self) -> &'static str {
- match self {
- MessageRole::System => "system",
- MessageRole::Assistant => "assistant",
- MessageRole::User => "user",
- }
- }
-}
-
-fn msg_to_value(msg: &Message) -> Value {
- json!({ "role": msg.role.name(), "content": msg.content })
-}