Resolve "Integrate Model Selection with Code Completion"

What does this MR do and why?

This MR integrates brings model switching capability to Code Completions. Model switching for code generation was completed already in !193266 (merged), and this change is also along similar lines, but with minor differences:

Decisions:

  • Direct connections won't be allowed when a model is "pinned" (ie, when a model is explicitly chosen for code completions). Reason in #548241 (comment 2555202872). MR in Disables direct connection model is pinned (!193629 - merged)
  • We to send the prompt from the monolith for code completion, even when model switching is used. This is because the legacy anthropic handler for code completion can't obtain the prompt from prompt registry, so we need to pass it in from the monolith. However, this prompt will be put to use by AIGW only for anthropic models, and not for Codestral on fireworks/vertex.
  • When no model is selected, and the customer chooses to keep "GitLab Default" as the model for code completions, the "default" model for code will continue to be decided by the monolith (using the saas_prompt method). This is in contrast to all other AI features where the "default" for the feature is decided by the AI Gateway, as mentioned in unit_primitives.yml with default_model

Testing

  • Execute gdk update so that corresponding changes in AIGW land on your local - gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!2741 (merged)
  • Install VSCode and install the GitLab workflow extension.
  • Login to the GitLab workflow extension with your GDK credentials (ie, PAT of root user)
  • Enable model switching globally: Feature.enable(:ai_model_switching)
  • In a root group where you want to use model switching, set Codestral 25.01 as the model for Code completion via http://gdk.test:3000/groups/group-name/-/settings/gitlab_duo/model_selection

Screenshot_2025-06-11_at_10.03.59_AM

  • Clone any project within this group locally, and open it in VSCode.
  • Make sure that VSCode correctly associates the opened folder to a GitLab project. If it does, you should be a able to see the following details on the GitLab Workflow tab like Issues assigned to me etc.

Screenshot_2025-06-11_at_10.04.44_AM

  • in any existing file or new file, write some code so that code completion requests are fired.
  • If you tail ai gateway logs with gdk tail gitlab-ai-gateway, you should be able to see that the completion request was powered by Codestral on fireworks and that it works.

Screenshot_2025-06-11_at_10.08.47_AM

  • Similarly, if we switch the model for code completion and execute a completion request, we can see in the logs that the request is indeed powered by the model we selected the second time.

Vertex:

Screenshot_2025-06-11_at_10.10.48_AM

Sonnet:

Screenshot_2025-06-11_at_10.11.32_AM

References

Screenshots or screen recordings

Before After

How to set up and validate locally

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Related to #545065 (closed)

Edited by Manoj M J [OOO]

Merge request reports

Loading