fix: capture ending of the assistant prompt too since we need to erase it from the reponse, ref #65

Author: pminev

ac-local-plugin/code/LocalLlama.cpp

@@ -59,6 +59,8 @@ class ChatSession {
     std::vector<llama::ChatMsg> m_chatMessages;
     size_t m_submittedMessages = 0;
 
+    ac::llama::AntipromptManager m_antiprompt;
+
 public:
     using Schema = sc::StateChatInstance;
 
@@ -103,6 +105,17 @@ class ChatSession {
         // user prefix should a substr before stop
         m_userPrefix = m_chatFormat->formatMsg({.role = m_roleUser, .text = "stop"}, {}, false);
         m_userPrefix = trim(m_userPrefix.substr(0, m_userPrefix.find("stop")));
+        m_antiprompt.addAntiprompt(m_userPrefix);
+
+        std::vector<llama::ChatMsg> msgs{
+            {.role = m_roleAsistant, .text = "pre"},
+            {.role = m_roleUser, .text = "post"},
+        };
+
+        auto assistantEnd = m_chatFormat->formatChat(msgs, false);
+        assistantEnd = assistantEnd.substr(assistantEnd.find("pre") + 3); // 3 because of the length of "pre"
+        assistantEnd = trim(assistantEnd.substr(0, assistantEnd.find("post")));
+        m_antiprompt.addAntiprompt(assistantEnd);
     }
 
     ~ChatSession() {
@@ -149,8 +162,7 @@ class ChatSession {
             m_submittedMessages = m_chatMessages.size();
         }
 
-        ac::llama::AntipromptManager antiprompt;
-        antiprompt.addAntiprompt(m_userPrefix);
+        m_antiprompt.reset();
 
         std::string fullResponse;
         Schema::OpGetChatResponse::Return ret;
@@ -167,7 +179,7 @@ class ChatSession {
             result += tokenStr;
             fullResponse += tokenStr;
 
-            auto matchedAntiPrompt = antiprompt.feedGeneratedText(tokenStr);
+            auto matchedAntiPrompt = m_antiprompt.feedGeneratedText(tokenStr);
             if (!matchedAntiPrompt.empty()) {
                 // and also hide it from the return value
                 // note that we assume that m_userPrefix is always the final piece of text in the response
@@ -176,7 +188,7 @@ class ChatSession {
                 break;
             }
 
-            if (isStreaming && !antiprompt.hasRunningAntiprompts()) {
+            if (isStreaming && !m_antiprompt.hasRunningAntiprompts()) {
                 co_await m_io.push(Frame_from(sc::StreamToken{}, result));
                 result = {};
             }

ac-local-plugin/example/ep-chat.cpp

@@ -46,10 +46,11 @@ int main() try {
                         "{{ '<|' + assistant_role + '|>\\n' }}"
                         "{% endif %}";
 
+    constexpr bool useChatTemplate = false;
     sid = llama.call<schema::StateModelLoaded::OpStartInstance>({
         .instanceType = "chat",
         .setup = "A chat between a human user and a helpful AI assistant.",
-        .chatTemplate = chatTemplate,
+        .chatTemplate = useChatTemplate ? chatTemplate : "",
         .roleUser = roleUser,
         .roleAssistant = roleAssistant,
     });

code/ac/llama/AntipromptManager.cpp

@@ -10,15 +10,22 @@ void AntipromptManager::addAntiprompt(std::string_view antiprompt) {
 }
 
 std::string AntipromptManager::feedGeneratedText(std::string_view text) {
+    std::vector<std::pair<std::string, size_t>> matchedAntiprompts;
     for (auto& ap : m_antiprompts) {
         int found = ap.feedText(text);
         if (found > 0) {
-            reset();
-            return found == 0 ?
-                        ap.getString():
-                        ap.getString() + std::string(text.substr(found, text.length()));
+            auto res = found == 0 ?
+                ap.getString():
+                ap.getString() + std::string(text.substr(found, text.length()));
+            matchedAntiprompts.push_back({res, found});
         }
     }
+    if (!matchedAntiprompts.empty()) {
+        reset();
+        std::sort(matchedAntiprompts.begin(), matchedAntiprompts.end());
+        auto& [res, found] = matchedAntiprompts.front();
+        return res;
+    }
 
     return {};
 }