|
|
import json |
|
|
from typing import Optional |
|
|
|
|
|
from injector import inject |
|
|
|
|
|
from taskweaver.code_interpreter.code_executor import CodeExecutor |
|
|
from taskweaver.code_interpreter.code_generator import CodeGeneratorPluginOnly |
|
|
from taskweaver.config.module_config import ModuleConfig |
|
|
from taskweaver.logging import TelemetryLogger |
|
|
from taskweaver.memory import Memory, Post |
|
|
from taskweaver.memory.attachment import AttachmentType |
|
|
from taskweaver.role import Role |
|
|
|
|
|
|
|
|
class CodeInterpreterConfig(ModuleConfig): |
|
|
def _configure(self): |
|
|
self._set_name("code_interpreter_plugin_only") |
|
|
self.use_local_uri = self._get_bool("use_local_uri", False) |
|
|
self.max_retry_count = self._get_int("max_retry_count", 3) |
|
|
|
|
|
|
|
|
class CodeInterpreterPluginOnly(Role): |
|
|
@inject |
|
|
def __init__( |
|
|
self, |
|
|
generator: CodeGeneratorPluginOnly, |
|
|
executor: CodeExecutor, |
|
|
logger: TelemetryLogger, |
|
|
config: CodeInterpreterConfig, |
|
|
): |
|
|
self.generator = generator |
|
|
self.executor = executor |
|
|
self.logger = logger |
|
|
self.config = config |
|
|
self.retry_count = 0 |
|
|
self.return_index = 0 |
|
|
|
|
|
self.logger.info("CodeInterpreter initialized successfully.") |
|
|
|
|
|
def reply( |
|
|
self, |
|
|
memory: Memory, |
|
|
event_handler: callable, |
|
|
prompt_log_path: Optional[str] = None, |
|
|
use_back_up_engine: Optional[bool] = False, |
|
|
) -> Post: |
|
|
response: Post = self.generator.reply( |
|
|
memory, |
|
|
event_handler, |
|
|
) |
|
|
|
|
|
if response.message is not None: |
|
|
return response |
|
|
|
|
|
functions = json.loads(response.get_attachment(type=AttachmentType.function)[0]) |
|
|
if len(functions) > 0: |
|
|
code = [] |
|
|
for i, f in enumerate(functions): |
|
|
function_name = f["name"] |
|
|
function_args = json.loads(f["arguments"]) |
|
|
function_call = ( |
|
|
f"r{self.return_index + i}={function_name}(" |
|
|
+ ", ".join( |
|
|
[ |
|
|
f'{key}="{value}"' if isinstance(value, str) else f"{key}={value}" |
|
|
for key, value in function_args.items() |
|
|
], |
|
|
) |
|
|
+ ")" |
|
|
) |
|
|
code.append(function_call) |
|
|
code.append(f'{", ".join([f"r{self.return_index + i}" for i in range(len(functions))])}') |
|
|
self.return_index += len(functions) |
|
|
|
|
|
event_handler("code", "\n".join(code)) |
|
|
exec_result = self.executor.execute_code( |
|
|
exec_id=response.id, |
|
|
code="\n".join(code), |
|
|
) |
|
|
|
|
|
response.message = self.executor.format_code_output( |
|
|
exec_result, |
|
|
with_code=True, |
|
|
use_local_uri=self.config.use_local_uri, |
|
|
) |
|
|
event_handler("CodeInterpreter-> Planner", response.message) |
|
|
else: |
|
|
response.message = "No code is generated because no function is selected." |
|
|
event_handler("CodeInterpreter-> Planner", response.message) |
|
|
|
|
|
return response |
|
|
|
