Add custom prompts and change translation behavior + Refactor

README.md

@@ -43,24 +43,45 @@ return CONFIGURATION
 
 Additionally, as other extra features are rolled out, they will be optional and can be set in the `features` table in the `configuration.lua` file.
 
-
 ### Translation
 
-To enable translation, you can set the `translate_to` parameter in the `features` table. For example, if you want to translate the text to French, you can set the `translate_to` parameter to `"French"`.
+To enable translation, simply add a translation prompt to your features.custom_prompts configuration. The prompt should specify the target language and any other translation preferences. For example this one is for French:
+
+```lua
+local CONFIGURATION = {
+    api_key = "YOUR_API_KEY",
+    model = "gpt-4o-mini",
+    base_url = "https://api.openai.com/v1/chat/completions",
+    features = {
+        custom_prompts = {
+            translation = "Please translate the following text to French with definition."
+        }
+    }
+}
+```
+
+### Custom Prompts
+
+You can customize the prompts used for different interactions by adding them to the `features.custom_prompts` section. Each prompt (except 'system') will automatically generate a button in the interface. The button's name will be the capitalized version of the prompt key.
 
-By setting the `translate_to` parameter, you can have the plugin translate the text to the language you specify. This is useful if you are reading a book in a language you are not fluent in and want to understand a chunk of text in a language you are more comfortable with.
+For example, this configuration:
 
 ```lua
 local CONFIGURATION = {
     api_key = "YOUR_API_KEY",
     model = "gpt-4o-mini",
     base_url = "https://api.openai.com/v1/chat/completions",
     features = {
-        translate_to = "French"
+        custom_prompts = {
+            summarize = "Please summarize the following text.",
+            translate = "Please translate the following text to French."
+        }
     }
 }
 ```
 
+Will generate buttons labeled "Summarize" and "Translate" in the interface.
+
 ## Installation
 
 If you clone this project, you should be able to put the directory, `askgpt.koplugin`, in the `koreader/plugins` directory and it should work. If you want to use the plugin without cloning the project, you can download the zip file from the releases page and extract the `askgpt.koplugin` directory to the `koreader/plugins` directory. If for some reason you extract the files of this repository in another directory, rename it before moving it to the `koreader/plugins` directory.

configuration.lua.sample

@@ -3,6 +3,14 @@ local CONFIGURATION = {
     model = "gpt-4o-mini",
     base_url = "https://api.openai.com/v1/chat/completions",
     additional_parameters = {},
+    features = {
+        custom_prompts = {
+            system = "You are a helpful AI assistant analyzing text from a book. Be concise and insightful in your responses.",
+            translation = "You are a skilled translator and language expert. For each text: 1) Provide an accurate translation to French, 2) Then explain its meaning and context in the original language to help with understanding.",
+            summarize = "Provide a concise summary of the following text, highlighting the key points.",
+            explain = "Explain this text as if teaching it to someone, breaking down complex concepts."
+        }
+    }
 }
 
 return CONFIGURATION

dialogs.lua

@@ -16,15 +16,25 @@ else
   print("configuration.lua not found, skipping...")
 end
 
-local function translateText(text, target_language)
+local function translateText(text)
+  if not text or text == "" then
+    return _("Error: No text provided for translation")
+  end
+
+  local translation_prompt = CONFIGURATION 
+    and CONFIGURATION.features 
+    and CONFIGURATION.features.custom_prompts 
+    and CONFIGURATION.features.custom_prompts.translation
+    or "You are a skilled translator and language expert. First explain the meaning of the text in the original language, then provide an accurate translation to English."
+
   local translation_message = {
     role = "user",
-    content = "Translate the following text to " .. target_language .. ": " .. text
+    content = "For the following text:\n\n\"" .. text .. "\"\n\n1. First explain what this means/expresses in the original language\n2. Then translate it as specified in your instructions"
   }
   local translation_history = {
     {
       role = "system",
-      content = "You are a helpful translation assistant. Provide direct translations without additional commentary."
+      content = translation_prompt
     },
     translation_message
   }
@@ -54,12 +64,27 @@ local function showLoadingDialog()
 end
 
 local function showChatGPTDialog(ui, highlightedText, message_history)
+  if not highlightedText or highlightedText == "" then
+    UIManager:show(InfoMessage:new{
+      text = _("Please highlight some text first."),
+    })
+    return
+  end
+
   local title, author =
     ui.document:getProps().title or _("Unknown Title"),
     ui.document:getProps().authors or _("Unknown Author")
+
+  local default_prompt = "The following is a conversation with an AI assistant. The assistant is helpful, creative, clever, and very friendly. Answer as concisely as possible."
+  local system_prompt = CONFIGURATION 
+    and CONFIGURATION.features 
+    and CONFIGURATION.features.custom_prompts 
+    and CONFIGURATION.features.custom_prompts.system
+    or default_prompt
+
   local message_history = message_history or {{
     role = "system",
-    content = "The following is a conversation with an AI assistant. The assistant is helpful, creative, clever, and very friendly. Answer as concisely as possible."
+    content = system_prompt
   }}
 
   local function handleNewQuestion(chatgpt_viewer, question)
@@ -97,7 +122,7 @@ local function showChatGPTDialog(ui, highlightedText, message_history)
         UIManager:scheduleIn(0.1, function()
           local context_message = {
             role = "user",
-            content = "I'm reading something titled '" .. title .. "' by " .. author ..
+            content = "I'm reading something titled '" .. title .. "' by " .. author .. 
               ". I have a question about the following highlighted text: " .. highlightedText
           }
           table.insert(message_history, context_message)
@@ -129,36 +154,56 @@ local function showChatGPTDialog(ui, highlightedText, message_history)
     }
   }
 
-  if CONFIGURATION and CONFIGURATION.features and CONFIGURATION.features.translate_to then
-    table.insert(buttons, {
-      text = _("Translate"),
-      callback = function()
-        showLoadingDialog()
-
-        UIManager:scheduleIn(0.1, function()
-          local translated_text = translateText(highlightedText, CONFIGURATION.features.translate_to)
-
-          table.insert(message_history, {
-            role = "user",
-            content = "Translate to " .. CONFIGURATION.features.translate_to .. ": " .. highlightedText
-          })
-
-          table.insert(message_history, {
-            role = "assistant",
-            content = translated_text
-          })
-
-          local result_text = createResultText(highlightedText, message_history)
-          local chatgpt_viewer = ChatGPTViewer:new {
-            title = _("Translation"),
-            text = result_text,
-            onAskQuestion = handleNewQuestion
-          }
-
-          UIManager:show(chatgpt_viewer)
-        end)
+  -- Add buttons for each custom prompt
+  if CONFIGURATION and CONFIGURATION.features and CONFIGURATION.features.custom_prompts then
+    for prompt_name, prompt in pairs(CONFIGURATION.features.custom_prompts) do
+      if prompt_name ~= "system" and type(prompt) == "string" then  -- Ensure prompt is valid
+        table.insert(buttons, {
+          text = _(prompt_name:gsub("^%l", string.upper)),  -- Capitalize first letter
+          callback = function()
+            showLoadingDialog()
+
+            UIManager:scheduleIn(0.1, function()
+              message_history = {  -- Reset message history for new prompt
+                {
+                  role = "system",
+                  content = prompt
+                },
+                {
+                  role = "user",
+                  content = "I'm reading something titled '" .. title .. "' by " .. author .. 
+                    "'. Here's the text I want you to process: " .. highlightedText
+                }
+              }
+
+              local success, response = pcall(queryChatGPT, message_history)
+
+              if not success then
+                UIManager:show(InfoMessage:new{
+                  text = _("Error: Failed to get response from ChatGPT"),
+                })
+                return
+              end
+
+              table.insert(message_history, { 
+                role = "assistant", 
+                content = response 
+              })
+
+              local result_text = createResultText(highlightedText, message_history)
+
+              local chatgpt_viewer = ChatGPTViewer:new {
+                title = _(prompt_name:gsub("^%l", string.upper)),
+                text = result_text,
+                onAskQuestion = handleNewQuestion
+              }
+
+              UIManager:show(chatgpt_viewer)
+            end)
+          end
+        })
       end
-    })
+    end
   end
 
   input_dialog = InputDialog:new{

gpt_query.lua

@@ -24,8 +24,16 @@ local ltn12 = require("ltn12")
 local json = require("json")
 
 local function queryChatGPT(message_history)
+  if not message_history or #message_history == 0 then
+    error("No message history provided")
+  end
+
   -- Use api_key from CONFIGURATION or fallback to the api_key module
   local api_key_value = CONFIGURATION and CONFIGURATION.api_key or api_key
+  if not api_key_value then
+    error("No API key configured")
+  end
+
   local api_url = CONFIGURATION and CONFIGURATION.base_url or "https://api.openai.com/v1/chat/completions"
   local model = CONFIGURATION and CONFIGURATION.model or "gpt-4o-mini"
 
@@ -65,11 +73,15 @@ local function queryChatGPT(message_history)
   }
 
   if code ~= 200 then
-    error("Error querying ChatGPT API: " .. code)
+    error("Error querying ChatGPT API: " .. (code or "unknown error"))
   end
 
   local response = json.decode(table.concat(responseBody))
-  return response.choices[1].message.content
+  if not response or not response.choices or not response.choices[1] or not response.choices[1].message then
+    error("Invalid response format from API")
+  end
+
+  return response.choices[1].message.content or "No response content"
 end
 
 return queryChatGPT