Import and initialize our tool spec

%pip install llama-index-agent-openai
%pip install llama-index-llms-openai
%pip install llama-index-tools-code-interpreter

!pip install llama-index

import os

os.environ["OPENAI_API_KEY"] = "sk-..."

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI

# Import and initialize our tool spec
from llama_index.tools.code_interpreter.base import CodeInterpreterToolSpec

code_spec = CodeInterpreterToolSpec()

tools = code_spec.to_tool_list()

# Create the Agent with our tools
agent = FunctionAgent(
    tools=tools,
    llm=OpenAI(model="gpt-4.1"),
)

# Context to store chat history
from llama_index.core.workflow import Context
ctx = Context(agent)

# Prime the Agent to use the tool
print(
    await agent.run(
        "Can you help me write some python code to pass to the code_interpreter tool",
        ctx=ctx
    )
)

print(
    await agent.run(
        """There is a world_happiness_2016.csv file in the `data` directory (relative path).
                 Can you write and execute code to tell me columns does it have?""",
        ctx=ctx,
    )
)

print(await agent.run("What are the top 10 happiest countries", ctx=ctx))

print(await agent.run("Can you make a graph of the top 10 happiest countries", ctx=ctx))

print(await agent.run("Can you make a graph of the top 10 happiest countries", ctx=ctx))

print(await agent.run("can you also plot the 10 lowest", ctx=ctx))

print(await agent.run("can you do it in one plot", ctx=ctx))