Skip to content

Commit

Permalink
add support for Codestral by Mistral (#17)
Browse files Browse the repository at this point in the history
  • Loading branch information
bmichotte authored May 30, 2024
1 parent 70ab6d9 commit 694a05c
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 2 deletions.
31 changes: 29 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,14 @@ AI source for [hrsh7th/nvim-cmp](https://github.com/hrsh7th/nvim-cmp)
This is a general purpose AI source for `cmp`, easily adapted to any restapi
supporting remote code completion.

For now, HuggingFace SantaCoder, OpenAI Chat and Google Bard are implemeted.
For now, HuggingFace SantaCoder, OpenAI Chat, Codestral and Google Bard are implemeted.

## Install

### Dependencies

- You will need `plenary.nvim` to use this plugin.
- For using OpenAI or HuggingFace, you will also need `curl`.
- For using Codestral, OpenAI or HuggingFace, you will also need `curl`.
- For using Google Bard, you will need [dsdanielpark/Bard-API](https://github.com/dsdanielpark/Bard-API).

### Using a plugin manager
Expand Down Expand Up @@ -93,6 +93,33 @@ environment, `OPENAI_API_KEY`.

Available models for OpenAI are `gpt-4` and `gpt-3.5-turbo`.

To use Codestral:

```lua
local cmp_ai = require('cmp_ai.config')

cmp_ai:setup({
max_lines = 1000,
provider = 'Codestral',
provider_options = {
model = 'codestral-latest',
},
notify = true,
notify_callback = function(msg)
vim.notify(msg)
end,
run_on_every_keystroke = true,
ignored_file_types = {
-- default is not to ignore
-- uncomment to ignore in lua:
-- lua = true
},
})
```

You will also need to make sure you have the Codestral api key in you
environment, `CODESTRAL_API_KEY`.

To use Google Bard:

```lua
Expand Down
59 changes: 59 additions & 0 deletions lua/cmp_ai/backends/codestral.lua
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
local requests = require('cmp_ai.requests')

Codestral = requests:new(nil)
BASE_URL = 'https://codestral.mistral.ai/v1/fim/completions'

function Codestral:new(o, params)
o = o or {}
setmetatable(o, self)
self.__index = self
self.params = vim.tbl_deep_extend('keep', params or {}, {
model = 'codestral-latest',
temperature = 0.1,
n = 1,
max_tokens = 100,
})

self.api_key = os.getenv('CODESTRAL_API_KEY')
if not self.api_key then
vim.schedule(function()
vim.notify('CODESTRAL_API_KEY environment variable not set', vim.log.levels.ERROR)
end)
self.api_key = 'NO_KEY'
end
self.headers = {
'Authorization: Bearer ' .. self.api_key,
}
return o
end

function Codestral:complete(lines_before, lines_after, cb)
if not self.api_key then
vim.schedule(function()
vim.notify('CODESTRAL_API_KEY environment variable not set', vim.log.levels.ERROR)
end)
return
end
local data = {
prompt = lines_before,
}
data = vim.tbl_deep_extend('keep', data, self.params)
self:Get(BASE_URL, self.headers, data, function(answer)
local new_data = {}
if answer.choices then
for _, response in ipairs(answer.choices) do
local entry = response.message.content
table.insert(new_data, entry)
end
end
cb(new_data)
end)
end

function Codestral:test()
self:complete('def factorial(n)\n if', ' return ans\n', function(data)
dump(data)
end)
end

return Codestral

0 comments on commit 694a05c

Please sign in to comment.