diff options
| author | sigoden <sigoden@gmail.com> | 2023-03-09 19:10:02 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-03-09 19:10:02 +0800 |
| commit | 9767c07eeeb20cb2f8386faf66cb88452987988b (patch) | |
| tree | fd6eec9a9ccad49798ab9de7329db507987d8a4d /src/config/conversation.rs | |
| parent | c45d71cdeabbced7a55381e3b431c3eae3799ac0 (diff) | |
| download | aichat-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.rs | 108 |
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 }) -} |
