🐛 fix: Fix context for MessageBuilder

Palm1r · Dec 17, 2024 · fc5e1ad · fc5e1ad
1 parent 93e59fb
commit fc5e1ad
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 71 deletions.
diff --git a/llmcore/MessageBuilder.cpp b/llmcore/MessageBuilder.cpp
@@ -1,3 +1,92 @@
+/* 
+ * Copyright (C) 2024 Petr Mironychev
+ *
+ * This file is part of QodeAssist.
+ *
+ * QodeAssist is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * QodeAssist is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with QodeAssist. If not, see <https://www.gnu.org/licenses/>.
+ */
+
 #include "MessageBuilder.hpp"
 
-namespace QodeAssist::LLMCore {} // namespace QodeAssist::LLMCore
+QodeAssist::LLMCore::MessageBuilder &QodeAssist::LLMCore::MessageBuilder::addSystemMessage(
+    const QString &content)
+{
+    m_systemMessage = content;
+    return *this;
+}
+
+QodeAssist::LLMCore::MessageBuilder &QodeAssist::LLMCore::MessageBuilder::addUserMessage(
+    const QString &content)
+{
+    m_messages.append({MessageRole::User, content});
+    return *this;
+}
+
+QodeAssist::LLMCore::MessageBuilder &QodeAssist::LLMCore::MessageBuilder::addSuffix(
+    const QString &content)
+{
+    m_suffix = content;
+    return *this;
+}
+
+QodeAssist::LLMCore::MessageBuilder &QodeAssist::LLMCore::MessageBuilder::addtTokenizer(
+    PromptTemplate *promptTemplate)
+{
+    m_promptTemplate = promptTemplate;
+    return *this;
+}
+
+QString QodeAssist::LLMCore::MessageBuilder::roleToString(MessageRole role) const
+{
+    switch (role) {
+    case MessageRole::System:
+        return ROLE_SYSTEM;
+    case MessageRole::User:
+        return ROLE_USER;
+    case MessageRole::Assistant:
+        return ROLE_ASSISTANT;
+    default:
+        return ROLE_USER;
+    }
+}
+
+void QodeAssist::LLMCore::MessageBuilder::saveTo(QJsonObject &request, ProvidersApi api)
+{
+    if (!m_promptTemplate) {
+        return;
+    }
+
+    ContextData context{
+        m_messages.isEmpty() ? QString() : m_messages.last().content, m_suffix, m_systemMessage};
+
+    if (api == ProvidersApi::Ollama) {
+        if (m_promptTemplate->type() == TemplateType::Fim) {
+            m_promptTemplate->prepareRequest(request, context);
+        } else {
+            QJsonArray messages;
+
+            messages.append(QJsonObject{{"role", "system"}, {"content", m_systemMessage}});
+            messages.append(QJsonObject{{"role", "user"}, {"content", m_messages.last().content}});
+            request["messages"] = messages;
+            m_promptTemplate->prepareRequest(request, context);
+        }
+    } else if (api == ProvidersApi::OpenAI) {
+        QJsonArray messages;
+
+        messages.append(QJsonObject{{"role", "system"}, {"content", m_systemMessage}});
+        messages.append(QJsonObject{{"role", "user"}, {"content", m_messages.last().content}});
+        request["messages"] = messages;
+        m_promptTemplate->prepareRequest(request, context);
+    }
+}
diff --git a/llmcore/MessageBuilder.hpp b/llmcore/MessageBuilder.hpp
@@ -47,76 +47,17 @@ struct Message
 class MessageBuilder
 {
 public:
-    MessageBuilder &addSystemMessage(const QString &content)
-    {
-        m_systemMessage = content;
-        return *this;
-    }
-
-    MessageBuilder &addUserMessage(const QString &content)
-    {
-        m_messages.append({MessageRole::User, content});
-        return *this;
-    }
-
-    MessageBuilder &addSuffix(const QString &content)
-    {
-        m_suffix = content;
-        return *this;
-    }
-
-    MessageBuilder &addtTokenizer(PromptTemplate *promptTemplate)
-    {
-        m_promptTemplate = promptTemplate;
-        return *this;
-    }
-
-    QString roleToString(MessageRole role) const
-    {
-        switch (role) {
-        case MessageRole::System:
-            return ROLE_SYSTEM;
-        case MessageRole::User:
-            return ROLE_USER;
-        case MessageRole::Assistant:
-            return ROLE_ASSISTANT;
-        default:
-            return ROLE_USER;
-        }
-    }
-
-    void saveTo(QJsonObject &request, ProvidersApi api)
-    {
-        if (!m_promptTemplate) {
-            return;
-        }
-
-        if (api == ProvidersApi::Ollama) {
-            ContextData context{
-                m_messages.isEmpty() ? QString() : m_messages.last().content,
-                m_suffix,
-                m_systemMessage};
-
-            if (m_promptTemplate->type() == TemplateType::Fim) {
-                m_promptTemplate->prepareRequest(request, context);
-            } else {
-                QJsonArray messages;
-
-                messages.append(QJsonObject{{"role", "system"}, {"content", m_systemMessage}});
-                messages.append(
-                    QJsonObject{{"role", "user"}, {"content", m_messages.last().content}});
-                request["messages"] = messages;
-                m_promptTemplate->prepareRequest(request, {});
-            }
-        } else if (api == ProvidersApi::OpenAI) {
-            QJsonArray messages;
-
-            messages.append(QJsonObject{{"role", "system"}, {"content", m_systemMessage}});
-            messages.append(QJsonObject{{"role", "user"}, {"content", m_messages.last().content}});
-            request["messages"] = messages;
-            m_promptTemplate->prepareRequest(request, {});
-        }
-    }
+    MessageBuilder &addSystemMessage(const QString &content);
+
+    MessageBuilder &addUserMessage(const QString &content);
+
+    MessageBuilder &addSuffix(const QString &content);
+
+    MessageBuilder &addtTokenizer(PromptTemplate *promptTemplate);
+
+    QString roleToString(MessageRole role) const;
+
+    void saveTo(QJsonObject &request, ProvidersApi api);
 
 private:
     QString m_systemMessage;