Skip to content

fix: add post-processing for codestral-on-vertex

Pam Artiaga requested to merge 634-add-post-processing-for-codestral into main

What does this merge request do and why?

We recently introduced Codestral (on Vertex AI) as a replacement for VertexAI-CodeGecko for Code Completions (see #564 (closed)).

Code Gecko, which uses the ai_gateway.code_suggestions.completions.CodeCompletionsLegacy class to generate suggestions, has post-processing for the suggestions. Codestral uses the ai_gateway.code_suggestions.completions.CodeCompletions class, which has no post-processing.

This MR introduces post-processor support to the ai_gateway.code_suggestions.completions.CodeCompletions class, then makes sure that post-processing is used for Codestral-on-Vertex.

This closes both Codestral post-processing (#634 - closed) and Investigate and fix Codestral bracket closing b... (#633 - closed)

How to set up and validate locally

Run this command to request a code suggestion for a Python function called def get_app_name_and_version():

curl "http://gdk.test:5052/v2/code/completions" \
-X POST \
--header "Content-Type: application/json" \
--data '{
    "current_file": {
      "file_name":"main.py",
      "language_identifier":"python",
      "content_above_cursor":"def get_app_name_and_version():\nreturn{",
      "content_below_cursor":"}"
    },
    "stream":false,
    "prompt_version":1,
    "model_provider":"vertex-ai",
    "model_name":"codestral@2405"
  }' \
| json_pp -json_opt pretty,canonical

Before the change / on main branch:

You will get a response where the suggestion (in the choices field) has a closing bracket:

{
   "choices" : [
      {
         "finish_reason" : "length",
         "index" : 0,
         "text" : "\n\"app_name\": \"MyApp\",\n\"version\": \"1.0.0\"\n}"
      }
   ],
   ...other response fields here...
}

After the change / on 634-add-post-processing-for-codestral branch:

The response should NOT have a closing bracket:

{
   "choices" : [
      {
         "finish_reason" : "length",
         "index" : 0,
         "text" : "\n\"app_name\": \"MyApp\",\n\"version\": \"1.0.0\"\n"
      }
   ],
   ...other response fields here...
}

Merge request checklist

  • Tests added for new functionality. If not, please raise an issue to follow up.
  • Documentation added/updated, if needed.

Closes #634 (closed)

Edited by Pam Artiaga

Merge request reports

Loading