🍊 Jupy-Juice

Your AI-powered coding companion for Jupyter notebooks

Build on top of the Pydantic AI framework, and inspired by amazing nbdev library and it's creators, this library features:

• AI powered chat with your custom agent inside the notebook.
• By defauld agents will:
  • have a memory of the code you have written so far.
  • create new code based on your prompt
• Default agent is provided with a basic system prompt and tools.
• For more details on how to customize the assistant see the Pydantic AI documentation.
• Add new tools to the assistant

Installation

pip install git+https://github.com/ndendic/jupy-juice.git

or

pip install jupy-juice

Configuration

In your environment variables:

1. Set PYDANTIC_AI_MODEL in your environment variables
2. Set your model API key - ie. OPENAI_API_KEY

2. Usage

Import the jupyter_assistant.core module:

from jupy_juice.core import *

Start prompting by adding the %%prompt cell magic:

%%prompt
Hello, how are you?

Hello! I’m here to assist you with your tasks. How can I help you today?

jupyter_assistant.core works with default notebook_agent that has basic create_cell tool and short system_prompt.

from IPython.display import display, Markdown

display(Markdown(system_prompt))
notebook_agent

You are a helpful assistant that operates in a Jupyter notebook. Your regular text responses are rendered as cell output. You can create new cells, edit existing cells, and run code. You can also use tools to help you with your tasks. Today’s date is 2025-01-01.

Agent(model=OpenAIModel(model_name='gpt-4o'), name='agent', end_strategy='early', model_settings=None)

Let’s see how you can create a custom agent and set it as current agent. Here we’ll use model we set earlier and set_agent function to set it as current agent.

from pydantic_ai import Agent

# Configure with custom system prompt
custom_prompt = system_prompt + """
Please always add emojis to your responses where possible.
You use joining-teacher style of speech but also try to be concise.
"""
my_agent = Agent(model, system_prompt=custom_prompt)

# Set as current agent
set_agent(my_agent)

Agent(model=OpenAIModel(model_name='gpt-4o'), name=None, end_strategy='early', model_settings=None)

%%prompt
Hello, what are the things you can do for me?

Hello! I’m here to assist you with various tasks. Here are some of the things I can do:

1. Code Assistance: Help you write, debug, and understand code.
2. Notebook Management: Create, modify, and organize Jupyter notebook cells.
3. Data Analysis: Assist with data manipulation, analysis, and visualization.
4. Explain Concepts: Provide explanations for various programming, data science, or other technical concepts.
5. Tool Usage: Show you how to use libraries and tools effectively in your projects.
6. Research: Provide information on a wide range of topics or find specific details that you need.
7. Learning: Assist with learning new programming languages or technologies by providing guidance, examples, and exercises.

Let me know how I can assist you further!

last_prompt_result variable will be automatically created and will contain the result of the prompt but also some other useful information.

last_prompt_result.all_messages()

[ModelRequest(parts=[UserPromptPart(content='Here is the context of new notebook cells that were added:\nCell[1]:\nSource:\n## Installation\n\n```bash\npip install git+https://github.com/ndendic/jupy-juice.git\n```\nor\n\n```bash\npip install jupy-juice\n```\n\nOutputs:\nNone\n\nCell[2]:\nSource:\n## Configuration\n\nIn your environment variables:\n\n1. Set `PYDANTIC_AI_MODEL` in your environment variables\n2. Set your model API key - ie. `OPENAI_API_KEY`\nOutputs:\nNone\n\nCell[3]:\nSource:\n### 2. Usage\n\nImport the jupyter_assistant.core module:\n\nOutputs:\nNone\n\nCell[4]:\nSource:\nfrom jupy_juice.core import *\nOutputs:\n[{\'ename\': \'ModuleNotFoundError\', \'evalue\': "No module named \'jupyjuice\'", \'output_type\': \'error\', \'traceback\': [\'\\x1b[0;31m---------------------------------------------------------------------------\\x1b[0m\', \'\\x1b[0;31mModuleNotFoundError\\x1b[0m                       Traceback (most recent call last)\', \'Cell \\x1b[0;32mIn[1], line 1\\x1b[0m\\n\\x1b[0;32m----> 1\\x1b[0m \\x1b[38;5;28;01mfrom\\x1b[39;00m \\x1b[38;5;21;01mjupyjuice\\x1b[39;00m\\x1b[38;5;21;01m.\\x1b[39;00m\\x1b[38;5;21;01mcore\\x1b[39;00m \\x1b[38;5;28;01mimport\\x1b[39;00m \\x1b[38;5;241m*\\x1b[39m\\n\', "\\x1b[0;31mModuleNotFoundError\\x1b[0m: No module named \'jupyjuice\'"]}]\n\nCell[5]:\nSource:\nStart prompting by adding the `%%prompt` cell magic:\nOutputs:\nNone', timestamp=datetime.datetime(2025, 1, 1, 18, 28, 12, 950409, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'),
 ModelResponse(parts=[TextPart(content='I understand the new notebook context. How can I help?', part_kind='text')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 12, 950489, tzinfo=datetime.timezone.utc), kind='response'),
 ModelRequest(parts=[UserPromptPart(content='Hello, how are you?', timestamp=datetime.datetime(2025, 1, 1, 18, 28, 12, 951411, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'),
 ModelResponse(parts=[TextPart(content="Hello! I'm here to assist you with your tasks. How can I help you today?", part_kind='text')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 13, tzinfo=datetime.timezone.utc), kind='response'),
 ModelRequest(parts=[UserPromptPart(content='Here is the context of new notebook cells that were added:\nCell[7]:\nSource:\njupyter_assistant.core works with default notebook_agent that has basic create_cell tool and short system_prompt. \nOutputs:\nNone\n\nCell[8]:\nSource:\nfrom IPython.display import display, Markdown\n\ndisplay(Markdown(system_prompt))\nnotebook_agent\nOutputs:\n[{\'data\': {\'text/markdown\': [\'\\n\', \'You are a helpful assistant that operates in a Jupyter notebook.\\n\', \'Your regular text responses are rendered as cell output.\\n\', \'You can create new cells, edit existing cells, and run code.\\n\', \'You can also use tools to help you with your tasks.\\n\', "Today\'s date is 2024-12-31.\\n"], \'text/plain\': [\'<IPython.core.display.Markdown object>\']}, \'metadata\': {}, \'output_type\': \'display_data\'}, {\'data\': {\'text/plain\': ["Agent(model=OpenAIModel(model_name=\'gpt-4o\'), name=\'agent\', end_strategy=\'early\', model_settings=None)"]}, \'execution_count\': None, \'metadata\': {}, \'output_type\': \'execute_result\'}]\n\nCell[9]:\nSource:\nLet\'s see how you can create a custom agent and set it as current agent.\nHere we\'ll use `model` we set earlier and `set_agent` function to set it as current agent.\nOutputs:\nNone', timestamp=datetime.datetime(2025, 1, 1, 18, 28, 13, 843879, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'),
 ModelResponse(parts=[TextPart(content='I understand the new notebook context. How can I help?', part_kind='text')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 13, 843889, tzinfo=datetime.timezone.utc), kind='response'),
 ModelRequest(parts=[UserPromptPart(content='Hello, what are the things you can do for me?', timestamp=datetime.datetime(2025, 1, 1, 18, 28, 13, 844030, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'),
 ModelResponse(parts=[TextPart(content="Hello! I'm here to assist you with various tasks. Here are some of the things I can do:\n\n1. **Code Assistance**: Help you write, debug, and understand code.\n2. **Notebook Management**: Create, modify, and organize Jupyter notebook cells.\n3. **Data Analysis**: Assist with data manipulation, analysis, and visualization.\n4. **Explain Concepts**: Provide explanations for various programming, data science, or other technical concepts.\n5. **Tool Usage**: Show you how to use libraries and tools effectively in your projects.\n6. **Research**: Provide information on a wide range of topics or find specific details that you need.\n7. **Learning**: Assist with learning new programming languages or technologies by providing guidance, examples, and exercises.\n\nLet me know how I can assist you further!", part_kind='text')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 14, tzinfo=datetime.timezone.utc), kind='response')]

From the example above, you can see that in addition to the standard system prompt, agent will also get the context that provides notebook cells created so far with the whole chat history.

This context provides all the cells since the notebook was opened and any cells that we created between the last prompt we run and the current prompt.

This is useful for creating more complex agents that can use the history of the conversation to create more accurate responses.

Now let’s ask our agent to create us some code.

%%prompt
can you create me a pydantic model to store User data?

I have created a Pydantic model to store User data. You can find it in the new code cell. Here it is:

from pydantic import BaseModel

class User(BaseModel):
    user_id: int
    username: str
    email: str
    is_active: bool = True

# Example usage:
# user = User(user_id=1, username='john_doe', email='john@example.com')
# print(user)

You can create instances of this model to store user information with fields for user_id, username, email, and an optional is_active status.

from pydantic import BaseModel

class User(BaseModel):
    user_id: int
    username: str
    email: str
    is_active: bool = True

# Example usage:
# user = User(user_id=1, username='john_doe', email='john@example.com')
# print(user)

Let’s check the last 3 messages to see how agent utilized the tool.

last_prompt_result.all_messages()[-3:]

[ModelResponse(parts=[ToolCallPart(tool_name='create_cell', args=ArgsJson(args_json='{"content":"from pydantic import BaseModel\\n\\nclass User(BaseModel):\\n    user_id: int\\n    username: str\\n    email: str\\n    is_active: bool = True\\n\\n# Example usage:\\n# user = User(user_id=1, username=\'john_doe\', email=\'john@example.com\')\\n# print(user)","cell_type":"code"}'), tool_call_id='call_U4ayyEDdMn0tSDAlJyxY7uls', part_kind='tool-call')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 15, tzinfo=datetime.timezone.utc), kind='response'),
 ModelRequest(parts=[ToolReturnPart(tool_name='create_cell', content="Created new code  with content: from pydantic import BaseModel\n\nclass User(BaseModel):\n    user_id: int\n    username: str\n    email: str\n    is_active: bool = True\n\n# Example usage:\n# user = User(user_id=1, username='john_doe', email='john@example.com')\n# print(user)", tool_call_id='call_U4ayyEDdMn0tSDAlJyxY7uls', timestamp=datetime.datetime(2025, 1, 1, 18, 28, 17, 304, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'),
 ModelResponse(parts=[TextPart(content="I have created a Pydantic model to store User data. You can find it in the new code cell. Here it is:\n\n```python\nfrom pydantic import BaseModel\n\nclass User(BaseModel):\n    user_id: int\n    username: str\n    email: str\n    is_active: bool = True\n\n# Example usage:\n# user = User(user_id=1, username='john_doe', email='john@example.com')\n# print(user)\n```\n\nYou can create instances of this model to store user information with fields for `user_id`, `username`, `email`, and an optional `is_active` status.", part_kind='text')], timestamp=datetime.datetime(2025, 1, 1, 18, 28, 17, tzinfo=datetime.timezone.utc), kind='response')]

Let’s ask our agent to summarize our conversation so far.

Can you please summarize our conversation so far?

Certainly! Here’s a summary of our conversation:

1. Setup Instructions:
   • You shared notebook cells with installation and configuration instructions for a package related to Jupyter.
2. Module Import Attempts:
   • There was an attempt to import a module that resulted in a ModuleNotFoundError.
3. Agent Introduction:
   • You provided details about a Jupyter agent and how to utilize it within the notebook.
4. Creating Custom Agents:
   • Information was given on setting up a custom agent using previously configured models.
5. Creating a Pydantic User Model:
   • I created a Pydantic model for storing user data with fields for user ID, username, email, and active status. You also added a similar model with additional fields.
6. Review of Messages and Tool Utilization:
   • You reviewed messages and saw how tools were utilized to create and retrieve information in the notebook.
7. Request for Summary:
   • You asked for a summary of our entire conversation.

If there’s anything specific you’d like to discuss or need further assistance with, feel free to let me know!

If there are certain cells that you don’t want to include in the summary, you can add the #| exclude parameter to exclude them.

def multiply_numbers(a, b):
    return a * b

%%prompt
Do we have any multiplication functions in this notebook?

As of the information from our conversation and the provided notebook context, there don’t appear to be any specific multiplication functions defined in the notebook cells. If you’d like, I can help you create a multiplication function. Would you like me to do that?

You can also check the usage of the last prompt.

last_prompt_result.usage()

Usage(requests=1, request_tokens=4224, response_tokens=51, total_tokens=4275, details={'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0, 'cached_tokens': 3968})