feat: add context and placeholders for system prompt (#994)

* feat: add context and placeholders for system prompt

* refactor: gemini suggestions

src/equicordplugins/triviaAI/index.tsx

@@ -12,13 +12,13 @@ import { findExportedComponentLazy } from "@webpack";
 import { ChannelStore, Menu } from "@webpack/common";
 
 import { settings } from "./settings";
-import { getResponse, handleResponse, parseMessageContent } from "./utils";
+import { getPayload, getResponse, handleResponse } from "./utils";
 
 const RobotIconLazy = findExportedComponentLazy("RobotIcon");
 const RobotIcon: IconComponent = props => <RobotIconLazy {...props} />;
 
 const messageCtxPatch: NavContextMenuPatchCallback = (children, { message }: { message: Message; }) => {
-    const payload = parseMessageContent(message);
+    const payload = getPayload(message);
     if (!payload) return;
 
     const group = findGroupChildrenByChildId("copy-text", children);
@@ -50,7 +50,7 @@ export default definePlugin({
     messagePopoverButton: {
         icon: RobotIcon,
         render(message: Message) {
-            const payload = parseMessageContent(message);
+            const payload = getPayload(message);
             if (!payload) return null;
 
             return {

src/equicordplugins/triviaAI/settings.ts

@@ -25,9 +25,10 @@ export const settings = definePluginSettings({
     },
     systemPrompt: {
         type: OptionType.STRING,
-        description: "System Prompt for the AI.",
-        default: "You are a helpful assistant who answers questions for the user in a concise and short way while using the least amount of words and punctuation.",
-        placeholder: "Enter system prompt."
+        description: "System Prompt for the AI. Placeholders: {current_user}, {current_time}",
+        default: "You are a helpful assistant who answers questions for the user in a concise and short way while using the least amount of words and punctuation.\nCurrent user: {current_user}\nCurrent time: {current_time}",
+        placeholder: "Enter system prompt.",
+        multiline: true
     },
     maxTokens: {
         type: OptionType.NUMBER,
@@ -40,6 +41,21 @@ export const settings = definePluginSettings({
         default: "https://openrouter.ai/api/v1/chat/completions",
         placeholder: "Enter your OpenAI compatible AI endpoint here."
     },
+    context: {
+        type: OptionType.NUMBER,
+        description: "Number of previous messages to include as context.",
+        default: 0
+    },
+    passMessageAuthorName: {
+        type: OptionType.BOOLEAN,
+        description: "Prepend the author's name to the message content when passing it to the AI. This can help the AI distinguish between different users in a conversation.",
+        default: true
+    },
+    treatSelfAsAssistant: {
+        type: OptionType.BOOLEAN,
+        description: "When enabled, your own messages will be treated as assistant messages in the context. This causes some models to start generating fanfic.",
+        default: false
+    },
     mode: {
         type: OptionType.SELECT,
         description: "How should answers be handled?",

src/equicordplugins/triviaAI/utils.ts

@@ -8,7 +8,7 @@ import { sendBotMessage } from "@api/Commands";
 import { insertTextIntoChatInputBox, sendMessage } from "@utils/discord";
 import { Logger } from "@utils/Logger";
 import { Message } from "@vencord/discord-types";
-import { showToast, Toasts } from "@webpack/common";
+import { MessageStore, showToast, Toasts, UserStore } from "@webpack/common";
 
 import { settings } from "./settings";
 
@@ -29,6 +29,66 @@ type ImagePart = {
 
 export type ContentPayload = string | (TextPart | ImagePart)[];
 
+export type ApiMessage = {
+    role: "user" | "assistant";
+    content: ContentPayload;
+};
+
+export function getPayload(message: Message): ApiMessage[] | null {
+    const prevMessages = getPreviousMessages(message, settings.store.context);
+    const allMessages = [...prevMessages, message];
+
+    const currentUserId = UserStore.getCurrentUser().id;
+
+    const payload: ApiMessage[] = [];
+
+    for (const msg of allMessages) {
+        const parsed = parseMessageContent(msg);
+        if (!parsed) continue;
+
+        const isOwn = msg.author?.id === currentUserId;
+        const isTargetMessage = msg.id === message.id;
+        const role = (isOwn && !isTargetMessage && settings.store.treatSelfAsAssistant) ? "assistant" : "user";
+
+        let content = parsed;
+
+        if (!isOwn && settings.store.passMessageAuthorName) {
+            const username = msg.author?.username ?? "Unknown";
+            const prefix = `${username}: `;
+
+            if (typeof parsed === "string") {
+                content = prefix + parsed;
+            } else if (Array.isArray(parsed)) {
+                content = [...parsed];
+                const firstTextIdx = content.findIndex(p => p.type === "text");
+
+                if (firstTextIdx !== -1) {
+                    content[firstTextIdx] = {
+                        type: "text",
+                        text: prefix + (content[firstTextIdx] as TextPart).text
+                    };
+                } else {
+                    content.unshift({
+                        type: "text",
+                        text: prefix
+                    });
+                }
+            }
+        }
+
+        payload.push({ role, content });
+    }
+
+    return payload.length > 0 ? payload : null;
+}
+
+export function getPreviousMessages(message: Message, count: number): Message[] {
+    const allMessages: Message[] = MessageStore.getMessages(message.channel_id)._array;
+    const idx = allMessages.findIndex(m => m.id === message.id);
+    if (idx <= 0 || count === 0) return [];
+    return allMessages.slice(Math.max(0, idx - count), idx);
+}
+
 export function parseMessageContent(message: Message): ContentPayload | null {
     const textParts: string[] = [];
 
@@ -123,19 +183,18 @@ export async function handleResponse(message: Message, response: string): Promis
     return response;
 }
 
-export async function getResponse(payload: ContentPayload): Promise<string> {
-    if (!settings.store.apiKey) {
-        showToast("TriviaAI: API Key is not set.", Toasts.Type.FAILURE);
-        return "";
-    }
+function getSystemPrompt() {
+    const currentUser = UserStore.getCurrentUser();
+    const currentTime = new Date().toString();
 
-    if (!settings.store.endpoint) {
-        showToast("TriviaAI: Endpoint is not set.", Toasts.Type.FAILURE);
-        return "";
-    }
+    return settings.store.systemPrompt
+        .replace(/{current_user}/g, currentUser?.username ?? "Unknown User")
+        .replace(/{current_time}/g, currentTime);
+}
 
-    if (!settings.store.model) {
-        showToast("TriviaAI: Model is not set.", Toasts.Type.FAILURE);
+export async function getResponse(payload: ApiMessage[]): Promise<string> {
+    if (!settings.store.apiKey || !settings.store.endpoint || !settings.store.model) {
+        showToast("TriviaAI: API settings are incomplete.", Toasts.Type.FAILURE);
         return "";
     }
 
@@ -151,12 +210,9 @@ export async function getResponse(payload: ContentPayload): Promise<string> {
                 messages: [
                     {
                         role: "system",
-                        content: settings.store.systemPrompt
+                        content: getSystemPrompt()
                     },
-                    {
-                        role: "user",
-                        content: payload
-                    }
+                    ...payload
                 ],
                 max_tokens: settings.store.maxTokens,
             })