(code: str, kernel: CodeKernel)
| 155 | |
| 156 | |
| 157 | def execute(code: str, kernel: CodeKernel) -> tuple[Literal["text", "image"] | None, str]: |
| 158 | res = "" |
| 159 | res_type = None |
| 160 | code = code.replace("<|observation|>", "") |
| 161 | code = code.replace("<|assistant|>python", "") |
| 162 | code = code.replace("<|assistant|>", "") |
| 163 | code = code.replace("<|user|>", "") |
| 164 | code = code.replace("<|system|>", "") |
| 165 | msg, output = kernel.execute(code) |
| 166 | |
| 167 | if msg["metadata"]["status"] == "timeout": |
| 168 | return res_type, "Timed out" |
| 169 | elif msg["metadata"]["status"] == "error": |
| 170 | return res_type, clean_ansi_codes("\n".join(kernel.get_error_msg(msg, verbose=True))) |
| 171 | |
| 172 | if "text" in output: |
| 173 | res_type = "text" |
| 174 | res = output["text"] |
| 175 | elif "data" in output: |
| 176 | for key in output["data"]: |
| 177 | if "text/plain" in key: |
| 178 | res_type = "text" |
| 179 | res = output["data"][key] |
| 180 | elif "image/png" in key: |
| 181 | res_type = "image" |
| 182 | res = output["data"][key] |
| 183 | break |
| 184 | |
| 185 | return res_type, res |
| 186 | |
| 187 | |
| 188 | @st.cache_resource |
no test coverage detected