MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/5/11 14:45
5
@Author  : alexanderwu
6
@File    : test_llm.py
7
"""
8

9
import pytest
10

11
from metagpt.llm import LLM
12

13
CODE_REVIEW_SMALLEST_CONTEXT = """
14
## game.js
15
```Code
16
// game.js
17
class Game {
18
    constructor() {
19
        this.board = this.createEmptyBoard();
20
        this.score = 0;
21
        this.bestScore = 0;
22
    }
23

24
    createEmptyBoard() {
25
        const board = [];
26
        for (let i = 0; i < 4; i++) {
27
            board[i] = [0, 0, 0, 0];
28
        }
29
        return board;
30
    }
31

32
    startGame() {
33
        this.board = this.createEmptyBoard();
34
        this.score = 0;
35
        this.addRandomTile();
36
        this.addRandomTile();
37
    }
38

39
    addRandomTile() {
40
        let emptyCells = [];
41
        for (let r = 0; r < 4; r++) {
42
            for (let c = 0; c < 4; c++) {
43
                if (this.board[r][c] === 0) {
44
                    emptyCells.push({ r, c });
45
                }
46
            }
47
        }
48
        if (emptyCells.length > 0) {
49
            let randomCell = emptyCells[Math.floor(Math.random() * emptyCells.length)];
50
            this.board[randomCell.r][randomCell.c] = Math.random() < 0.9 ? 2 : 4;
51
        }
52
    }
53

54
    move(direction) {
55
        // This function will handle the logic for moving tiles
56
        // in the specified direction and merging them
57
        // It will also update the score and add a new random tile if the move is successful
58
        // The actual implementation of this function is complex and would require
59
        // a significant amount of code to handle all the cases for moving and merging tiles
60
        // For the purposes of this example, we will not implement the full logic
61
        // Instead, we will just call addRandomTile to simulate a move
62
        this.addRandomTile();
63
    }
64

65
    getBoard() {
66
        return this.board;
67
    }
68

69
    getScore() {
70
        return this.score;
71
    }
72

73
    getBestScore() {
74
        return this.bestScore;
75
    }
76

77
    setBestScore(score) {
78
        this.bestScore = score;
79
    }
80
}
81

82
```
83
"""
84

85
MOVE_DRAFT = """
86
## move function draft
87

88
```javascript
89
move(direction) {
90
    let moved = false;
91
    switch (direction) {
92
        case 'up':
93
            for (let c = 0; c < 4; c++) {
94
                for (let r = 1; r < 4; r++) {
95
                    if (this.board[r][c] !== 0) {
96
                        let row = r;
97
                        while (row > 0 && this.board[row - 1][c] === 0) {
98
                            this.board[row - 1][c] = this.board[row][c];
99
                            this.board[row][c] = 0;
100
                            row--;
101
                            moved = true;
102
                        }
103
                        if (row > 0 && this.board[row - 1][c] === this.board[row][c]) {
104
                            this.board[row - 1][c] *= 2;
105
                            this.board[row][c] = 0;
106
                            this.score += this.board[row - 1][c];
107
                            moved = true;
108
                        }
109
                    }
110
                }
111
            }
112
            break;
113
        case 'down':
114
            // Implement logic for moving tiles down
115
            // Similar to the 'up' case but iterating in reverse order
116
            // and checking for merging in the opposite direction
117
            break;
118
        case 'left':
119
            // Implement logic for moving tiles left
120
            // Similar to the 'up' case but iterating over columns first
121
            // and checking for merging in the opposite direction
122
            break;
123
        case 'right':
124
            // Implement logic for moving tiles right
125
            // Similar to the 'up' case but iterating over columns in reverse order
126
            // and checking for merging in the opposite direction
127
            break;
128
    }
129

130
    if (moved) {
131
        this.addRandomTile();
132
    }
133
}
134
```
135
"""
136

137
FUNCTION_TO_MERMAID_CLASS = """
138
## context
139
```
140
class UIDesign(Action):
141
    #Class representing the UI Design action.
142
    def __init__(self, name, context=None, llm=None):
143
        super().__init__(name, context, llm)  # 需要调用LLM进一步丰富UI设计的prompt
144
    @parse
145
    def parse_requirement(self, context: str):
146
        #Parse UI Design draft from the context using regex.
147
        pattern = r"## UI Design draft.*?\n(.*?)## Anything UNCLEAR"
148
        return context, pattern
149
    @parse
150
    def parse_ui_elements(self, context: str):
151
        #Parse Selected Elements from the context using regex.
152
        pattern = r"## Selected Elements.*?\n(.*?)## HTML Layout"
153
        return context, pattern
154
    @parse
155
    def parse_css_code(self, context: str):
156
        pattern = r"```css.*?\n(.*?)## Anything UNCLEAR"
157
        return context, pattern
158
    @parse
159
    def parse_html_code(self, context: str):
160
        pattern = r"```html.*?\n(.*?)```"
161
        return context, pattern
162
    async def draw_icons(self, context, *args, **kwargs):
163
        #Draw icons using SDEngine.
164
        engine = SDEngine()
165
        icon_prompts = self.parse_ui_elements(context)
166
        icons = icon_prompts.split("\n")
167
        icons = [s for s in icons if len(s.strip()) > 0]
168
        prompts_batch = []
169
        for icon_prompt in icons:
170
            # fixme: 添加icon lora
171
            prompt = engine.construct_payload(icon_prompt + ".<lora:WZ0710_AW81e-3_30e3b128d64T32_goon0.5>")
172
            prompts_batch.append(prompt)
173
        await engine.run_t2i(prompts_batch)
174
        logger.info("Finish icon design using StableDiffusion API")
175
    async def _save(self, css_content, html_content):
176
        save_dir = CONFIG.workspace_path / "resources" / "codes"
177
        if not os.path.exists(save_dir):
178
            os.makedirs(save_dir, exist_ok=True)
179
        # Save CSS and HTML content to files
180
        css_file_path = save_dir / "ui_design.css"
181
        html_file_path = save_dir / "ui_design.html"
182
        with open(css_file_path, "w") as css_file:
183
            css_file.write(css_content)
184
        with open(html_file_path, "w") as html_file:
185
            html_file.write(html_content)
186
    async def run(self, requirements: list[Message], *args, **kwargs) -> ActionOutput:
187
        #Run the UI Design action.
188
        # fixme: update prompt (根据需求细化prompt）
189
        context = requirements[-1].content
190
        ui_design_draft = self.parse_requirement(context=context)
191
        # todo: parse requirements str
192
        prompt = PROMPT_TEMPLATE.format(context=ui_design_draft, format_example=FORMAT_EXAMPLE)
193
        logger.info(prompt)
194
        ui_describe = await self._aask_v1(prompt, "ui_design", OUTPUT_MAPPING)
195
        logger.info(ui_describe.content)
196
        logger.info(ui_describe.instruct_content)
197
        css = self.parse_css_code(context=ui_describe.content)
198
        html = self.parse_html_code(context=ui_describe.content)
199
        await self._save(css_content=css, html_content=html)
200
        await self.draw_icons(ui_describe.content)
201
        return ui_describe
202
```
203
-----
204
## format example
205
[CONTENT]
206
{
207
    "ClassView": "classDiagram\n        class A {\n        -int x\n        +int y\n        -int speed\n        -int direction\n        +__init__(x: int, y: int, speed: int, direction: int)\n        +change_direction(new_direction: int) None\n        +move() None\n    }\n    "
208
}
209
[/CONTENT]
210
## nodes: "<node>: <type>  # <comment>"
211
- ClassView: <class 'str'>  # Generate the mermaid class diagram corresponding to source code in "context."
212
## constraint
213
- Language: Please use the same language as the user input.
214
- Format: output wrapped inside [CONTENT][/CONTENT] as format example, nothing else.
215
## action
216
Fill in the above nodes(ClassView) based on the format example.
217
"""
218

219
MOVE_FUNCTION = """
220
## move function implementation
221

222
```javascript
223
move(direction) {
224
    let moved = false;
225
    switch (direction) {
226
        case 'up':
227
            for (let c = 0; c < 4; c++) {
228
                for (let r = 1; r < 4; r++) {
229
                    if (this.board[r][c] !== 0) {
230
                        let row = r;
231
                        while (row > 0 && this.board[row - 1][c] === 0) {
232
                            this.board[row - 1][c] = this.board[row][c];
233
                            this.board[row][c] = 0;
234
                            row--;
235
                            moved = true;
236
                        }
237
                        if (row > 0 && this.board[row - 1][c] === this.board[row][c]) {
238
                            this.board[row - 1][c] *= 2;
239
                            this.board[row][c] = 0;
240
                            this.score += this.board[row - 1][c];
241
                            moved = true;
242
                        }
243
                    }
244
                }
245
            }
246
            break;
247
        case 'down':
248
            for (let c = 0; c < 4; c++) {
249
                for (let r = 2; r >= 0; r--) {
250
                    if (this.board[r][c] !== 0) {
251
                        let row = r;
252
                        while (row < 3 && this.board[row + 1][c] === 0) {
253
                            this.board[row + 1][c] = this.board[row][c];
254
                            this.board[row][c] = 0;
255
                            row++;
256
                            moved = true;
257
                        }
258
                        if (row < 3 && this.board[row + 1][c] === this.board[row][c]) {
259
                            this.board[row + 1][c] *= 2;
260
                            this.board[row][c] = 0;
261
                            this.score += this.board[row + 1][c];
262
                            moved = true;
263
                        }
264
                    }
265
                }
266
            }
267
            break;
268
        case 'left':
269
            for (let r = 0; r < 4; r++) {
270
                for (let c = 1; c < 4; c++) {
271
                    if (this.board[r][c] !== 0) {
272
                        let col = c;
273
                        while (col > 0 && this.board[r][col - 1] === 0) {
274
                            this.board[r][col - 1] = this.board[r][col];
275
                            this.board[r][col] = 0;
276
                            col--;
277
                            moved = true;
278
                        }
279
                        if (col > 0 && this.board[r][col - 1] === this.board[r][col]) {
280
                            this.board[r][col - 1] *= 2;
281
                            this.board[r][col] = 0;
282
                            this.score += this.board[r][col - 1];
283
                            moved = true;
284
                        }
285
                    }
286
                }
287
            }
288
            break;
289
        case 'right':
290
            for (let r = 0; r < 4; r++) {
291
                for (let c = 2; c >= 0; c--) {
292
                    if (this.board[r][c] !== 0) {
293
                        let col = c;
294
                        while (col < 3 && this.board[r][col + 1] === 0) {
295
                            this.board[r][col + 1] = this.board[r][col];
296
                            this.board[r][col] = 0;
297
                            col++;
298
                            moved = true;
299
                        }
300
                        if (col < 3 && this.board[r][col + 1] === this.board[r][col]) {
301
                            this.board[r][col + 1] *= 2;
302
                            this.board[r][col] = 0;
303
                            this.score += this.board[r][col + 1];
304
                            moved = true;
305
                        }
306
                    }
307
                }
308
            }
309
            break;
310
    }
311

312
    if (moved) {
313
        this.addRandomTile();
314
    }
315
}
316
```
317
"""
318

319

320
@pytest.fixture()
321
def llm():
322
    return LLM()
323

324

325
@pytest.mark.asyncio
326
async def test_llm_code_review(llm):
327
    choices = [
328
        "Please review the move function code above. Should it be refactor?",
329
        "Please implement the move function",
330
        "Please write a draft for the move function in order to implement it",
331
    ]
332
    # prompt = CODE_REVIEW_SMALLEST_CONTEXT+ "\n\n" + MOVE_DRAFT + "\n\n" + choices[1]
333
    # rsp = await llm.aask(prompt)
334

335
    prompt = CODE_REVIEW_SMALLEST_CONTEXT + "\n\n" + MOVE_FUNCTION + "\n\n" + choices[0]
336
    prompt = FUNCTION_TO_MERMAID_CLASS
337

338
    _ = await llm.aask(prompt)
339

340

341
# if __name__ == "__main__":
342
#     pytest.main([__file__, "-s"])
343