From 4f93c5fad3f03fd0302e3a93760216fc9ca58ffa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= Date: Mon, 15 Jan 2024 16:59:49 +0800 Subject: [PATCH] add only_code arg for WriteCodeByGenerate. --- metagpt/actions/write_analysis_code.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py index 04cad34a5..76d47ba28 100644 --- a/metagpt/actions/write_analysis_code.py +++ b/metagpt/actions/write_analysis_code.py @@ -88,8 +88,14 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode): ) -> str: # context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user")) prompt = self.process_msg(context, system_msg) + is_only_code = kwargs.pop("only_code", True) + code_content = await self.llm.aask_code(prompt, **kwargs) - return code_content["code"] + if is_only_code: + return code_content["code"] + else: + return code_content + class WriteCodeWithTools(BaseWriteAnalysisCode):