| import logging |
| import os |
| import time |
| from typing import Any, Dict, Union |
|
|
| from werkzeug.datastructures import FileStorage |
|
|
| from ..agent import BaseAgent |
| from ..agent.react import AsyncReactAgent |
| from ..schemas import AgentRequest, MediaFile, Message, RoleType |
| from ..utils import generate_random_string, get_logger, get_model_config_path |
|
|
| logger = get_logger() |
|
|
|
|
| class CodeInterpreterSession: |
| def __init__( |
| self, |
| session_id: Union[None, str] = None, |
| model_name: Union[None, str] = "openai", |
| config_path: Union[None, str] = None, |
| agent: AsyncReactAgent = None, |
| **kwargs): |
| self.session_id = session_id |
| self.config_path = config_path |
| self.input_files = [] |
| self.output_files = [] |
| self.messages = [] |
| self.agent = agent |
| self.llm_model_name = self.agent.llm.model_name |
|
|
| logger.info("Use model {} and llm in config {} for conversation {}" |
| .format(model_name, self.llm_model_name, self.config_path, self.session_id)) |
|
|
| @classmethod |
| async def create(cls, |
| model_name: Union[None, str] = "openai", |
| config_path: Union[None, str] = None, |
| **kwargs: Dict[str, Any]): |
| if config_path is None: |
| config_path = get_model_config_path(model_name) |
| logger.info(f"Use Config Path: {config_path}") |
|
|
| sandbox_id = generate_random_string(12) |
|
|
| |
| agent = await BaseAgent.async_from_config_path_and_kwargs(config_path, **kwargs) |
| await agent.plugins_map["python_code_sandbox"].set_sandbox_id(sandbox_id) |
|
|
| return cls(session_id=sandbox_id, |
| model_name=model_name, |
| config_path=config_path, |
| agent=agent) |
|
|
| async def upload_to_sandbox(self, file: Union[str, FileStorage]): |
| dst_path = await self.agent.sync_to_sandbox(file) |
| message = f'User uploaded the following files: {dst_path}\n' |
| logging.info(f"The file path {file} has been synced to sandbox with file path {dst_path}") |
| self.messages.append(Message(RoleType.System, message)) |
| self.input_files.append(MediaFile(file_name=os.path.basename(dst_path), sandbox_path=dst_path)) |
|
|
| async def chat(self, user_messages, input_files=None): |
| start_time = time.time() |
|
|
| self.messages.extend(user_messages) |
| agent_request = AgentRequest( |
| messages=self.messages, |
| input_files=self.input_files, |
| sandbox_id=self.session_id |
| ) |
| logger.info(f"Agent request: {agent_request.__dict__}") |
|
|
| async for agent_response in self.agent.async_run(agent_request): |
| logger.info(f"Agent response:\n{agent_response.output_text}") |
| self.messages.append(Message(RoleType.System, agent_response.output_text)) |
| yield agent_response |
|
|
| exec_time = time.time() |
| logger.info( |
| f'Agent Execution Latency: {exec_time - start_time}' |
| ) |
|
|
| def __enter__(self): |
| pass |
|
|
| def __exit__(self, exc_type, exc_value, traceback) -> None: |
| pass |
|
|
