feat: add Completions agent for codestral on vertex

What does this merge request do and why?

We are introducing Codestral as the new model for Code Completions, and we decided to use Vertex as the provider for Codestral. (See https://gitlab.com/gitlab-org/gitlab/-/issues/470171+).

To add support for Codestral-on-Vertex on AI Gateway, we:

• accept a new provider and model combination for the /code/completions endpoint: vertex-ai and codestral@2405
• use litellm SDK since this supports Codestral-on-Vertex
• implement this through the Agent registry so that we can specify the prompt_template on AI Gateway.
  • Since we are using Agents, we do not use litellm sdk directly, but through langchain_community.chat_models.ChatLiteLLM. See !1172 (comment 2028386691)

How to set up and validate locally

Validate successful request

expand for details

Send a Code Completions request to the AI Gateway with "vertex-ai" as the provider and "codestral@2405" as the model name. The payload should look like:

{
  "current_file": {
    "content_below_cursor": "}",
    "file_name": "main.go",
    "language_identifier": "go",
    "content_above_cursor": "func parseEnvFile(filename string) (map[string]string, error) {"
  },
  "stream": false,
  "prompt_version": 2,
  "model_provider": "vertex-ai",
  "model_name": "codestral"
}

And the whole request would look like:

curl "<your-local-ai-gateway-host>/v2/code/completions" \
-X POST \
--header "Content-Type: application/json" \
--data "{\"current_file\":{\"content_below_cursor\":\"<the content below cursor>\",\"file_name\":\"<the file name>\",\"language_identifier\":\"<the language identifier>\",\"content_above_cursor\":\"<the content above cursor>\"},\"stream\":false,\"prompt_version\":2,\"model_provider\":\"vertex-ai\",\"model_name\":\"codestral@2405\"}" \
| json_pp -json_opt pretty,canonical # this is just to have the response in a human-readable format

The response should look something like:

{
   "choices" : [
      {
         "finish_reason" : "length",
         "index" : 0,
         "text" : "\n\n\t// Read the contents of the .env file\n\tdata, err := os.ReadFile(filename)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Split the contents into lines\n\tlines := strings.Split(string(data), \"\\n\")\n\n\t// Initialize an empty map to store the key-value pairs\n\tenvMap := make(map[string]string)\n\n\t// Iterate over each line\n\tfor _, line := range lines {\n\t\t// Split the line into key and value\n\t\t"
      }
   ],
   "created" : 1724063761,
   "experiments" : [],
   "id" : "id",
   "model" : {
      "engine" : "vertex-ai",
      "lang" : "go",
      "name" : "vertex_ai/codestral@2405",
      "tokens_consumption_metadata" : {
         "context_tokens_sent" : 0,
         "context_tokens_used" : 0,
         "input_tokens" : 20,
         "output_tokens" : 0
      }
   },
   "object" : "text_completion"
}

Validate failed request

expand for details

AI Gateway handles the prompt_template for Vertex-on-Codestral (defined in the Agent definition file). If you set a prompt (essentially the prompt_template) in the payload, this should return an error:

curl "http://gdk.test:5052/v2/code/completions" \
--include \ # we are setting this so we can see the status code in the response
-X POST \
--header "Content-Type: application/json" \
--data "{\"current_file\":{\"content_below_cursor\":\"}\",\"file_name\":\"main.go\",\"language_identifier\":\"go\",\"content_above_cursor\":\"func parseEnvFile(filename string) (map[string]string, error) {\"},\"stream\":false,\"prompt_version\":2,\"model_provider\":\"vertex-ai\",\"model_name\":\"codestral@2405\",\"prompt\":\"blah\"}"

# response
HTTP/1.1 400 Bad Request
date: Wed, 07 Aug 2024 10:34:14 GMT
server: uvicorn
content-length: 86
content-type: application/json
x-process-time: 0.0034730409970507026
x-request-id: d2289a9dc47e44ea8fd3e13a659daa50

{"detail":"You cannot specify a prompt with the given provider and model combination"}%

Merge request checklist

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