Add Codestral edit predictions provider

[dvcrn]

Attempt of adding Codestral through Mistral API as edit predictions provider. Heavily inspired (or more like copy pasted) by the Supermaven provider

Config

Get API key from https://console.mistral.ai/codestral

  "features": {
    "edit_prediction_provider": "codestral"
  },
  "edit_predictions": {
    "codestral": {
      "api_key": "xxx",
      "model": "codestral-latest",
      "max_tokens": 150
    }
  },

Disclaimer: I'm not a rust engineer and this is my first time dealing with the Zed repo. AI helped me a lot with figuring things out.

[cla-bot]

We require contributors to sign our Contributor License Agreement, and we don't have @dvcrn on file. You can sign our CLA at https://zed.dev/cla. Once you've signed, post a comment here that says '@cla-bot check'.

[dvcrn]

@cla-bot check

[cla-bot]

The cla-bot has been summoned, and re-checked this pull request!

[dvcrn]

• Updated to use TreeSitter
• Add support for subtle mode

Found a bug: using multiple cursors is messing with things and generates a bunch of duplicates stuff

[remorses]

Thank you for this PR! It works very well

Here are some suggestions

• Replace <|end_of_file|> tokens from the output. I noticed <|end_of_file|> gets added often for predictions at the end of the file
• Use the Mistral API key from the existing API settings panel instead of having to add the API key in the settings.json

[dvcrn]

Use the Mistral API key from the existing API settings panel instead of having to add the API key in the settings.json

The issue here is that Mistral API and Codestral use different API keys. It works if the user already has their Codestral API key in Zed, but this usually means that the normal Mistral functionality doesn't work. Maybe we need to add a separate key for Codestral specific?

[remorses]

You are right. I think it's fine to keep the API key in the settings. Maybe someone from Mistral can open a PR to add their own OAuth system like Supermaven some day

I am using this branch as my daily driver, this is the best inline completion provider available in Zed right now, really thank you

[dvcrn]

Thank you for this PR! It works very well

Here are some suggestions

• Replace <|end_of_file|> tokens from the output. I noticed <|end_of_file|> gets added often for predictions at the end of the file

• Use the Mistral API key from the existing API settings panel instead of having to add the API key in the settings.json

Do you have some examples of the end_of_file tokens getting inserted?

[dvcrn]

• Rebased this on top of latest main
• Updated the provider to use edit_prediction instead of inline_completion
• Added light event tracking

Please give this a go! What changes are missing to get this merged?

[mgsloan]

@dvcrn Thanks for working on this! Sorry about the huge review delay. I've pushed some improvements. Overall it seems pretty reasonable, and seems to mostly work (I have seen some funky large predictions that seem like bugs, but not worth blocking merge on it).

I think it'd be good to remove the context header for now to make this change smaller / more reviewable. The logic differs from zeta and some of the constants seem wild - like 80 lines of diff is a lot (arbitrarily large if there are long lines). Also not sure about wrapping it in /*. Can add something like it back in in a followup. Also in our work on zeta2 there will probably be some new code for handling edit history.

The API key really shouldn't be in the settings, it should probably be an additional API key entry in the Mistral section of the agent settings. That would be a fair bit of work to add, though, so not going to block this PR on it.

[mgsloan]

Going to go ahead and make the changes I described

[dvcrn]

Please feel free to adjust whatever you see fit, I trust your judgement on that 😄

With /* wrapping you mean for the additional context, right? I wanted to make sure the model understands that this is just additional context, and since it's trained on code, using codeblocks felt right

The API key really shouldn't be in the settings, it should probably be an additional API key entry in the Mistral section of the agent settings. That would be a fair bit of work to add, though, so not going to block this PR on it.

Yes agree. I was considering re-using the mistral API key settings, but since codestral wants a different endpoint, it started getting a little confusing

[zed-industries-bot]

	Warnings
⚠️	This PR is missing release notes. Please add a "Release Notes" section that describes the change: Release Notes: - Added/Fixed/Improved ... If your change is not user-facing, you can use "N/A" for the entry: Release Notes: - N/A

Generated by 🚫 dangerJS against 16459c1

[dvcrn]

@mgsloan thanks for taking this PR until merge!

[tredondo]

1. Documentation is missing at https://zed.dev/docs/ai/edit-prediction.
2. How about using Codestral locally, e.g. through Ollama?
3. Does this continue the work from Codestral (Mistral code suggestion) #12519 (comment)? (that issue was closed brutally)

[pedrxd]

Is there any reason why the api.mistral.ia endpoint is not used? Currently, the codestral endpoint only has the free option, and due to their privacy policy, they use your data to improve their models.

I understand that for most people this may be a benefit, but considering that they may handle sensitive data, it would be advisable to use the paid option. Apart from not having to configure several tokens, as one is sufficient.

Thank you very much for your work on this pull request.

[dvcrn]

@pedrxd Is Codestral now available through normal mistral API? When I started implementation on this, it sounded like the separate Codestral endpoint was the only real way to use Codestral FIM, hence why I couldn't just re-use the existing mistral provider in Zed. Let me look into this

/EDIT: Okay we can swap to api.mistral.ai - https://docs.mistral.ai/capabilities/code_generation

But the docs still say:

• If you're a user, wanting to query Codestral as part of an IDE plugin, codestral.mistral.ai is recommended.
• If you're building a plugin, or anything that exposes these endpoints directly to the user, and expect them to bring their own API keys, you should also target codestral.mistral.ai

So both endpoints make sense here. Maybe we could check if the user already has the mistral provider setup? But I don't think we should just start using it without any indication what's happening and ramp up cost on that key.

How about making the API endpoint configurable? The API key was already configurable in my initial PR, so it would just be another field. Users that want api.mistral.ai can then swap the endpoint + API key out to explicitly use api.mistral.ai over codestral.mistral.ai

[pedrxd]

@dvcrn
Exactly, we should not change the current behaviour. What's more, using the free Mistral option is perfectly acceptable; I even think that endpoint is used for the business plan.

Mistral allows on-premise deployments, so allowing you to write the endpoint you want to use could be the best option.

The configuration interface could be left as it is now, allowing endpoints to be overwritten via settings.json (like a advance option). This would allow someone with an on-premises deployment to change the endpoints, and someone who wants to use API.mistral.ai instead of codestral to overwrite them.

Either option would be fine. If you think it is more appropriate to put it in the interface, that is also fine.

[dvcrn]

The configuration interface could be left as it is now, allowing endpoints to be overwritten via settings.json (like a advance option). This would allow someone with an on-premises deployment to change the endpoints, and someone who wants to use API.mistral.ai instead of codestral to overwrite them.

Yes I think if anything, this option would be the best. We default to codestral and if someone really wants to, they can override the base_url through the settings.json file, in a similar way to how we can already do it for other providers. (Wonder if that would even enable local ollama codestral, if the API interface is identical)

[tredondo]

@dvcrn: maybe remove the settings.json block from your original comment while the only documentation about configuring Codestral is the link to this issue until #40949 is fixed?

@remorses:

I am using this branch as my daily driver, this is the best inline completion provider available in Zed right now, really thank you

It's still surprisingly unintelligent. In a JS file, it completes "// Generate a random string of 100 characters" with "Math.random().toString(36).substring(2, 100)".