diff --git a/src/equicordplugins/triviaAI/index.tsx b/src/equicordplugins/triviaAI/index.tsx index afa941109..2de85c1a9 100644 --- a/src/equicordplugins/triviaAI/index.tsx +++ b/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 => ; 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 { diff --git a/src/equicordplugins/triviaAI/settings.ts b/src/equicordplugins/triviaAI/settings.ts index 1ccccd6ec..1729a05aa 100644 --- a/src/equicordplugins/triviaAI/settings.ts +++ b/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?", diff --git a/src/equicordplugins/triviaAI/utils.ts b/src/equicordplugins/triviaAI/utils.ts index f9f000e9e..b13386e30 100644 --- a/src/equicordplugins/triviaAI/utils.ts +++ b/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 { - 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 { + 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 { messages: [ { role: "system", - content: settings.store.systemPrompt + content: getSystemPrompt() }, - { - role: "user", - content: payload - } + ...payload ], max_tokens: settings.store.maxTokens, })