Skip to content

feat: enable code completion support with CodeGemma

Bruno Cardoso requested to merge bc/add_code_gemma_2b into main

What does this merge request do and why?

Adds support for codegemma (either 2b or 7b-code) to the /v2/code/completions endpoint using the litellm client (introduced in !742 (merged)).

Part of Code Gemma 2B prompt: Code Completion (gitlab-org/gitlab#462483 - closed)

How to set up and validate locally

  1. Download a model and start your ollama server:
ollama pull codegemma:2b
ollama serve
  1. Start the litellm proxy server (to allow ai-gateway to use the openai-compatible api):
# config.yaml
model_list:
  - model_name: codegemma
    litellm_params:
      model: ollama/codegemma:2b
      api_base: http://localhost:11434
litellm --config config.yaml
  1. Run ai-gateway and go to http://localhost:5052/docs and make a request to the /v2/code/completions endpoint using this example payload (make sure you have AIGW_CUSTOM_MODELS__ENABLED=True in your ai-gateway .env file):
{
  "current_file": {
    "file_name": "app.py",
    "language_identifier": "python",
    "content_above_cursor": "def hello_world():\\n",
    "content_below_cursor": "    return "
  },
  "model_provider": "litellm",
  "model_endpoint": "http://0.0.0.0:4000",
  "model_name": "codegemma",
  "telemetry": [],
  "prompt_version": 2,
  "prompt": ""
}

Screen_Recording_2024-05-28_at_15.24.03

Merge request checklist

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

Merge request reports