Fix assistant role's name map.

Ai00-X · Nov 10, 2023 · 2ef9986 · 2ef9986
1 parent c6be083
commit 2ef9986
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/oai/chat.rs b/src/oai/chat.rs
@@ -109,7 +109,11 @@ impl From<ChatRequest> for GenerateRequest {
             .map(|record| record.content)
             .join("\n\n");
 
-        let assistant = Role::Assistant.to_string();
+        let assistant = Role::Assistant;
+        let assistant = names
+            .get(&assistant)
+            .cloned()
+            .unwrap_or(assistant.to_string());
         let prompt = prompt + &format!("\n\n{assistant}:");
 
         let max_tokens = max_tokens.min(crate::MAX_TOKENS);