Zsh Codex 代码解析一

我码玄黄
320 阅读3分钟

Zsh Codex 代码解析一

前言

在前面的文章中，我在 MacOS 中使用了 Zsh Codex 插件来实现 AI 命令自动补全。

我对这个利用AI实现的补全小工具的实现过程非常感兴趣，它的代码也写得非常清晰。

因此，接下来的文章，将一步步对 Zsh Codex 的代码进行解析，更好地理解其工作原理。

仓库地址：github.com/tom-doerr/z…

整体代码

目录结构

打开所有的目录，我们可以看到，它的代码非常清晰，总共有用的就三个文件，分别是：

  • create_completion.py
  • zsh_codex.plugin.zsh
  • services/services.py

根据文件名基本可以看出上述三个文件的作用：

  • create_completion.py：主要是实现了自动补全的逻辑，它会根据用户输入的内容，调整之后调用 service，然后将自动补全的结果输出到终端。
  • zsh_codex.plugin.zsh：主要是实现了 Zsh Codex 插件在 zsh 终端 中的实现逻辑，它会调用 create_completion.py 文件中的 main 函数，然后将自动补全的结果输出到终端。
  • services/services.py：主要是实现了 Zsh Codex 插件的服务逻辑，很显然这个文件就是用来接入 AI 的请求和放入提示词的地方。

create_completion.py

这篇文章我们首先来看中间的文件，不考虑在 Zsh 中如何调用它，以及AI的对接部分。

我们忽略 import 以及 默认的 main 部分，只看核心代码。

def main():
    parser = argparse.ArgumentParser(
        description="Generate command completions using AI."
    )
    parser.add_argument(
        "cursor_position", type=int, help="Cursor position in the input buffer"
    )
    args = parser.parse_args()

    client = ClientFactory.create()
  • 创建一个命令行参数解析器，并设置工具的描述信息。
  • 添加一个必需的参数 cursor_position，用于指定输入缓冲区中的光标位置。
  • 解析命令行参数，并将结果存储在 args 对象中。
  • 使用 ClientFactory 创建一个客户端对象，用于与AI服务进行交互，从 import 中可知，这个对象来自于 services.services
# Read the input prompt from stdin.
buffer = sys.stdin.read()
zsh_prefix = "#!/bin/zsh\n\n"
buffer_prefix = buffer[: args.cursor_position]
buffer_suffix = buffer[args.cursor_position :]
full_command = zsh_prefix + buffer_prefix + buffer_suffix

completion = client.get_completion(full_command)

if completion.startswith(zsh_prefix):
    completion = completion[len(zsh_prefix) :]

line_prefix = buffer_prefix.rsplit("\n", 1)[-1]
  • 代码首先读取标准输入的内容
  • 然后根据光标位置将输入内容拆分成前缀和后缀，这个光标位置来自于 args.cursor_position
  • 接着将这些内容与一个 Zsh脚本 的前缀组合成一个完整的命令
  • 最后调用一个 AI服务 来获取命令的补全建议
  • 并对补全建议进行处理，去除多余的前缀。
# 1
for prefix in [buffer_prefix, line_prefix]:
    if completion.startswith(prefix):
        completion = completion[len(prefix) :]
        break
# 2
if buffer_suffix and completion.endswith(buffer_suffix):
    completion = completion[: -len(buffer_suffix)]
# 3
completion = completion.strip("\n")

# 4
if line_prefix.strip().startswith("#"):
    completion = "\n" + completion

sys.stdout.write(completion)
  • 代码处理了不同的情况，根据前缀和后缀来处理补全建议
  • 第一种情况：如果补全建议以前缀开头，则去除前缀
  • 第二种情况：如果补全建议以后缀结尾，则去除后缀
  • 第三种情况：如果补全建议以换行符开头，则去除换行符
  • 第四种情况：如果前缀以注释开头，则在补全建议前添加一个换行符，这种情况就是上一篇文章中的那个动图，在命令行中写一个注释，可以看到它补全的内容是换行的，这样就可以直接回车执行了。
  • 最后将处理后的补全建议写入标准输出。

总结

总的来说，这个脚本的主要功能是：

  1. 处理标准输出并处理成需要使用的几个参数，包括前缀、后缀、完整命令等。
  2. 调用AI服务获取补全建议。
  3. 对补全建议进行处理，去除多余的前缀和后缀。
  4. 处理不同的情况，根据前缀和后缀来处理补全建议。
  5. 将处理后的补全建议写入标准输出。

可以看出这个文件是整个插件的逻辑的核心部分，它的代码非常清晰，非常容易理解。

下一篇文章我们来分析 services 部分的代码。

– 欢迎点赞、关注、转发、收藏【我码玄黄】，各大平台同名。

avatar
文章
阅读
粉丝