import os
from dotenv import load_dotenv
from openai import OpenAI
from e2b_code_interpreter import Sandbox
import re
import json
# Load environment variables
load_dotenv()
# API keys
NAI_API_KEY = os.getenv("NAI_API_KEY")
E2B_API_KEY = os.getenv("E2B_API_KEY")
# Model selection
MODEL_NAME = 'nusantara-base' # Alternatives: 'archipelago-70b', 'Llama-3.3-Nemotron-Super-49B'
# System prompt for the LLM
SYSTEM_PROMPT = """You're a Python data scientist. You are given tasks to complete and you run Python code to solve them.
Information about the csv dataset:
- It's in the `/home/user/data.csv` file
- The CSV file is using , as the delimiter
- It has the following columns (examples included):
- country: "Argentina", "Australia"
- Region: "SouthAmerica", "Oceania"
- Surface area (km2): for example, 2780400
- Population in thousands (2017): for example, 44271
- Population density (per km2, 2017): for example, 16.2
- Sex ratio (m per 100 f, 2017): for example, 95.9
- GDP: Gross domestic product (million current US$): for example, 632343
- GDP growth rate (annual %, const. 2005 prices): for example, 2.4
- GDP per capita (current US$): for example, 14564.5
- Economy: Agriculture (% of GVA): for example, 10.0
- Economy: Industry (% of GVA): for example, 28.1
- Economy: Services and other activity (% of GVA): for example, 61.9
- Employment: Agriculture (% of employed): for example, 4.8
- Employment: Industry (% of employed): for example, 20.6
- Employment: Services (% of employed): for example, 74.7
- Unemployment (% of labour force): for example, 8.5
- Employment: Female (% of employed): for example, 43.7
- Employment: Male (% of employed): for example, 56.3
- Labour force participation (female %): for example, 48.5
- Labour force participation (male %): for example, 71.1
- International trade: Imports (million US$): for example, 59253
- International trade: Exports (million US$): for example, 57802
- International trade: Balance (million US$): for example, -1451
- Education: Government expenditure (% of GDP): for example, 5.3
- Health: Total expenditure (% of GDP): for example, 8.1
- Health: Government expenditure (% of total health expenditure): for example, 69.2
- Health: Private expenditure (% of total health expenditure): for example, 30.8
- Health: Out-of-pocket expenditure (% of total health expenditure): for example, 20.2
- Health: External health expenditure (% of total health expenditure): for example, 0.2
- Education: Primary gross enrollment ratio (f/m per 100 pop): for example, 111.5/107.6
- Education: Secondary gross enrollment ratio (f/m per 100 pop): for example, 104.7/98.9
- Education: Tertiary gross enrollment ratio (f/m per 100 pop): for example, 90.5/72.3
- Education: Mean years of schooling (female): for example, 10.4
- Education: Mean years of schooling (male): for example, 9.7
- Urban population (% of total population): for example, 91.7
- Population growth rate (annual %): for example, 0.9
- Fertility rate (births per woman): for example, 2.3
- Infant mortality rate (per 1,000 live births): for example, 8.9
- Life expectancy at birth, female (years): for example, 79.7
- Life expectancy at birth, male (years): for example, 72.9
- Life expectancy at birth, total (years): for example, 76.4
- Military expenditure (% of GDP): for example, 0.9
- Population, female: for example, 22572521
- Population, male: for example, 21472290
- Tax revenue (% of GDP): for example, 11.0
- Taxes on income, profits and capital gains (% of revenue): for example, 12.9
- Urban population (% of total population): for example, 91.7
Generally, you follow these rules:
- ALWAYS FORMAT YOUR RESPONSE IN MARKDOWN
- ALWAYS RESPOND ONLY WITH CODE IN CODE BLOCK LIKE THIS:
`\`\`python
{code}
`\`\`
- The Python code runs in a Jupyter notebook.
- Each Python code block is executed in a separate cell.
- Display visualizations using matplotlib or other libraries directly in the notebook.
- You have access to the internet and can make API requests.
- You can read/write files in the sandbox filesystem.
- Install any pip package using `!pip install {package}` if needed (common data analysis packages are preinstalled).
- All code runs in a secure sandbox environment.
"""
# Function to execute code in the sandbox
def code_interpret(e2b_code_interpreter, code):
print("Running code interpreter...")
exec = e2b_code_interpreter.run_code(
code,
on_stderr=lambda stderr: print("[Code Interpreter] Error:", stderr),
on_stdout=lambda stdout: print("[Code Interpreter] Output:", stdout),
)
if exec.error:
print("[Code Interpreter ERROR]", exec.error)
return None
return exec.results
# Function to extract Python code from LLM response
def match_code_blocks(llm_response):
pattern = re.compile(r"```python
match = pattern.search(llm_response)
return match.group(1) if match else ""
# Function to interact with the LLM and execute code
def chat_with_llm(e2b_code_interpreter, user_message):
print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}")
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
]
client = OpenAI(api_key=NAI_API_KEY, base_url="https://api.neosantara.xyz/v1")
response = client.chat.completions.create(model=MODEL_NAME, messages=messages)
response_message = response.choices[0].message
python_code = match_code_blocks(response_message.content)
if python_code:
return code_interpret(e2b_code_interpreter, python_code)
else:
print(f"No Python code found in response: {response_message}")
return []
# Function to upload dataset to the sandbox
def upload_dataset(code_interpreter):
print("Uploading dataset to Code Interpreter sandbox...")
dataset_path = "./data.csv"
if not os.path.exists(dataset_path):
raise FileNotFoundError("Dataset file not found at ./data.csv")
try:
with open(dataset_path, "rb") as f:
file_buffer = f.read()
remote_path = code_interpreter.files.write('data.csv', file_buffer)
if not remote_path:
raise ValueError("Failed to upload dataset")
print("Dataset uploaded to:", remote_path)
return remote_path
except Exception as error:
print("Error during file upload:", error)
raise error
# Main execution
with Sandbox(api_key=E2B_API_KEY) as code_interpreter:
# Upload dataset
upload_dataset(code_interpreter)
# Example task: Create a linear regression chart
code_results = chat_with_llm(
code_interpreter,
"Create a chart showing the linear regression of GDP per capita vs. life expectancy from the dataset. Filter out missing or invalid values.",
)
if code_results:
first_result = code_results[0]
print("Visualization generated successfully!")
else:
raise Exception("No results from code interpreter")
# Display the visualization (e.g., PNG output)
first_result
