Skip to content

Mock `TextGenBaseModel` and `ChatModelBase` to recover the fake model logic

Alexander Chueshev requested to merge ac/fix-fake-models into main

What does this merge request do and why?

This MR recovers the existing fake model logic by mocking the Anthropic and Vertex clients instead of using FakePalmTextGenModel. More details about why we need to recover this logic - https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/415

How to set up and validate locally

  1. Run the AIGW setting the env variable AIGW_MOCK_MODEL_RESPONSES=true
  2. Request any model
  3. The AIGW will echo a given prompt
    {
     "id": "id",
     "model": {
       "engine": "anthropic",
       "name": "claude-instant-1.2",
       "lang": ""
     },
     "experiments": [],
     "object": "text_completion",
     "created": 1711545755,
     "choices": [
       {
         "text": "echo: [\"hello\"]",
         "index": 0,
         "finish_reason": "length"
       }
     ]
    }

Merge request checklist

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

Closes https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/415

Edited by Alexander Chueshev

Merge request reports