Proposal: Code Suggestions Context

Code Suggestions Context

Overview

The Code Suggestions Context aims to enhance the code suggestions feature by providing more context. It involves gathering information from multiple files open in the code editor, storing it as context, and reusing it across several code suggestion requests.

Problem description

Currently, when the editor extension requests code suggestions, it only provides information about the current file, such as the filename and the context above and below the cursor in the file. The prompt generated for the AI model contains this information for more relevant code suggestions. However, there could be more relevant information outside of the current file that can improve the quality of suggested results and make them more relevant.

One way to provide more information about the current state is to extend the ``/code_suggestions/completions` API endpoint to receive content from neighboring or related files.

{
  "current_file": {
    "file_name": "main.go",
    "content_above_cursor": "...",
    "content_below_cursor": "..."
  },
  "neighbours": [
    { "file_name": "user.go", "content": "..." },
    { "file_name": "user_test.go", "content": "..." },
  ]
}

There is a disadvantage to this approach. The API endpoint will receive a larger payload since neighbours can contain several files with their content. More importantly, when a developer works in a file without changing context, the neighbours body will not change. An editor extension will send dozens or more requests with the same neighbours body, which can be a significant part of the whole request body.

Proposal

Introduce a new API endpoint to cache the current context and access it every time code suggestions are requested.

The context API endpoint receives information about names and content of relative files, stores information in cache, and responds with a cache key (context_id), and optionaly expiration time.

We introduce a separate API endpoint to provide a neighbours context and get a context ID back. We are going to feed to the API the list of neighbours:

{
  "neighbours": [
    { "file_name": "user.go", "content": "..." },
    { "file_name": "user_test.go", "content": "..." },
  ]
}

Redis can be used as a cache storage, since the context is relatively short-lived.

When an editor extension receives context_id it caches it locally, and pass as a new parameter to /code_suggestions/completions API endpoint.

{
  "current_file": {
    "file_name": "main.go",
    "content_above_cursor": "...",
    "content_below_cursor": "..."
  },
  "contextId": "some-unique-context-id"
}

When GitLab monolith receives code suggestion request, it tries to retrive additional information from the context and enchance code suggestions prompt.

<content_above_cursor>
<content_below_cursor>

<probably_related_files>
 <file filename="user.go">
   ...
 </file>
 <file filename="user_test.go">
   ...
 </file>
</probably_related_files>

With this approach, the request body sent by the editor extension will remain around the same size, but the payload between the monolith and AI Gateway (and therefore between AI Gateway and Anthropic) will grow in size, potentially affecting latency.

When editor extension updates context

An editor extension can update the context when events occur, such as

• On editor launch
• Current file changed (user opened different tab or a new file)
• Context expired
• Etc.

How editor extension gathers information

There could be different ways to detect neighbors, starting from the simplest way to more advanced:

• Currently open tabs in the code editor
• Related files, such as production code and their test files
• Related files from a vector database

Requests diagram

sequenceDiagram
    actor USER as User
    participant IDE as Editor Extension
    participant GITLAB as GitLab Monolith
    participant RD as Redis
    participant AIGW as AI Gateway
    participant ANT as Anthropic
    USER->>IDE: Opens `user.go` file
    USER->>IDE: Opens `user_test.go` file
    IDE->>GITLAB: Sends new context request
    GITLAB->>RD: Stores context
    RD->>GITLAB: Generates ctx_id
    GITLAB->>IDE: Responds with ctx_id
    USER->>IDE: Opens `main.go` file
    IDE->>GITLAB: Sends new context request
    GITLAB->>RD: Stores context
    RD->>GITLAB: Generates ctx_id
    GITLAB->>IDE: Responds with `ctx_id`
    USER->>IDE: Types `// print user's fullname`
    IDE->>GITLAB: Requests code suggestions, passes `ctx_id`
    GITLAB->>RD: Fetches context info
    RD->>GITLAB: Provides context info
    GITLAB->>AIGW: Sends request with enchansed prompt
    AIGW->>ANT: Send requests
    ANT->>AIGW: Response
    AIGW->>GITLAB: Response
    GITLAB->>IDE: Response
    IDE->>USER: Displays a suggestion

Additional thoughts

This could be an intermediate solution or a solution when there is no vertex DB related to a project (for example, not scanned by Repository X-Ray yet). The solution can work with projects not hosted on GitLab. It requires a GitLab token to perform API requests, which is already done as an Authorization step from GitLab editor extension.

Redis memory usage estimation

More details in this internal comment #440646 (comment 1827360208)

PoC Merge Request

POC: Code Suggestions Context (!144879 - closed)

PoC Demo

https://youtu.be/MFke56qvFes