using openai api directly

4 files changed, 42 insertions, 11 deletions

diff --git a/README.md b/README.md
index 59d559d..20475c8 100644
--- a/README.md
+++ b/README.md
@@ -9,9 +9,6 @@ Complete text in vim using OpenAI.
 Prerequisites:
 
 ```sh
-# install https://github.com/peterdemin/openai-cli
-pip install openai-cli
-
 # configure openai api key https://platform.openai.com/account/api-keys
 echo "YOUR_OPENAPI_TOKEN" > ~/.config/openai.token
 
@@ -19,12 +16,14 @@ echo "YOUR_OPENAPI_TOKEN" > ~/.config/openai.token
 export OPENAI_API_TOKEN="YOUR_OPENAPI_TOKEN"
 ```
 
-Add plugin to your `.vimrc`, e.g. using `vim-plug`:
+Add plugin to your `.vimrc` using `vim-plug`:
 
 ```vim
-Plug 'madox2/vim-ai'
+Plug 'madox2/vim-ai', { 'do': './install.sh' }
 ```
 
+The plugin requires `python3` and `pip` to be installed.
+
 ## Usage
 
 ### Basic usage
diff --git a/install.sh b/install.sh
new file mode 100755
index 0000000..ed2d5b5
--- /dev/null
+++ b/install.sh
@@ -0,0 +1,9 @@
+#!/bin/bash
+
+if [ -z "$(which pip)" ]; then
+  echo pip not found
+  exit 1
+fi
+
+echo installing deps using pip
+pip install openai
diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index 974fa19..134f878 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -1,8 +1,5 @@
-if $OPENAI_API_TOKEN != ''
-  let g:openaiToken = $OPENAI_API_TOKEN
-else
-  let g:openaiToken = system("cat ~/.config/openai.token")
-endif
+let s:plugin_root = expand('<sfile>:p:h:h')
+let s:complete_script_file = s:plugin_root . "/py/complete.py"
 
 function! AIRun(...) range
   let prompt = getline(a:firstline, a:lastline)
@@ -16,7 +13,7 @@ function! AIRun(...) range
   let prompt = join(prompt, "\n")
 
   echo "Completing..."
-  let output = system("echo " . shellescape(prompt) . " | openai complete - -t " . g:openaiToken)
+  let output = system("echo " . shellescape(prompt) . " | python3 " . s:complete_script_file . " ")
   let output = trim(output)
 
   execute a:firstline . ',' . a:lastline . 'd'
diff --git a/py/complete.py b/py/complete.py
new file mode 100644
index 0000000..ba460c4
--- /dev/null
+++ b/py/complete.py
@@ -0,0 +1,26 @@
+import sys
+import os
+import openai
+
+config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")
+
+api_key = os.getenv("OPENAI_API_KEY")
+
+try:
+    with open(config_file_path, 'r') as file:
+        api_key = file.read()
+except Exception:
+    pass
+
+openai.api_key = api_key.strip()
+
+prompt = "".join(sys.stdin.readlines())
+
+completion = openai.Completion.create(
+    model="text-davinci-003",
+    prompt=prompt,
+    max_tokens=1000,
+    temperature=0.1
+)
+
+print(completion.choices[0].text)