from e2b_code_interpreter import Sandbox
import os

# Define your agent application
agent_code = """
import os
from smolagents import CodeAgent, InferenceClientModel

# Initialize the agents
# Note: Using a specific model that is known to work with fireworks-ai
model_id = "deepseek-ai/DeepSeek-V4-Flash"

agent = CodeAgent(
    model=InferenceClientModel(model_id=model_id, token=os.getenv("HF_TOKEN")),
    tools=[],
    name="coder_agent",
    description="This agent takes care of your difficult algorithmic problems using code."
)

manager_agent = CodeAgent(
    model=InferenceClientModel(model_id=model_id, token=os.getenv("HF_TOKEN")),
    tools=[],
    managed_agents=[agent],
)

# Run the agent
response = manager_agent.run("What's the 5th Fibonacci number?")
print(response)
"""

def run_code_raise_errors(sandbox, code: str) -> str:
    execution = sandbox.run_code(
        code,
        envs={'HF_TOKEN': os.getenv('HF_TOKEN')}
    )
    if execution.error:
        execution_logs = "\n".join([str(log) for log in execution.logs.stdout])
        logs = execution_logs
        logs += execution.error.traceback
        raise ValueError(logs)
    return "\n".join([str(log) for log in execution.logs.stdout])

# Use context manager to ensure sandbox is closed
with Sandbox.create() as sandbox:
    print("Sandbox created. Installing smolagents...")
    # Install required packages
    sandbox.commands.run("pip install smolagents")
    
    print("Running agent code...")
    # Run the agent code in the sandbox
    execution_logs = run_code_raise_errors(sandbox, agent_code)
    print("Execution Result:")
    print(execution_logs)