Chat with Codes

This example illustrates how to read a repository from file system and use that as contexts when chatting with the AI assistant. The prompt interleaved with the program codes behaves like comments that explain the codes after that.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
# Usage: python examples/advanced/chat_with_codes.py

import glob
import os
from argparse import ArgumentParser

import seedir as sd
from loguru import logger
from prompt_toolkit.shortcuts import prompt
from rich.console import Console

from appl import AIRole, convo, gen, ppl
from appl.compositor import Tagged
from appl.core import load_file, make_panel
from appl.utils import get_num_tokens

parser = ArgumentParser()
parser.add_argument("--intro-file", type=str, default="./README.md")
parser.add_argument("--source", type=str, default="./examples")
parser.add_argument("--ext", type=str, default=".py")

args = parser.parse_args()


def read_file(file: str):
    with open(file, "r", encoding="utf-8") as f:
        return f.read()


@ppl
def chat(intro_file: str, source_folder: str, ext: str = ".py"):
    with Tagged("introduction"):
        read_file(intro_file)

    with Tagged("directory", attrs={"name": source_folder}):
        "The source code is organized as follows:"
        (
            tree := sd.seedir(
                source_folder,
                style="spaces",
                printout=False,
                exclude_folders=["__pycache__"],
                exclude_files=[".DS_Store"],
            )
        )
        print(f"Chatting with {ext} files in the directory:\n{tree}")

    with Tagged("source"):
        "The contents of the source code are as follows:"

        for name in glob.glob(
            os.path.join(source_folder, "**", f"*{ext}"), recursive=True
        ):
            with Tagged("file", attrs={"name": name}):
                "```python"
                read_file(name)  # load the file contents and put in the prompt
                "```"

    "Now begin the chat about the project:\n"
    print(f"num_tokens: {get_num_tokens(str(convo()))}")

    console = Console()
    while True:
        (query := prompt("User: "))
        if query.startswith("exit"):
            break
        console.print(make_panel(query, title="User"))
        logger.info(f"User: {query}")
        with AIRole():
            (
                res := str(
                    gen(stream=True, max_relay_rounds=10).streaming(title="Assistant")
                )
            )
        logger.info(f"Assistant: {res}")


if __name__ == "__main__":
    chat(args.intro_file, args.source, args.ext)