load token with fn/expression, closes #122

Author: madox2

README.md

@@ -343,6 +343,7 @@ END
 " - options.request_timeout: request timeout in seconds
 " - options.auth_type: API authentication method (bearer, api-key, none)
 " - options.token_file_path: override global token configuration
+" - options.token_load_fn: expression/vim function to load token
 " - options.selection_boundary: selection prompt wrapper (eliminates empty responses, see #20)
 " - ui.paste_mode: use paste mode (see more info in the Notes below)
 let g:vim_ai_complete = {
@@ -358,6 +359,7 @@ let g:vim_ai_complete = {
 \    "stream": 1,
 \    "auth_type": "bearer",
 \    "token_file_path": "",
+\    "token_load_fn": "",
 \    "selection_boundary": "#####",
 \    "initial_prompt": s:initial_complete_prompt,
 \  },
@@ -375,6 +377,7 @@ let g:vim_ai_complete = {
 " - options.request_timeout: request timeout in seconds
 " - options.auth_type: API authentication method (bearer, api-key, none)
 " - options.token_file_path: override global token configuration
+" - options.token_load_fn: expression/vim function to load token
 " - options.selection_boundary: selection prompt wrapper (eliminates empty responses, see #20)
 " - ui.paste_mode: use paste mode (see more info in the Notes below)
 let g:vim_ai_edit = {
@@ -390,6 +393,7 @@ let g:vim_ai_edit = {
 \    "stream": 1,
 \    "auth_type": "bearer",
 \    "token_file_path": "",
+\    "token_load_fn": "",
 \    "selection_boundary": "#####",
 \    "initial_prompt": s:initial_complete_prompt,
 \  },
@@ -415,6 +419,7 @@ END
 " - options.request_timeout: request timeout in seconds
 " - options.auth_type: API authentication method (bearer, api-key, none)
 " - options.token_file_path: override global token configuration
+" - options.token_load_fn: expression/vim function to load token
 " - options.selection_boundary: selection prompt wrapper (eliminates empty responses, see #20)
 " - ui.open_chat_command: preset (preset_below, preset_tab, preset_right) or a custom command
 " - ui.populate_options: dump [chat] config to the chat header
@@ -434,6 +439,7 @@ let g:vim_ai_chat = {
 \    "stream": 1,
 \    "auth_type": "bearer",
 \    "token_file_path": "",
+\    "token_load_fn": "",
 \    "selection_boundary": "",
 \    "initial_prompt": s:initial_chat_prompt,
 \  },
@@ -452,6 +458,7 @@ let g:vim_ai_chat = {
 " - options.request_timeout: request timeout in seconds
 " - options.auth_type: API authentication method (bearer, api-key, none)
 " - options.token_file_path: override global token configuration
+" - options.token_load_fn: expression/vim function to load token
 " - options.download_dir: path to image download directory, `cwd` if not defined
 let g:vim_ai_image = {
 \  "provider": "openai",
@@ -465,6 +472,7 @@ let g:vim_ai_image = {
 \    "request_timeout": 40,
 \    "auth_type": "bearer",
 \    "token_file_path": "",
+\    "token_load_fn": "",
 \  },
 \  "ui": {
 \    "download_dir": "",
@@ -477,6 +485,9 @@ let g:vim_ai_roles_config_file = s:plugin_root . "/roles-example.ini"
 " custom token file location
 let g:vim_ai_token_file_path = "~/.config/openai.token"
 
+" custom fn to load token, e.g. "g:GetAIToken()"
+let g:vim_ai_token_load_fn = ""
+
 " enables/disables full markdown highlighting in aichat files
 " NOTE: code syntax highlighting works out of the box without this option enabled
 " NOTE: highlighting may be corrupted when using together with the `preservim/vim-markdown`

autoload/vim_ai_config.vim

@@ -74,6 +74,7 @@ let g:vim_ai_openai_complete = {
 \  "stream": 1,
 \  "auth_type": "bearer",
 \  "token_file_path": "",
+\  "token_load_fn": "",
 \  "selection_boundary": "#####",
 \  "initial_prompt": s:initial_complete_prompt,
 \}
@@ -88,6 +89,7 @@ let g:vim_ai_openai_chat = {
 \  "stream": 1,
 \  "auth_type": "bearer",
 \  "token_file_path": "",
+\  "token_load_fn": "",
 \  "selection_boundary": "",
 \  "initial_prompt": s:initial_chat_prompt,
 \}
@@ -100,6 +102,7 @@ let g:vim_ai_openai_image = {
 \  "request_timeout": 40,
 \  "auth_type": "bearer",
 \  "token_file_path": "",
+\  "token_load_fn": "",
 \}
 
 if !exists("g:vim_ai_open_chat_presets")
@@ -124,6 +127,9 @@ endif
 if !exists("g:vim_ai_token_file_path")
   let g:vim_ai_token_file_path = "~/.config/openai.token"
 endif
+if !exists("g:vim_ai_token_load_fn")
+  let g:vim_ai_token_load_fn = ""
+endif
 if !exists("g:vim_ai_roles_config_file")
   let g:vim_ai_roles_config_file = s:plugin_root . "/roles-example.ini"
 endif

doc/vim-ai.txt

@@ -42,6 +42,7 @@ Options: >
   \    "stream": 1,
   \    "auth_type": "bearer",
   \    "token_file_path": "",
+  \    "token_load_fn": "",
   \    "selection_boundary": "#####",
   \    "initial_prompt": s:initial_complete_prompt,
   \  },
@@ -82,6 +83,7 @@ Options: >
   \    "stream": 1,
   \    "auth_type": "bearer",
   \    "token_file_path": "",
+  \    "token_load_fn": "",
   \    "selection_boundary": "#####",
   \    "initial_prompt": s:initial_complete_prompt,
   \  },
@@ -120,6 +122,7 @@ Options: >
   \    "stream": 1,
   \    "auth_type": "bearer",
   \    "token_file_path": "",
+  \    "token_load_fn": "",
   \    "selection_boundary": "",
   \    "initial_prompt": s:initial_chat_prompt,
   \  },
@@ -171,6 +174,7 @@ Options: >
   \    "request_timeout": 40,
   \    "auth_type": "bearer",
   \    "token_file_path": "",
+  \    "token_load_fn": "",
   \  },
   \  "ui": {
   \    "download_dir": "",

py/providers/openai.py

@@ -30,6 +30,7 @@ def request(self, messages: Sequence[AIMessage]) -> Iterator[AIResponseChunk]:
             'request_timeout': options['request_timeout'],
             'auth_type': options['auth_type'],
             'token_file_path': options['token_file_path'],
+            'token_load_fn': options['token_load_fn'],
         }
 
         def _flatten_content(messages):
@@ -70,7 +71,11 @@ def _filter_valid_chunks(chunk):
         return filter(_filter_valid_chunks, map(_map_chunk, response))
 
     def _load_api_key(self):
-        raw_api_key = self.utils.load_api_key("OPENAI_API_KEY", self.options['token_file_path'])
+        raw_api_key = self.utils.load_api_key(
+            "OPENAI_API_KEY",
+            token_file_path=self.options['token_file_path'],
+            token_load_fn=self.options['token_load_fn'],
+        )
         # The text is in format of "<api key>,<org id>" and the
         # <org id> part is optional
         elements = raw_api_key.strip().split(",")
@@ -118,6 +123,7 @@ def request_image(self, prompt: str) -> list[AIImageResponseChunk]:
             'request_timeout': options['request_timeout'],
             'auth_type': options['auth_type'],
             'token_file_path': options['token_file_path'],
+            'token_load_fn': options['token_load_fn'],
         }
         openai_options = {
             'model': options['model'],

py/types.py

@@ -22,7 +22,7 @@ def print_debug(self, text: str, *args: Any):
         pass
     def make_known_error(self, message: str):
         pass
-    def load_api_key(self, env_variable: str, file_path: str):
+    def load_api_key(self, env_variable: str, token_file_path: str = "", token_load_fn: str = ""):
         pass
 
 class AIResponseChunk(TypedDict):

py/utils.py

@@ -28,28 +28,44 @@ def print_debug(text, *args):
 class KnownError(Exception):
     pass
 
+def load_token_from_env_variable(env_variable_name):
+    return os.getenv(env_variable_name).strip()
+
+def load_token_from_file_path(path):
+    if not path:
+        return None
+    with open(os.path.expanduser(path), 'r') as file:
+        return file.read().strip()
+
+def load_token_from_fn(expression):
+    if not expression:
+        return None
+    return vim.eval(expression).strip()
+
 class AIProviderUtils():
     def print_debug(self, text, *args):
         print_debug(text, *args)
 
     def make_known_error(self, message: str):
         return KnownError(message)
 
-    def load_api_key(self, env_variable_name: str, token_file_path: str):
-        # token precedence: env variable, config file path, global file path
-        api_key = os.getenv(env_variable_name)
-        if api_key:
-            return api_key
-        try:
-            global_token_file_path = vim.eval("g:vim_ai_token_file_path")
-            token_file_path = token_file_path or global_token_file_path
-            with open(os.path.expanduser(token_file_path), 'r') as file:
-                api_key = file.read()
-        except Exception:
-            pass
-        if not api_key:
-            raise KnownError("Missing API key")
-        return api_key
+    def load_api_key(self, env_variable_name: str, token_file_path: str = "", token_load_fn: str = ""):
+        loaders = (
+            lambda: load_token_from_file_path(token_file_path),
+            lambda: load_token_from_fn(token_load_fn),
+            lambda: load_token_from_env_variable(env_variable_name),
+            lambda: load_token_from_file_path(vim.eval("g:vim_ai_token_file_path")),
+            lambda: load_token_from_fn(vim.eval("g:vim_ai_token_load_fn")),
+        )
+        for loader in loaders:
+            try:
+                api_key = loader()
+                if api_key:
+                    return api_key
+            except Exception:
+                pass
+
+        raise KnownError("Missing API key")
 
 ai_provider_utils = AIProviderUtils()
 

tests/context_test.py

@@ -7,6 +7,7 @@
     "max_tokens": "0",
     "temperature": "1",
     "token_file_path": "",
+    "token_load_fn": "",
     "selection_boundary": "",
     "initial_prompt": "You are a general assistant.",
   },
@@ -27,6 +28,7 @@
     "size": "1024x1024",
     "style": "vivid",
     "token_file_path": "",
+    "token_load_fn": "",
   },
   "ui": {
     "paste_mode": "1",

tests/deprecated_role_syntax_test.py

@@ -11,6 +11,7 @@
     "stream": "1",
     "auth_type": "bearer",
     "token_file_path": "",
+    "token_load_fn": "",
     "selection_boundary": "",
     "initial_prompt": "You are a general assistant.",
   },

tests/mocks/vim.py

@@ -12,6 +12,8 @@ def eval(cmd):
             return os.path.abspath(os.path.join(dirname, '../..'))
         case 'getcwd()':
             return os.path.abspath(os.path.join(dirname, '../..'))
+        case 'g:LoadToken()':
+            return 'fn.secret'
         case _:
             return None
 

tests/resources/example.token

@@ -0,0 +1 @@
+file.secret