instructor
132 строки · 3.9 Кб
1import enum2import instructor3
4from openai import OpenAI5from pydantic import BaseModel, Field6
7
8client = instructor.from_openai(OpenAI())9
10
11class NodeType(str, enum.Enum):12"""Enumeration representing the types of nodes in a filesystem."""13
14FILE = "file"15FOLDER = "folder"16
17
18class Node(BaseModel):19"""20Class representing a single node in a filesystem. Can be either a file or a folder.
21Note that a file cannot have children, but a folder can.
22
23Args:
24name (str): The name of the node.
25children (List[Node]): The list of child nodes (if any).
26node_type (NodeType): The type of the node, either a file or a folder.
27
28Methods:
29print_paths: Prints the path of the node and its children.
30"""
31
32name: str = Field(..., description="Name of the folder")33children: list["Node"] = Field(34default_factory=list,35description="List of children nodes, only applicable for folders, files cannot have children",36)37node_type: NodeType = Field(38default=NodeType.FILE,39description="Either a file or folder, use the name to determine which it could be",40)41
42def print_paths(self, parent_path=""):43"""Prints the path of the node and its children."""44
45if self.node_type == NodeType.FOLDER:46path = f"{parent_path}/{self.name}" if parent_path != "" else self.name47
48print(path, self.node_type)49
50if self.children is not None:51for child in self.children:52child.print_paths(path)53else:54print(f"{parent_path}/{self.name}", self.node_type)55
56
57class DirectoryTree(BaseModel):58"""59Container class representing a directory tree.
60
61Args:
62root (Node): The root node of the tree.
63
64Methods:
65print_paths: Prints the paths of the root node and its children.
66"""
67
68root: Node = Field(..., description="Root folder of the directory tree")69
70def print_paths(self):71"""Prints the paths of the root node and its children."""72
73self.root.print_paths()74
75
76Node.model_rebuild()77DirectoryTree.model_rebuild()78
79
80def parse_tree_to_filesystem(data: str) -> DirectoryTree:81"""82Convert a string representing a directory tree into a filesystem structure
83using OpenAI's GPT-3 model.
84
85Args:
86data (str): The string to convert into a filesystem.
87
88Returns:
89DirectoryTree: The directory tree representing the filesystem.
90"""
91
92completion = client.chat.completions.create(93model="gpt-3.5-turbo-0613",94response_model=DirectoryTree,95messages=[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],105max_tokens=1000,106)107root = DirectoryTree.from_response(completion)108return root109
110
111if __name__ == "__main__":112root = parse_tree_to_filesystem(113"""114root
115├── folder1
116│ ├── file1.txt
117│ └── file2.txt
118└── folder2
119├── file3.txt
120└── subfolder1
121└── file4.txt
122"""
123)124root.print_paths()125# >>> root NodeType.FOLDER126# >>> root/folder1 NodeType.FOLDER127# >>> root/folder1/file1.txt NodeType.FILE128# >>> root/folder1/file2.txt NodeType.FILE129# >>> root/folder2 NodeType.FOLDER130# >>> root/folder2/file3.txt NodeType.FILE131# >>> root/folder2/subfolder1 NodeType.FOLDER132# >>> root/folder2/subfolder1/file4.txt NodeType.FILE133