diff --git a/src/components/Option/Settings/general-settings.tsx b/src/components/Option/Settings/general-settings.tsx index 21f622a..c37b551 100644 --- a/src/components/Option/Settings/general-settings.tsx +++ b/src/components/Option/Settings/general-settings.tsx @@ -29,6 +29,8 @@ export const GeneralSettings = () => { false ) + const [generateTitle, setGenerateTitle] = useStorage("titleGenEnabled", false) + const [hideCurrentChatModelSettings, setHideCurrentChatModelSettings] = useStorage("hideCurrentChatModelSettings", false) @@ -141,6 +143,19 @@ export const GeneralSettings = () => { /> +
+
+ + {t("generalSettings.settings.generateTitle.label")} + +
+ + setGenerateTitle(checked)} + /> +
+
{t("generalSettings.settings.darkMode.label")} diff --git a/src/hooks/chat-helper/index.ts b/src/hooks/chat-helper/index.ts index f2426ee..18f89be 100644 --- a/src/hooks/chat-helper/index.ts +++ b/src/hooks/chat-helper/index.ts @@ -1,5 +1,6 @@ import { saveHistory, saveMessage } from "@/db" import { setLastUsedChatModel } from "@/services/model-settings" +import { generateTitle } from "@/services/title" import { ChatHistory } from "@/store/option" export const saveMessageOnError = async ({ @@ -14,7 +15,7 @@ export const saveMessageOnError = async ({ setHistoryId, isRegenerating, message_source = "web-ui", - message_type + message_type }: { e: any setHistory: (history: ChatHistory) => void @@ -73,7 +74,8 @@ export const saveMessageOnError = async ({ ) await setLastUsedChatModel(historyId, selectedModel) } else { - const newHistoryId = await saveHistory(userMessage, false, message_source) + const title = await generateTitle(selectedModel, userMessage, userMessage) + const newHistoryId = await saveHistory(title, false, message_source) if (!isRegenerating) { await saveMessage( newHistoryId.id, @@ -154,7 +156,8 @@ export const saveMessageOnSuccess = async ({ ) await setLastUsedChatModel(historyId, selectedModel!) } else { - const newHistoryId = await saveHistory(message, false, message_source) + const title = await generateTitle(selectedModel, message, message) + const newHistoryId = await saveHistory(title, false, message_source) await saveMessage( newHistoryId.id, selectedModel, diff --git a/src/models/index.ts b/src/models/index.ts index d627a5d..4048cb3 100644 --- a/src/models/index.ts +++ b/src/models/index.ts @@ -13,12 +13,12 @@ export const pageAssistModel = async ({ }: { model: string baseUrl: string - keepAlive: string - temperature: number - topK: number - topP: number - numCtx: number - seed: number + keepAlive?: string + temperature?: number + topK?: number + topP?: number + numCtx?: number + seed?: number }) => { switch (model) { case "chrome::gemini-nano::page-assist": diff --git a/src/services/title.ts b/src/services/title.ts new file mode 100644 index 0000000..c7a9cd4 --- /dev/null +++ b/src/services/title.ts @@ -0,0 +1,72 @@ +import { pageAssistModel } from "@/models" +import { Storage } from "@plasmohq/storage" +import { getOllamaURL } from "./ollama" +import { cleanUrl } from "@/libs/clean-url" +import { HumanMessage } from "langchain/schema" +const storage = new Storage() + +// this prompt is copied from the OpenWebUI codebase +export const DEFAULT_TITLE_GEN_PROMPT = `Here is the query: + +-------------- + +{{query}} + +-------------- + +Create a concise, 3-5 word phrase as a title for the previous query. Suitable emojis for the summary can be used to enhance understanding. Avoid quotation marks or special formatting. RESPOND ONLY WITH THE TITLE TEXT. ANSWER USING THE SAME LANGUAGE AS THE QUERY. + + +Examples of titles: + +Stellar Achievement Celebration +Family Bonding Activities +🇫🇷 Voyage à Paris +🍜 Receta de Ramen Casero +Shakespeare Analyse Literarische +日本の春祭り体験 +Древнегреческая Философия Обзор + +Response:` + + +export const isTitleGenEnabled = async () => { + const enabled = await storage.get("titleGenEnabled") + return enabled ?? false +} + +export const setTitleGenEnabled = async (enabled: boolean) => { + await storage.set("titleGenEnabled", enabled) +} + + +export const generateTitle = async (model: string, query: string, fallBackTitle: string) => { + + const isEnabled = await isTitleGenEnabled() + + if (!isEnabled) { + return fallBackTitle + } + + try { + const url = await getOllamaURL() + + const titleModel = await pageAssistModel({ + baseUrl: cleanUrl(url), + model + }) + + const prompt = DEFAULT_TITLE_GEN_PROMPT.replace("{{query}}", query) + + const title = await titleModel.invoke([ + new HumanMessage({ + content: prompt + }) + ]) + + return title.content.toString() + } catch (error) { + console.log(`Error generating title: ${error}`) + return fallBackTitle + } +} \ No newline at end of file