instructor

Форк
0
/
parse_recursive_paths.py 
132 строки · 3.9 Кб
1
import enum
2
import instructor
3

4
from openai import OpenAI
5
from pydantic import BaseModel, Field
6

7

8
client = instructor.from_openai(OpenAI())
9

10

11
class NodeType(str, enum.Enum):
12
    """Enumeration representing the types of nodes in a filesystem."""
13

14
    FILE = "file"
15
    FOLDER = "folder"
16

17

18
class Node(BaseModel):
19
    """
20
    Class representing a single node in a filesystem. Can be either a file or a folder.
21
    Note that a file cannot have children, but a folder can.
22

23
    Args:
24
        name (str): The name of the node.
25
        children (List[Node]): The list of child nodes (if any).
26
        node_type (NodeType): The type of the node, either a file or a folder.
27

28
    Methods:
29
        print_paths: Prints the path of the node and its children.
30
    """
31

32
    name: str = Field(..., description="Name of the folder")
33
    children: list["Node"] = Field(
34
        default_factory=list,
35
        description="List of children nodes, only applicable for folders, files cannot have children",
36
    )
37
    node_type: NodeType = Field(
38
        default=NodeType.FILE,
39
        description="Either a file or folder, use the name to determine which it could be",
40
    )
41

42
    def print_paths(self, parent_path=""):
43
        """Prints the path of the node and its children."""
44

45
        if self.node_type == NodeType.FOLDER:
46
            path = f"{parent_path}/{self.name}" if parent_path != "" else self.name
47

48
            print(path, self.node_type)
49

50
            if self.children is not None:
51
                for child in self.children:
52
                    child.print_paths(path)
53
        else:
54
            print(f"{parent_path}/{self.name}", self.node_type)
55

56

57
class DirectoryTree(BaseModel):
58
    """
59
    Container class representing a directory tree.
60

61
    Args:
62
        root (Node): The root node of the tree.
63

64
    Methods:
65
        print_paths: Prints the paths of the root node and its children.
66
    """
67

68
    root: Node = Field(..., description="Root folder of the directory tree")
69

70
    def print_paths(self):
71
        """Prints the paths of the root node and its children."""
72

73
        self.root.print_paths()
74

75

76
Node.model_rebuild()
77
DirectoryTree.model_rebuild()
78

79

80
def parse_tree_to_filesystem(data: str) -> DirectoryTree:
81
    """
82
    Convert a string representing a directory tree into a filesystem structure
83
    using OpenAI's GPT-3 model.
84

85
    Args:
86
        data (str): The string to convert into a filesystem.
87

88
    Returns:
89
        DirectoryTree: The directory tree representing the filesystem.
90
    """
91

92
    completion = client.chat.completions.create(
93
        model="gpt-3.5-turbo-0613",
94
        response_model=DirectoryTree,
95
        messages=[
96
            {
97
                "role": "system",
98
                "content": "You are a perfect file system parsing algorithm. You are given a string representing a directory tree. You must return the correct filesystem structure.",
99
            },
100
            {
101
                "role": "user",
102
                "content": f"Consider the data below:\n{data} and return the correctly labeled filesystem",
103
            },
104
        ],
105
        max_tokens=1000,
106
    )
107
    root = DirectoryTree.from_response(completion)
108
    return root
109

110

111
if __name__ == "__main__":
112
    root = parse_tree_to_filesystem(
113
        """
114
        root
115
        ├── folder1
116
        │   ├── file1.txt
117
        │   └── file2.txt
118
        └── folder2
119
            ├── file3.txt
120
            └── subfolder1
121
                └── file4.txt
122
        """
123
    )
124
    root.print_paths()
125
    # >>> root                                  NodeType.FOLDER
126
    # >>> root/folder1                          NodeType.FOLDER
127
    # >>> root/folder1/file1.txt                NodeType.FILE
128
    # >>> root/folder1/file2.txt                NodeType.FILE
129
    # >>> root/folder2                          NodeType.FOLDER
130
    # >>> root/folder2/file3.txt                NodeType.FILE
131
    # >>> root/folder2/subfolder1               NodeType.FOLDER
132
    # >>> root/folder2/subfolder1/file4.txt     NodeType.FILE
133

Использование cookies

Мы используем файлы cookie в соответствии с Политикой конфиденциальности и Политикой использования cookies.

Нажимая кнопку «Принимаю», Вы даете АО «СберТех» согласие на обработку Ваших персональных данных в целях совершенствования нашего веб-сайта и Сервиса GitVerse, а также повышения удобства их использования.

Запретить использование cookies Вы можете самостоятельно в настройках Вашего браузера.