From 9767c07eeeb20cb2f8386faf66cb88452987988b Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 9 Mar 2023 19:10:02 +0800
Subject: 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.
```
---
 src/config/message.rs | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)
 create mode 100644 src/config/message.rs

(limited to 'src/config/message.rs')

diff --git a/src/config/message.rs b/src/config/message.rs
new file mode 100644
index 0000000..0a06a73
--- /dev/null
+++ b/src/config/message.rs
@@ -0,0 +1,34 @@
+use serde::{Deserialize, Serialize};
+
+pub const MESSAGE_EXTRA_TOKENS: usize = 6;
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct Message {
+    pub role: MessageRole,
+    pub content: String,
+}
+
+impl Message {
+    pub fn new(content: &str) -> Self {
+        Self {
+            role: MessageRole::User,
+            content: content.to_string(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum MessageRole {
+    System,
+    Assistant,
+    User,
+}
+
+#[test]
+fn test_serde() {
+    assert_eq!(
+        serde_json::to_string(&Message::new("Hello World")).unwrap(),
+        "{\"role\":\"user\",\"content\":\"Hello World\"}"
+    )
+}
-- 
cgit v1.2.3