Skip to content

Extend VertexAi::Client to add text/code/message_chat methods

George Koltsov requested to merge georgekoltsov/vertex-text-bison into master

What does this MR do and why?

This MR updates Gitlab::Llm::VertexAi::Client to:

  • Add text method to utilize text-bison model
  • Add code method to utilize code-bison-001 model
  • Add messages_chat method to utilize multiturn chat-bison-001 model

The newest models endpoints live under different URLs comparing to the default codechat endpoint that makes a request to Gitlab::CurrentSettings.current_application_settings.tofa_url. Since the endpoint URL is currently hardcoded in the application settings, I am replacing parts of the string during runtime to make each corresponding url valid. For example, for :text endpoint:

  1. Removing -preprod- from hostname
  2. Replacing endpoints with publishers/google/models
  3. Model name from codechat-bison-001 to text-bison

Each model now has slightly different URL endpoint updates and has a corresponding ModelConfigurations::* class.

Mentions https://gitlab.com/gitlab-org/gitlab/-/issues/407336 & https://gitlab.com/gitlab-org/gitlab/-/issues/410112

Screenshots or screen recordings

Screenshots are required for UI changes, and strongly recommended for all other merge requests.

How to set up and validate locally

Assuming you have Vertex AI API setup locally:

  1. Open rails console
  2. > Gitlab::Llm::VertexAi::Client.new(User.first).chat(content: 'Hi!')
  3. > Gitlab::Llm::VertexAi::Client.new(User.first).messages_chat(content: [{author: 'user', content: 'Who is GitLab CEO?'},{author: 'content', 'content': 'Sytse Sijbrandij'}, {author:'user', content:'what is his middlename'}])
  4. > Gitlab::Llm::VertexAi::Client.new(User.first).code(content: 'function()')
  5. > Gitlab::Llm::VertexAi::Client.new(User.first).text(content: 'Hi!')
  6. Verify responses are coming from the corresponding models

MR acceptance checklist

This checklist encourages us to confirm any changes have been analyzed to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Vasilii Iakliushin

Merge request reports