import uuid, base64, re

from io import BytesIO

from typing import Optional

from codeboxapi import CodeBox # type: ignore

from codeboxapi.schema import CodeBoxOutput # type: ignore

from langchain.tools import StructuredTool, BaseTool

from langchain.chat_models import ChatOpenAI

from langchain.chat_models.base import BaseChatModel

from langchain.prompts.chat import MessagesPlaceholder

from langchain.agents import AgentExecutor, BaseSingleActionAgent

from langchain.memory import ConversationBufferMemory

from codeinterpreterapi.schema import CodeInterpreterResponse, CodeInput, File, UserRequest

from codeinterpreterapi.config import settings

from codeinterpreterapi.chains.functions_agent import OpenAIFunctionsAgent

from codeinterpreterapi.prompts import code_interpreter_system_message

from codeinterpreterapi.callbacks import CodeCallbackHandler

from codeinterpreterapi.chains.modifications_check import get_file_modifications

from codeinterpreterapi.chains.remove_download_link import remove_download_link

class CodeInterpreterSession:

def __init__(

self,

model=None,

openai_api_key=settings.OPENAI_API_KEY,

verbose=settings.VERBOSE,

tools: list[BaseTool] = None

) -> None:

self.codebox = CodeBox()

self.verbose = verbose

self.tools: list[BaseTool] = self._tools(tools)

self.llm: BaseChatModel = self._llm(model, openai_api_key)

self.agent_executor: AgentExecutor = self._agent_executor()

self.input_files: list[File] = []

self.output_files: list[File] = []

async def astart(self) -> None:

await self.codebox.astart()

def _tools(self, additional_tools: list[BaseTool] = None) -> list[BaseTool]:

additional_tools = additional_tools or []

return additional_tools + [

StructuredTool(

name="python",

description=

# TODO: variables as context to the agent

# TODO: current files as context to the agent

"Input a string of code to a python interpreter (jupyter kernel). "

"Variables are preserved between runs. ",

func=self.run_handler,

coroutine=self.arun_handler,

args_schema=CodeInput,

),

]

def _llm(self, model: Optional[str] = None, openai_api_key: Optional[str] = None) -> BaseChatModel:

if model is None:

model = "gpt-4"

if openai_api_key is None:

raise ValueError(

"OpenAI API key missing. Set OPENAI_API_KEY env variable or pass `openai_api_key` to session."

)

return ChatOpenAI(

temperature=0.03,

model=model,

openai_api_key=openai_api_key,

max_retries=3,

request_timeout=60 * 3,

) # type: ignore

def _agent(self) -> BaseSingleActionAgent:

return OpenAIFunctionsAgent.from_llm_and_tools(

llm=self.llm,

tools=self.tools,

system_message=code_interpreter_system_message,

extra_prompt_messages=[MessagesPlaceholder(variable_name="memory")],

)

def _agent_executor(self) -> AgentExecutor:

return AgentExecutor.from_agent_and_tools(

agent=self._agent(),

callbacks=[CodeCallbackHandler(self)],

max_iterations=9,

tools=self.tools,

verbose=self.verbose,

memory=ConversationBufferMemory(memory_key="memory", return_messages=True),

)

async def show_code(self, code: str) -> None:

"""Callback function to show code to the user."""

if self.verbose:

print(code)

def run_handler(self, code: str):

raise NotImplementedError("Use arun_handler for now.")

async def arun_handler(self, code: str):

"""Run code in container and send the output to the user"""

output: CodeBoxOutput = await self.codebox.arun(code)

if not isinstance(output.content, str):

raise TypeError("Expected output.content to be a string.")

if output.type == "image/png":

filename = f"image-{uuid.uuid4()}.png"

file_buffer = BytesIO(base64.b64decode(output.content))

file_buffer.name = filename

self.output_files.append(File(name=filename, content=file_buffer.read()))

return f"Image {filename} got send to the user."

elif output.type == "error":

if "ModuleNotFoundError" in output.content:

if package := re.search(

r"ModuleNotFoundError: No module named '(.*)'", output.content

):

await self.codebox.ainstall(package.group(1))

return f"{package.group(1)} was missing but got installed now. Please try again."

else: pass

# TODO: preanalyze error to optimize next code generation

if self.verbose:

print("Error:", output.content)

elif modifications := await get_file_modifications(code, self.llm):

for filename in modifications:

if filename in [file.name for file in self.input_files]:

continue

fileb = await self.codebox.adownload(filename)

if not fileb.content:

continue

file_buffer = BytesIO(fileb.content)

file_buffer.name = filename

self.output_files.append(

File(name=filename, content=file_buffer.read())

)

return output.content

async def input_handler(self, request: UserRequest):

if not request.files:

return

if not request.content:

request.content = (

"I uploaded, just text me back and confirm that you got the file(s)."

)

request.content += "\n**The user uploaded the following files: **\n"

for file in request.files:

self.input_files.append(file)

request.content += f"[Attachment: {file.name}]\n"

await self.codebox.aupload(file.name, file.content)

request.content += "**File(s) are now available in the cwd. **\n"

async def output_handler(self, final_response: str) -> CodeInterpreterResponse:

"""Embed images in the response"""

for file in self.output_files:

if str(file.name) in final_response:

# rm  from the response

final_response = re.sub(rf"\n\n!\[.*\]\(.*\)", "", final_response)

if self.output_files and re.search(rf"\n\[.*\]\(.*\)", final_response):

final_response = await remove_download_link(final_response, self.llm)

return CodeInterpreterResponse(content=final_response, files=self.output_files)

async def generate_response(

self,

user_msg: str,

files: list[File] = [],

detailed_error: bool = False,

) -> CodeInterpreterResponse:

"""Generate a Code Interpreter response based on the user's input."""

user_request = UserRequest(content=user_msg, files=files)

try:

await self.input_handler(user_request)

response = await self.agent_executor.arun(input=user_request.content)

return await self.output_handler(response)

except Exception as e:

if self.verbose:

import traceback

traceback.print_exc()

if detailed_error:

return CodeInterpreterResponse(

content=f"Error in CodeInterpreterSession: {e.__class__.__name__} - {e}"

)

else:

return CodeInterpreterResponse(

content="Sorry, something went while generating your response."

"Please try again or restart the session."

)

async def is_running(self) -> bool:

return await self.codebox.astatus() == "running"

async def astop(self) -> None:

await self.codebox.astop()

async def __aenter__(self) -> "CodeInterpreterSession":

await self.astart()

return self

async def __aexit__(self, exc_type, exc_value, traceback) -> None:

await self.astop()