#!/usr/bin/env python
# coding: utf-8

# # Jupyter AI
# 
# Jupyter AI is a JupyterLab extension that provides a friendly user interface between users and AI models.
# 
# This demo showcases the IPython magics Jupyter AI provides out-of-the-box.
# 
# Load the IPython extension:

# In[1]:


get_ipython().run_line_magic('load_ext', 'jupyter_ai')


# ## Getting started
# 
# To start using Jupyter AI, use the `%%ai` cell magic with the model specified via the syntax `<provider-id>:<model-id>`. Every line after the first should contain the prompt.
# 
# Currently, we support the following providers:
# 
# - `ai21`
# - `anthropic`
# - `cohere`
# - `huggingface_hub`
# - `openai`
# - `openai-chat`
# - `sagemaker-endpoint`
# 
# 
# These are LangChain LLM providers (`langchain.llms`). Any model ID that is valid for the corresponding provider is valid in the AI magics. For example, `claude-v1.2` is a valid model ID for the `anthropic` provider, and one can invoke it via Jupyter AI magics like so:

# In[2]:


get_ipython().run_cell_magic('ai', 'anthropic:claude-v1.2', 'Write a poem about C++.\n')


# ## Implicit providers
# 
# If your model ID is unique, then we can also infer the provider ID implicitly:

# In[3]:


get_ipython().run_cell_magic('ai', 'j2-jumbo-instruct # infers AI21 provider', 'Write some JavaScript code that prints "hello world" to the console.\n')


# ## Rapidly experimenting with different HF Hub models
# We can call models on HuggingFace Hub directly:

# In[4]:


get_ipython().run_cell_magic('ai', 'huggingface_hub:google/flan-t5-xl', 'What is the capital of New York state?\n')


# In[5]:


get_ipython().run_cell_magic('ai', 'huggingface_hub:gpt2', 'What is the square root of 2?\n')


# ## Formatting via `-f/--format`
# 
# You can also pass the `-f/--format` argument to specify an IPython display to use to render the output. Valid formats:
# 
# - `markdown`
# - `math`
# - `html`
# - `json`

# In[6]:


get_ipython().run_cell_magic('ai', 'anthropic:claude-v1.2 -f html', 'Create a square using SVG with a black border and white fill. \n')


# In[8]:


get_ipython().run_cell_magic('ai', 'chatgpt -f math', 'Generate the 2D heat equation.\n')


# ## IPython interpolation
# 
# You can also interpolate IPython scope into your prompt via curly braces.

# In[9]:


poet = "Walt Whitman"


# In[10]:


get_ipython().run_cell_magic('ai', 'chatgpt', 'Write a poem in the style of {poet}\n')


# In[11]:


for i in range(0, 5):
  print(i)


# In[12]:


get_ipython().run_cell_magic('ai', 'cohere:command-xlarge-nightly', 'Please explain the code below:\n--\n{In[11]}\n')


# # Full example: The Laplace Equation
# 
# Say I just got a homework assignment about "the Laplace equation" from my professor.
# 
# 1. What is the Laplace equation?
# 2. What is the form of the Laplace equation in polar coordinates?
# 3. Write a Python algorithm to solve the Laplace equation numerically.
# 
# The problem is: I have absolutely no idea what they're talking about.
# 
# Let's ask an LLM what the assignment is about.

# In[13]:


get_ipython().run_cell_magic('ai', 'claude-v1.2', 'What is the 2d Laplace equation and when is it used? Use LaTeX for equations delimited by `$`.\n')


# In[14]:


get_ipython().run_cell_magic('ai', 'j2-jumbo-instruct --format math', 'Write the 2d Laplace equation in polar coordinates.\n')


# In[15]:


get_ipython().run_cell_magic('ai', 'chatgpt', 'Please generate the Python code to solve the 2D Laplace equation in cartesian coordinates.\nSolve the equation on the square domain x=(0,1) and y=(0,1) with vanishing boundary conditions.\nPlot the solution using Matplotlib.\nPlease also provide an explanation.\n')