diff --git a/README.md b/README.md index ce097f3..62407b3 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,17 @@ let g:ollama_model = 'qwen2.5-coder:3b' " Deepseek-coder-v2 let g:ollama_model = 'deepseek-coder-v2:16b-lite-base-q4_0' + +OpenAI-like FIM interface + +``` +let g:ollama_host = 'https://api.siliconflow.cn/v1' +let g:ollama_model = 'deepseek-ai/DeepSeek-V2.5' +let g:ollama_key = 'sk-xxxxxxxxxxxxx' +let g:ollama_context_lines = 100 +``` + +./autoload/ollama.vim ``` | Variable | Default | Description | diff --git a/autoload/ollama.vim b/autoload/ollama.vim index c1e8d09..a9c4645 100644 --- a/autoload/ollama.vim +++ b/autoload/ollama.vim @@ -107,7 +107,7 @@ function! ollama#GetSuggestion(timer) let s:timer_id = -1 let l:current_line = line('.') let l:current_col = col('.') - let l:context_lines = 30 + let l:context_lines = g:ollama_context_lines " Get the lines before and after the current line let l:prefix_lines = getline(max([1, l:current_line - l:context_lines]), l:current_line - 1) @@ -129,7 +129,7 @@ function! ollama#GetSuggestion(timer) let l:prompt = l:prefix . '' . l:suffix let l:model_options = substitute(json_encode(g:ollama_model_options), "\"", "\\\"", "g") - call ollama#logger#Debug("Connecting to Ollama on ".g:ollama_host." using model ".g:ollama_model) + call ollama#logger#Debug("Connecting to Ollama on ".g:ollama_host." using model ".g:ollama_model." api_type".g:ollama_api_type) call ollama#logger#Debug("model_options=".l:model_options) " Convert plugin debug level to python logger levels let l:log_level = ollama#logger#PythonLogLevel(g:ollama_debug) @@ -138,7 +138,9 @@ function! ollama#GetSuggestion(timer) \ "-m", g:ollama_model, \ "-u", g:ollama_host, \ "-o", l:model_options, - \ "-l", l:log_level + \ "-l", l:log_level, + \ "-t", g:ollama_api_type, + \ "-s", g:ollama_key, \ ] call ollama#logger#Debug("command=". join(l:command, " ")) let l:job_options = { diff --git a/plugin/ollama.vim b/plugin/ollama.vim index cb5df22..9ed336f 100644 --- a/plugin/ollama.vim +++ b/plugin/ollama.vim @@ -40,6 +40,21 @@ if !exists('g:ollama_timeout') let g:ollama_timeout = 10 endif +if !exists('g:ollama_key') + " default code completion model + let g:ollama_key = '' +endif + +if !exists('g:ollama_api_type') + " default to ollama + let g:ollama_api_type = 'ollama' +endif + +if !exists('g:ollama_context_lines') + " default to 30 + let g:ollama_context_lines = 30 +endif + " Defines the color scheme for ollama suggestions function! s:ColorScheme() abort if &t_Co == 256 diff --git a/python/complete.py b/python/complete.py index f302613..72d763d 100755 --- a/python/complete.py +++ b/python/complete.py @@ -10,7 +10,7 @@ # Default values DEFAULT_HOST = 'http://localhost:11434' DEFAULT_MODEL = 'codellama:code' -DEFAULT_OPTIONS = '{ "temperature": 0, "top_p": 0.95 }' +DEFAULT_OPTIONS = '{ "temperature": 0, "top_p": 0.95}' # create logger log = OllamaLogger('ollama.log') @@ -93,12 +93,45 @@ def generate_code_completion(config, prompt, baseurl, model, options): else: raise Exception(f"Error: {response.status_code} - {response.text}") +def generate_code_completion_openai(prompt, url, key, model, options, max_tokens=256): + parts = prompt.split('') + + if len(parts) != 2: + log.error("Prompt does not contain ''.") + sys.exit(1) + + newprompt = parts[0] + new_suffix = parts[1] + log.debug(f"newprompt: {newprompt}, new_suffix: {new_suffix}") + + from openai import OpenAI + client = OpenAI(api_key=key, base_url=url) + + try: + response = client.completions.create( + model=model, + prompt=newprompt, + suffix=new_suffix, + stream=False, + max_tokens=max_tokens, + **options + + ) + + text = response.choices[0].text.strip() + return text + except Exception as e: + log.error(f"Error: {e}") + sys.exit(1) + if __name__ == "__main__": parser = argparse.ArgumentParser(description="Complete code with Ollama LLM.") parser.add_argument('-m', '--model', type=str, default=DEFAULT_MODEL, help="Specify the model name to use.") parser.add_argument('-u', '--url', type=str, default=DEFAULT_HOST, help="Specify the base endpoint URL to use (default="+DEFAULT_HOST+")") parser.add_argument('-o', '--options', type=str, default=DEFAULT_OPTIONS, help="Specify the Ollama REST API options.") parser.add_argument('-l', '--log-level', type=int, default=OllamaLogger.ERROR, help="Specify log level") + parser.add_argument('-t', '--api-type', type=str, default='ollama', help="ollama or openai") + parser.add_argument('-s', '--key', type=str, default='', help="please specify the key") args = parser.parse_args() log.setLevel(args.log_level) @@ -112,8 +145,14 @@ def generate_code_completion(config, prompt, baseurl, model, options): # strip suffix (e.g ':7b-code') from modelname modelname = args.model modelname = modelname.rsplit(':', 1)[0] - config = load_config(modelname) prompt = sys.stdin.read() - response = generate_code_completion(config, prompt, args.url, args.model, options) + if args.api_type == 'ollama': + config = load_config(modelname) + response = generate_code_completion(config, prompt, args.url, args.model, options) + elif args.api_type == 'openai': + response = generate_code_completion_openai(prompt, args.url, args.key, args.model, options) + else: + log.error("Please specify the correct api type") + sys.exit(1) print(response, end='')