Как работает Code Completion?

from models import User
from telemetry import telemetry_for_user_event

def create_user(data: dict) -> User:
    user = User.model_validate(**data)
    <CURSOR>
    return user

Там, где <CURSOR> находится каретка пользователя. Все что до — PREFIX, а все, что после — SUFFIX. Конечно, это самая важная инфа, ведь модель должна в своей подсказке продолжить PREFIX, причем так, чтобы после логично шел SUFFIX. В импорте есть функция telemetry_for_user_event, ее LLM и должна использовать в подсказке.

Отправим всю эту информацию в модель. Для примера рассмотрим LLM, которая специально обучена дописывать код. Например, starcoder2. Промпт для нашей задачи будет выглядеть вот так

<fim_prefix>
from models import User
from telemetry import telemetry_for_user_event

def create_user(data: dict) -> User:
    user = User.model_validate(**data)
    
<fix_suffix>
    return user
<file_sep>
<|endoftext|>

где <fim_prefix>, <fix_suffix> и т.д. — это специальные токены, на которых обучалась модель (подробнее можно почитать в самом пэйпере про starcoder2). Ответ будет следующим

telemetry_for_user_event(user)
<|endoftext|>

Выглядит хорошо, но есть одно но — посмотрим на определение функции telemetry_for_user_event

def telemetry_for_user_event(user: User, event: str):
    ...

Она принимает на вход два параметра, но LLM ничего об этом не знает. Поэтому нам нужно добавить в запрос CODE SNIPPETS, которые содержат вспомогательную информацию из других файлов проекта. Их можно собирать, например, добавив определения всех функций или классов, которые импортируются в наш целевой файл. Теперь наша схема имеет следующий вид:

PREFIX+SUFFIX + CODE SNIPPETS -> Prompt builder -> LLM -> Response

Отправим новый запрос в модель

<file_sep>telemetry.py
def telemetry_for_user_event(user: User, event: str):
    ...
<file_sep>main.py
<fim_prefix>
from models import User
from telemetry import telemetry_for_user_event

def create_user(data: dict) -> User:
    user = User.model_validate(**data)
    
<fix_suffix>
    return user
<file_sep>
<|endoftext|>

И получим следующий ответ

telemetry_for_user_event(user, event="user_created")
<|endoftext|>

Выглядит как то, что нам нужно. Но и в таком варианте есть «Но». Вот как должен выглядит тот код, который нужен пользователю

promo_available = check_user_promo(user)
    if promo_available:
        user.activate_promo()
    telemetry_for_user_event(user, event="user_created")

А почему? Потому что так принято в нашем игрушечном проекте. И мы не сможем просто по импортам нагрепать такую информацию. Но мы можем использовать USER BEHAVIOR — информацию о поведении нашего пользователя в IDE. Вот статья про использование USER BEHAVIOR в компании ByteDance (tiktok).

Например, мы можем использовать историю последних просмотренных файлов — наш пользователь перед тем, как начать реализацию create_user, посмотрел в проекте как реализуются другие похожие хэндлеры. Там может быть следующий код

from telemetry import telemetry_for_user_event
from models import User
from promotions import check_user_promo

def handle_auth(user: User):
    promo_available = check_user_promo(user)
    if promo_available:
        user.activate_promo()
    telemetry_for_user_event(user, event="auth_success")

Теперь наша схема работы Code Completion имеет вид

PREFIX+SUFFIX + CODE SNIPPETS + USER BEHAVIOR -> Prompt builder -> LLM -> Response

Новый запрос к модели

<file_sep>auth_service.py
…
def handle_auth(user: User):
…
<file_sep>telemetry.py
…
<file_sep>main.py
<fim_prefix>
...
<fix_suffix>
…
<file_sep>
<|endoftext|>

Ответ будет следующим

promo_available = check_user_promo(user)
    if promo_available:
        user.activate_promo()
    telemetry_for_user_event(user, event="user_created")
<|endoftext|>

И это уже то, что нам нужно. В нашей схеме осталось еще одно «Но».

Мы смотрим на код, в котором не более 10 строк. В реальных проектах обычно не так — тысячи или десятки тысяч строк кода, иногда даже в одном файле. Так же пользователь может посещать десятки файлов и производит множество действий в проекте. Это порождает десятки или даже сотни CODE SNIPPETS и USER BEHAVIOR. Чтобы не забивать весь контекст LLM этой информацией — необходимо ранжировать все сниппеты и выбирать только наиболее релевантные (для текущих PREFIX и SUFFIX). Вот финальная схема

PREFIX+SUFFIX + Ranking(CODE SNIPPETS + USER BEHAVIOR) -> Prompt builder -> LLM -> Response

Примерно так и работает Code Completion. Общая схема предельно понятная, но каждый отдельный пункт — содержательная инженерная задача.