From 2ef9986d95b1f6508ab438b619684e8fde87442d Mon Sep 17 00:00:00 2001
From: cryscan <crycan@umich.edu>
Date: Fri, 10 Nov 2023 13:47:25 +0800
Subject: [PATCH] Fix assistant role's name map.

---
 src/oai/chat.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/oai/chat.rs b/src/oai/chat.rs
index 6dbe6d96..8776a5f3 100644
--- 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);