Локальные ламы в VSCode

Некоторое время назад я написал плагин VSCode, используя уже устаревший Codex API от OpenAI, и он работал довольно хорошо, выполняя простые инструкции для выбранного кода в VSCode. Например, он может добавить строку документации к выбранной мной функции.

Поскольку я много возился с локальными LLM (большими языковыми моделями), следующим шагом, естественно, было посмотреть, как я могу создать что-то подобное, не зависящее от OpenAI.

Для этого я запустил проект оазис. Когда я только начинал, идея была очень простой: просто напрямую вызывать text-generation-webui API, чтобы сгенерировать строки документации для моего собственного кода.

Мне не потребовалось много времени, чтобы добиться аналогичных результатов с моим старым плагином, но теперь с использованием WizardLM вместо OpenAI. Это было довольно интересно, однако эта модель все еще ограничена в плане лицензии: она предназначена только для исследований.

И на этот раз я хотел что-то, что я мог бы использовать в коммерческих целях или для своей собственной работы, не имея проблем с объяснением лицензии модели.

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

Поскольку в последнее время я тратил больше времени на написание подсказок и ознакомился с библиотекой руководств Microsoft, мне пришло в голову, что я могу создать жизнеспособный продукт, объединив эти две вещи:

  1. Salesforce Codegen для создания кода
  2. Библиотека рекомендаций Microsoft, чтобы направить модель на создание того, что мне действительно нужно, вместо случайного ввода

Поэтому я потратил несколько дней на создание новой версии плагина, сочетающей в себе этот подход, и был вполне доволен первоначальными результатами. По общему признанию, решение оказалось намного сложнее, чем я ожидал, но оно работает довольно хорошо! И что более важно, мы можем использовать даже версию модели с 350 миллионами параметров и при этом корректно генерировать некоторые строки документации!

Конечно, это не всегда работает, и более крупные модели по-прежнему работают намного лучше, но использование меньших моделей дает нам огромное улучшение с точки зрения доступности и производительности логического вывода — почти любой может раскрутить версию модели 350 м.

Текущая версия этого плагина доступна на моем GitHub.

Сегодня мы рассмотрим, как был создан этот плагин:

  1. Написание расширения VSCode.
  2. Настройка HTTP-сервера с библиотекой инструкций.
  3. Настройка сервера подсказок промежуточного программного обеспечения для динамического анализа кода и выбора подсказок.
  4. Подробности о том, как сделать так, чтобы такая команда, как add_docstring, эффективно работала с менее мощными моделями. Это будет включать в себя, как написать подсказку и убедиться, что ввод / вывод выглядит так, как ожидалось. Этот последний раздел будет далее подразделяться на:

а) Разбор ввода

b) Подсказка «Добавить строку документации»

c) Построение вывода

Примечание. формат строки документации довольно прост. Ничто не мешает нам улучшать сгенерированный формат.

Вот схема высокого уровня того, что мы строим:

Основы написания расширения VSCode

Первый шаг — следовать официальной документации hello world. Это дает вам хороший образец расширения с начальной загрузкой.

Затем мы можем изменить наш extension.ts, чтобы добавить основную логику. Первый шаг — определить, как мы будем инициализировать наше расширение. После инициализации мы зарегистрируем существующие команды. Вот соответствующий код:

export function activate(context: vscode.ExtensionContext) {
 const commands = [
  ["addDocstring", "add_docstring"],
  ["addTypeHints", "add_type_hints"],
  ["addUnitTest", "add_unit_test"],
  ["fixSyntaxError", "fix_syntax_error"],
  ["customPrompt", "custom_prompt"],
 ];
 commands.forEach(tuple_ => {
  const [commandName, oasisCommand] = tuple_;
  const command = vscode.commands.registerCommand(`oasis.${commandName}`, () => {
   useoasis(oasisCommand);
  });
  context.subscriptions.push(command);
 });
}

Это связывает каждую ожидаемую команду с вызовом основной логической функции, называемой useoasis. Мы рассмотрим эту функцию позже.

Эти команды должны соответствовать тому, что определено в проекте package.json — мы определим несколько блоков для объявления команд, предоставляемых нашим расширением — начнем с определения contributes.commands:

[
      {
        "command": "oasis.addDocstring",
        "title": "Add Docstring to Selection"
      },
      {
        "command": "oasis.addTypeHints",
        "title": "Add type hints to selection"
      },
      {
        "command": "oasis.fixSyntaxError",
        "title": "Fix syntax error"
      },
      {
        "command": "oasis.customPrompt",
        "title": "Custom Prompt"
      }
]

Мы также хотим, чтобы пользователь мог выбрать блок кода с помощью мыши, найти эту команду с помощью щелчка правой кнопкой мыши, поэтому мы используем конфигурацию contributes.menus.editor/context, например:

"editor/context": [
  {
    "command": "oasis.addDocstring",
    "when": "editorHasSelection",
    "group": "7_modification"
  },
  {
    "command": "oasis.addTypeHints",
    "when": "editorHasSelection",
    "group": "7_modification"
  },
  {
    "command": "oasis.fixSyntaxError",
    "when": "editorHasSelection",
    "group": "7_modification"
  },
  {
    "command": "oasis.customPrompt",
    "when": "editorHasSelection",
    "group": "7_modification"
  }
]

Хорошо, этого достаточно, чтобы начать работу с командами привязки. Вернемся назад и посмотрим на нашу основную функцию расширения:

async function useoasis(command: string) {
 const activeEditor = vscode.window.activeTextEditor;
 if (!activeEditor) {
  return;
 }
 console.log("Reading config")
 const oasisUrl = vscode.workspace.getConfiguration('oasis').get("prompt_server_url")
 console.log("Oasis URL:", oasisUrl);
 const document = activeEditor.document;
 const selection = activeEditor.selection;

 const text = document.getText(selection);

 const requestBody = JSON.stringify({
  data: text
 });
 const url = `${oasisUrl}/${command}`;

 console.log("Calling API", url, "with body: ", requestBody);

 let response: oasisResponse | undefined = undefined;
 try {
  response = await got(url, {
   method: "POST",
   headers: {
    // eslint-disable-next-line @typescript-eslint/naming-convention
    "Content-Type": "application/json",
    // eslint-disable-next-line @typescript-eslint/naming-convention
   },
   body: requestBody,
   timeout: {
    request: 300000  // 5 minutes max
   }
  }).json();
 } catch (e: any) {
  vscode.window.showErrorMessage("Oasis Plugin: error calling the API")
  try {
   const apiStatusCode = `Error calling API: ${e.response.statusCode}`;
   vscode.window.showErrorMessage(apiStatusCode);
  } catch (error) {
   console.error("Error parsing error response code", error);
  }
 }

 if (response) {
  console.log("From got", response.text);
  const editedText = response.text;

  activeEditor.edit(editBuilder => {
   console.log("Edit builder", editBuilder);
   editBuilder.replace(selection, editedText);
  });
 }
};

Это немного дольше, но все еще довольно просто. Во-первых, мы проверяем наличие активного редактора, чтобы избежать ссылки на null. Вот простой трюк (слегка переформатированный):

const oasisUrl = vscode.workspace
  .getConfiguration('oasis')
  .get("prompt_server_url")

Это считывает конфигурацию рабочей области, чтобы узнать, где обслуживается сервер подсказок. Итак, мы должны вернуться к конфигурации расширения (package.json) и добавить следующее:

  "contributes": {
    "configuration": {
      "title": "Oasis",
      "properties": {
        "oasis.prompt_server_url": {
          "type": "string",
          "default": "http://0.0.0.0:9000",
          "description": "The URL where the extension can find the prompt server. Defaults to: 'http://0.0.0.0:9000'"
        }
      }
    },

Это не только определяет значение по умолчанию, но также предоставляет конфигурацию, чтобы пользователь мог переопределить ее либо через пользовательский интерфейс, либо через редактор JSON. Вот пример скриншота пользовательского интерфейса:

Вернемся к коду расширения:

const text = document.getText(selection);

const requestBody = JSON.stringify({
data: text
});
const url = `${oasisUrl}/${command}`;

console.log("Calling API", url, "with body: ", requestBody);

Довольно просто, читаем выделенный текст в редакторе и готовим запрос. Напомним, что command является параметром замыкания, встроенным в функцию activate.

Мы используем здесь библиотеку got (без уважительной причины) и делаем запрос:

 let response: oasisResponse | undefined = undefined;
 try {
  response = await got(url, {
   method: "POST",
   headers: {
    // eslint-disable-next-line @typescript-eslint/naming-convention
    "Content-Type": "application/json",
    // eslint-disable-next-line @typescript-eslint/naming-convention
   },
   body: requestBody,
   timeout: {
    request: 300000  // 5 minutes max
   }
  }).json();
 } catch (e: any) {
  vscode.window.showErrorMessage("Oasis Plugin: error calling the API")
  try {
   const apiStatusCode = `Error calling API: ${e.response.statusCode}`;
   vscode.window.showErrorMessage(apiStatusCode);
  } catch (error) {
   console.error("Error parsing error response code", error);
  }
 }

Если запрос не выполнен, мы отображаем ошибку в пользовательском интерфейсе, которая отображается следующим образом:

Не очень информативно, но, по крайней мере, пользователь знает, что ждать больше не стоит. Затем, если все работает, мы парсим и затем используем активный редактор, чтобы заменить выделение ответом от API:

 if (response) {
  console.log("From got", response.text);
  const editedText = response.text;

  activeEditor.edit(editBuilder => {
   console.log("Edit builder", editBuilder);
   editBuilder.replace(selection, editedText);
  });
 }

Библиотека руководств: Настройка HTTP-сервера с помощью библиотеки руководств

Guidance — очень интересная библиотека от Microsoft для работы с LLM. Я рекомендую быстро прочитать его Readme.

Основное преимущество использования инструкций заключается в том, что вы можете выполнять определенные задачи с гораздо менее мощными моделями, следуя инструкциям.

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

При разработке под руководством может быть очень сложно перезагрузить сервер, особенно если используемая модель тяжелая. По этой причине я обернул загрузку модели и библиотеку инструкций внутри сервера, который не нужно часто менять.

Исходный код этого сервера очень прост и помещается в один файл:

class Request(BaseModel):
    input_vars: Dict[str, str]
    output_vars: List[str]
    guidance_kwargs: Dict[str, str]
    prompt_template: str


app = FastAPI()

print("Loading model, this may take a while...")
# model = "TheBloke/wizardLM-7B-HF"
model = "Salesforce/codegen-350m-mono"
# model = "Salesforce/codegen-2b-mono"
# model = "Salesforce/codegen-6b-mono"
# model = "Salesforce/codegen-16B-mono"

llama = guidance.llms.Transformers(model, quantization_config=nf4_config, revision="main")
print("Server loaded!")


@app.post("/")
def call_llama(request: Request):
    input_vars = request.input_vars
    kwargs = request.guidance_kwargs
    output_vars = request.output_vars

    guidance_program: Program = guidance(request.prompt_template)
    program_result = guidance_program(
        **kwargs,
        stream=False,
        async_mode=False,
        caching=False,
        **input_vars,
        llm=llama,
    )
    output = {}
    for output_var in output_vars:
        output[output_var] = program_result[output_var]
    return output

Вот и все! Он берет шаблон приглашения с входными переменными и ожидаемыми выходными переменными из HTTP-запроса и направляет его через библиотеку инструкций. Затем он извлекает ожидаемые выходные переменные и возвращает их в ответе HTTP.

Если вам интересно, quantization_config, Я использую новейшую технику 4-битного квантования, выпущенную Hugging Face.

Нам также понадобится клиент для использования этого сервера:

import requests

guidance_url = "http://0.0.0.0:9090"
def call_guidance(prompt_template, output_vars, input_vars=None, guidance_kwargs=None):
    if input_vars is None:
        input_vars = {}
    if guidance_kwargs is None:
        guidance_kwargs = {}
    
    data = {
        "prompt_template": prompt_template,
        "output_vars": output_vars,
        "guidance_kwargs": guidance_kwargs,
        "input_vars": input_vars,
    }
    
    response = requests.post(
        guidance_url, 
        json=data
    )
    
    response.raise_for_status()
    
    return response.json()

Ничего слишком интересного, чтобы увидеть здесь.

Обратите внимание, что этот сервер не взаимодействует напрямую с расширением VScode.

Настройка сервера подсказок промежуточного программного обеспечения для динамического анализа кода и выбора подсказок

Теперь давайте посмотрим на промежуточное ПО, которое фактически получает запрос от VSCode. Давайте перейдем к коду основного сервера:

from fastapi import FastAPI, HTTPException
import logging
import codegen_guidance_prompts
import wizard_lm_guidance_prompts

from pydantic import BaseModel
from guidance_client import call_guidance
from commands.commands import build_command_mapping


logger = logging.getLogger("uvicorn")
logger.setLevel(logging.DEBUG)

class Request(BaseModel):
    data: str


app = FastAPI()


prompts_module = codegen_guidance_prompts
commands_mapping = build_command_mapping(prompts_module)


@app.post("/{command}/")
def read_root(command, request: Request):
    logger.info("Received command: '%s'", command)
    logger.debug("Received data: '%s'", request)
    received_code = request.data
    try:
        command_to_apply = commands_mapping[command]
        logger.info("Loaded command: '%s'", command_to_apply)

    except KeyError:
        raise HTTPException(status_code=404, detail=f"Command not supported: {command}")

    prompt_key, prompt_to_apply, extracted_input = command_to_apply.prompt_picker(received_code)
    logger.info("Extracted input: %s", extracted_input)

    keys_difference = set(prompt_to_apply.input_vars) - set(extracted_input.keys())

    if keys_difference:
        error_msg = f"Missing input keys for the prompt: {keys_difference}"
        logger.error(error_msg)
        raise HTTPException(status_code=500, detail=error_msg)


    logger.info("Loaded command: '%s'", command_to_apply.__class__.__name__)
    logger.info("Parsed input :")
    for key, item in extracted_input.items():
        logger.info("(%s): '%s'", key, item)
    logger.info("Calling LLM...")

    result = call_guidance(
        prompt_template=prompt_to_apply.prompt_template,
        input_vars=extracted_input,
        output_vars=prompt_to_apply.output_vars,
        guidance_kwargs={}
    )
    logger.info("LLM output: '%s'", result)
    
    result = command_to_apply.output_extractor(prompt_key, extracted_input, result)

    logger.info("parsed output: '%s'", result)

    return {"text": result}

Это может выглядеть просто, но скрывает много неприятных деталей. Однако общая идея контроллера действительно проста:

  1. Взять входную пару (команда, код ввода)
  2. Найдите подходящую подсказку
  3. Разобрать ввод в определенный формат
  4. Запустить через сервер наведения
  5. Разобрать вывод в правильный формат
  6. Ответить обратно

Детали команды add_docstring

Разбор ввода

Хорошо, мы прошли через простые части. Теперь давайте рассмотрим более сложные части. Давайте расширим каждый шаг выше.

Во-первых, мы узнаем, какую команду нас просят применить:

logger.info("Received command: '%s'", command)
logger.debug("Received data: '%s'", request)
received_code = request.data
try:
    command_to_apply = commands_mapping[command]
    logger.info("Loaded command: '%s'", command_to_apply)

Это словарь, созданный с помощью этого вызова:

commands_mapping = build_command_mapping(prompts_module)

Здесь мы определяем доступные команды. Обратите внимание, что пока доступна только одна команда:

def build_command_mapping(prompt_module: PromptModuleInterface):
    add_docstring_command = DocStringCommand(
        prompt={
            "generic_prompt": prompt_module.doc_string_guidance_prompt,
            "function_prompt": prompt_module.function_doc_string_guidance_prompt
        }
    )
    commands_mapping = {
          "add_docstring": add_docstring_command,
          # "add_type_hints": ADD_TYPE_HINTS,
          # "fix_syntax_error": FIX_SYNTAX_ERROR,
          # "custom_prompt": CUSTOM_PROMPT
    }

  return commands_mapping

Класс classDocStringCommand реализует абстрактный класс Command:

@dataclass
class Command:
    prompt: Dict[str, GuidancePrompt]

    @abc.abstractclassmethod
    def prompt_picker(self, input_: str) -> Tuple[str, GuidancePrompt, Dict[str, str]]:
        raise NotImplementedError()

    @abc.abstractclassmethod
    def output_extractor(self, prompt_key: str, extracted_input: Dict[str, str], result: Dict[str, str]) -> str:
        raise NotImplementedError()


class DocStringCommand(Command):
    def __init__(self, *args, **kwargs):
        super().__init__(*args,**kwargs)

Таким образом, этот класс должен содержать словарь подсказок и реализовывать две функции, чтобы он мог выбрать правильную подсказку и извлечь результат. GuidancePrompt покажется вам знакомым, если вы вспомните, как устроен наш сервер навигации:

@dataclass
class GuidancePrompt:
    prompt_template: str
    input_vars: Dict[str, str]
    output_vars: Dict[str, str]
    guidance_kwargs: Dict[str, str]

Давайте посмотрим, как мы выбираем подсказку:

def prompt_picker(self, input_: str) -> Tuple[str, GuidancePrompt, Dict[str, str]]:
    prompt_key = "None"
    try:
        function_header, function_body, leading_indentation, indentation_type = function_parser(input_)
        prompt_key = "function_prompt"
        return_value = prompt_key, self.prompt[prompt_key], {
            "function_header": function_header,
            "function_body": function_body,
            "leading_indentation": leading_indentation,
            "indentation_type": indentation_type
        }

    except (FailedToParseFunctionException):
        logger.warn("Failed to identify specific type of code block, falling back to generic prompt")
        prompt_key = "generic_prompt"
        return_value = prompt_key, self.prompt[prompt_key], {"input": input_}

    logger.info("Chosen prompt: %s", prompt_key)
    return return_value

Итак, если вы заметили, здесь две подсказки: generic_prompt и function_prompt. Если мы идентифицируем функцию, мы применяем более позднюю, специализированную, иначе мы возвращаемся к универсальной.

Возвращаемое значение представляет собой кортеж формата (prompt_key, prompt, input_dict). Мы будем использовать эти значения позже при вызове LLM с помощью этой подсказки и при анализе выходных данных LLM.

Возвращаясь немного назад, если мы посмотрим, как мы определяем, является ли это функцией, мы вызываем метод function_parser. Давайте пройдемся по этой кроличьей норе:

def function_parser(input_code_str: str) -> Tuple[str, str, str]:
    leading_indentation = _get_leading_indentation(input_code_str)
    simple_indented_code = _remove_extra_indentation(input_code_str, leading_indentation)
    indentation_type = _get_indentation_type(input_code_str)

    parsed = ast.parse(simple_indented_code, filename="<string>")
    parsed
    try:
        first_node = parsed.body[0]
    except IndexError:
        raise FailedToParseFunctionException from IndexError
    
    if not isinstance(first_node, ast.FunctionDef):
        raise FailedToParseFunctionException(f"Parsed type is not a function: '{type(first_node)}'")
    
    function_body = _extract_function_body(first_node, leading_indentation, indentation_type)
    function_header = _extract_function_header(first_node)
    return function_header, function_body, leading_indentation, indentation_type

Итак, основная идея здесь — использовать ast.parse (документация), чтобы мы могли посмотреть на AST и найти тело/заголовок функции. Мы также делаем некоторые странные манипуляции с отступами, почему?

Хорошо, если вы попытаетесь передать следующую строку в ast.parse:

"""  
  def hello():
    print('hello!')
""" 

Вас встретит приятный IndentationError. Почему? Я не совсем уверен, но кажется, что этот синтаксический анализатор предполагает, что он всегда читает файл. Это означает, что эта строка интерпретируется как исходный код уровня модуля и, следовательно, имеет неправильный отступ.

Работа над этим привела меня к нескольким обходным путям с отступами, которые я не буду подробно описывать — вы всегда можете погрузиться в полный исходный код, если вам интересно.

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

Подсказка по строке документации

Прежде чем мы начнем, давайте воспользуемся одной из функций Oasis, чтобы увидеть, как работает это поколение. Учитывая эту функцию в качестве входных данных:

def _extract_function_header(fun_code: ast.FunctionDef) -> str:
    full = ast.unparse(fun_code)
    body = ast.unparse(fun_code.body)
    return full.split(body[0:USED_BODY_CHARS_TO_SPLIT])[0].strip()

Это строка документации, сгенерированная Oasis:

def _extract_function_header(fun_code: ast.FunctionDef) -> str:
    """This function extracts the function header from the given function definition

    Parameters: 
        fun_code (ast.FunctionDef): The function definition to extract the function header from.

    Returns: 
        str: The function header.

    """

    full = ast.unparse(fun_code)
    body = ast.unparse(fun_code.body)
    return full.split(body[0:USED_BODY_CHARS_TO_SPLIT])[0].strip()

Так как же он генерируется? Подсказка несколько длинная, поэтому давайте рассмотрим ее блоками. Во-первых, мы начинаем приводить примеры с использованием старой доброй техники подсказок.

function_doc_string_guidance_prompt = GuidancePrompt(
    prompt_template="""
def print_hello_world():
    print("hello_world")

def print_hello_world():
# Docstring below
    \"\"\"This functions print the string 'hello_world' to the standard output.\"\"\"

def sum_2(x, y):
    return x + y

def sum_2(x, y):
# Docstring below
    \"\"\"This functions receives two parameters and returns the sum.
    
    Parameters:
        int: x - first number to sum
        int: y - second number to sum

    Returns:
        int: sum of the two given integers
    \"\"\"
    return x + y

Обратите внимание, что мы не даем явно инструкции модели, в отличие от того, что мы делаем с моделями WizardLM. Здесь мы просто предоставляем шаблоны, чтобы модель могла автозаполняться аналогичным образом. После нескольких примеров мы определяем, как должна выполняться генерация:

"""
(few shot examples above...)

{{leading_indentation}}{{function_header}}
{{function_body}}

{{leading_indentation}}{{function_header}}
# Docstring below
{{leading_indentation}}\"\"\"{{gen 'description' temperature=0.1 max_tokens=128 stop='.'}}

Parameters: {{gen 'parameters' temperature=0.1 max_tokens=128 stop='Returns:'}}
Returns: {{gen 'returns' temperature=0.1 max_tokens=128 stop='\"\"\"'}}
""",
    guidance_kwargs={},
    input_vars=["function_header", "function_body", "leading_indentation"],
    output_vars=["description", "parameters", "returns"],
)

В самом конце мы определяем ожидаемые входные и выходные переменные. Это не является строго частью шаблона, но помогает нам при взаимодействии с этой подсказкой из высокоуровневого вызова в кодовой базе. Давайте посмотрим на конец реальной подсказки.

Первый блок добавляет входной исходный код, как в нескольких примерах:

"""
{{leading_indentation}}{{function_header}}
{{function_body}}
"""

Синтаксис может напомнить вам механизм шаблонов, такой как Jinja2, и, как вы могли ожидать, он действительно представляет переменные, которые будут заменены перед отправкой входных данных в модель большого языка.

Затем здесь мы вставляем только заголовок функции и комментарий # Docstring below, как и в нескольких примерах, чтобы вызвать автодополнение, как мы ожидаем от модели:

{{leading_indentation}}{{function_header}}
# Docstring below

Мы запускаем нашу первую генерацию с помощью этой команды (давайте пока проигнорируем отступ):

\"\"\"{{gen 'description' temperature=0.1 max_tokens=128 stop='.'}}

Итак, здесь мы запускаем строку документации, а затем просим LLM генерировать ее до тех пор, пока она не найдет символ «.». Давайте посмотрим генерацию образца. Обратите внимание, что в выводе нет тройных кавычек:

This function extracts the function header from the given function definition

Затем следующая сгенерированная переменная:

Parameters: {{gen 'parameters' temperature=0.1 max_tokens=128 stop='Returns:'}}

Который расширяется в:

fun_code (ast.FunctionDef): The function definition to extract the function header from.

И наконец:

{{gen 'returns' temperature=0.1 max_tokens=128 stop='\"\"\"'}}

Становится:

str: The function header.

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

Создание вывода

Мы подошли к концу нашего серверного контроллера промежуточного слоя:

result = call_guidance(
        prompt_template=prompt_to_apply.prompt_template,
        input_vars=extracted_input,
        output_vars=prompt_to_apply.output_vars,
        guidance_kwargs={}
    )
logger.info("LLM output: '%s'", result)
    
result = command_to_apply.output_extractor(prompt_key, extracted_input, result)

Мы используем приглашение выше с нашим клиентом руководства, чтобы создать словарь, который будет содержать поколения из предыдущего раздела. Он придет в следующем формате:

{
  "description": "Func description",
  "parameters": "a (int): a number to use in the sum",
  "returns": "(int): returns the sum"
}

Затем мы можем предоставить его выходному экстрактору DocStringCommand — это длинная функция, так что давайте рассмотрим только некоторые ее части.

Во-первых, мы решаем, как анализировать вывод в зависимости от используемой подсказки.

def output_extractor(self, prompt_key, extracted_input, result: Dict[str, str]) -> str:
    if prompt_key == "generic_prompt":
        return result["output"]
    elif prompt_key == "function_prompt":
        ind = extracted_input["leading_indentation"]

Затем, когда мы анализируем вывод function_prompt, есть некоторая логика для преобразования сгенерированных параметров в часть строки документации:

# (boring indentation workaround black magic above...)
try:
    parameters = result["parameters"].strip().split("\n")
except (KeyError, ValueError):
    parameters = []

# (handles the same for the 'returns' part...)
parameters_string = ""
if parameters:
    parameters_string = body_indentation + "Parameters: \n"
    for param in parameters:
        parameters_string += f"{body_indentation}{indentation_type}{param.lstrip().strip()}\n"

    parameters_string += "\n"

Позже мы используем заголовок функции и добавляем описание и строку параметров, которую мы построили ранее (это все еще в той же функции):

(...)
code_with_docstring = (header_indentation + extracted_input["function_header"] + "\n") 
code_with_docstring += (body_indentation + '"""')
code_with_docstring += (result["description"].lstrip().strip() + "\n\n")

logger.info("Header with description: %s", code_with_docstring)
if parameters_string:
    code_with_docstring += (parameters_string)
if returns_string:
    code_with_docstring += (returns_string)


code_with_docstring += (body_indentation + '"""\n\n')
    
code_with_docstring += extracted_input["function_body"]
logger.info("Generated code with docstring: %s", code_with_docstring)
return code_with_docstring

Обратите внимание, как мы «вручную» восстанавливаем строку документации с поколениями из LLM. Это неединственный способ, также можно использовать всю строку непосредственно из руководства и частично избежать этой гимнастики (нам все равно нужно определить начало строки документации и немного изменить подсказку).

Заключение и следующие шаги

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

Однако в некоторых случаях дополнительные усилия окупаются: повышенная производительность подсказок определенно того стоит, когда требуется точность.

Я был довольно разочарован тем, насколько сложно было использовать ast.parse для моего варианта использования, в частности, это привело меня к двум проблемам:

  1. Раздражающий IndentationError, который заставил меня написать много лишнего кода.
  2. Удаление повторяющихся новых строк из проанализированного кода.

Вторая проблема — текущий источник ошибок в Oasis, а это значит, что мне может понадобиться полностью удалить ast.parse и написать свой собственный синтаксический анализатор.

В целом, однако, я думаю, что мы получили хорошую отправную точку для написания управляемых подсказок, которые помогают с генерацией кода.