Jupyternaut V3: Agentic tool-calling #1479

dlqqq · 2025-09-08T17:47:42Z

Description

This PR adds tool-calling capabilities to Jupyternaut. See the "code changes" section on implementation details. Key points:

Tool calls & their outputs are rendered using a new <jai-tool-call> custom web component. This leverages the custom elements API available in browsers, which lets us define custom HTML elements.
- To help with this, I've added r2wc, a new NPM dependency that stands for "React to Web Component", which provides a simple utility to define a custom element's rendering lifecycle as a React component.
BasePersona.stream_message() has been removed. This was previously the main method used by personas to write back to the YChat.
All of the streaming, output rendering, and agentic tool-calling logic has been unified under the default_flow module, which uses pocketflow to define all of the "core logic" as an agent flow (a graph of nodes that can call each other).
pocketflow has been added as a new PyPI dependency.
Defines a new default toolkit under jupyter_ai.tools.default_toolkit. This is a very minimal "starter set" of tools which may be upgraded/superseded in a future release.

Demo

Screen.Recording.2025-09-08.at.10.51.52.AM.mov

Code changes

Frontend

Adds the r2wc NPM dependency.
Adds a new web-components labextension plugin that provides the <jai-tool-call> custom element, defined in a new JaiToolCall React component.
Each <JaiToolCall /> component takes the following props:

type JaiToolCallProps = {
  id: string;
  type: string;
  function: string;
  index: number;
  output?: {
    tool_call_id: string;
    role: string;
    name: string;
    content: string | null;
  };
};

When a message contains tool calls, each tool call is rendered in one of these elements with output being unset. This shows the tool name & a loading spinner.
After Jupyternaut runs each tool, the output property is set, which hides the loading spinner and shows a green check mark.

Backend

Adds a jupyter_ai.litellm_lib module:
- Provides a ToolCallList class that aggregates the tool calls specified in each stream chunk. It provides a render() method that can be used to convert the list into a string of <jai-tool-call> elements so its consumers can easily render tool calls in real time.
- Provides a run_tools() function which accepts a ToolCallList & a toolkit.
Adds a new jupyter_ai.default_flow module:
- This new module now defines all of the logic that is run when a user invokes an LLM by @-mentioning it. This includes chat streaming, output rendering, tool-calling capabilities, and agentic recursion.
- This replaces BasePersona.stream_message().
- The main export of this module is run_default_flow(), which accepts a dictionary of parameters that govern how the flow is run. Custom personas can tweak our default flow without re-inventing the wheel by changing the parameters.
Removes BasePersona.stream_message().
Edits Jupyternaut.process_message() to just call await run_default_flow() to have everything handled automatically.
Implements a default toolkit under jupyter_ai.tools.default_toolkit. This is what Jupyternaut will use for now.

User-facing changes

Jupyternaut is now way more useful.

Backwards-incompatible changes

BasePersona.stream_message() has been removed.

Follow-up items

Error handling in response - Jupyternaut just halts whenever this occurs, with no message surfaced to the user.
Tool call success/failure state - right now each tool call is either "loading" or "success". We should allow it to show a red X if a tool call raised an exception.
Can't open the tool element while the AI response is streaming.

dlqqq · 2025-09-10T17:19:34Z

Hey folks, I will be AFK from tomorrow to Monday. Feel free to review & test, but note that I may not be able to reply until Monday. Thanks!

dlqqq · 2025-09-18T17:36:52Z

I have completed refactoring this branch using pocketflow. Need to address some minor regressions as well & this PR will be ready for review.

dlqqq · 2025-09-18T18:09:55Z

@Zsailer Hey Zach! This PR is mostly complete right now. I've unified all of the agent logic under a new "agent flow" that uses the pocketflow library.

Brian would love it if you can give this PR a review. He had mentioned that you were meeting w/ the Anthropic team about Jupyter AI, and seeing how outputs are rendered here might be helpful.

3coins · 2025-09-18T19:10:42Z

packages/jupyter-ai/jupyter_ai/default_flow/default_flow.py

+    It should take `content: str` and `tool_call_ui_elements: str` as format arguments.
+    """
+
+    toolkit: Toolkit | None


Is this too limiting? For example, there could be tools from multiple toolkits that I might end up using.

3coins · 2025-09-18T19:15:33Z

packages/jupyter-ai/jupyter_ai/litellm_lib/toolcall_list.py

+    """
+    Name of the tool function to be called.
+
+    TODO: Check if this attribute is defined for non-function tools, e.g. tools


https://modelcontextprotocol.info/docs/concepts/tools/

3coins · 2025-09-18T20:30:05Z

@dlqqq
I am getting this error when running Jupyternaut. Are we missing any dependency?

[E 2025-09-18 13:21:19.490 AiExtension] Exception occurred while running default agent flow:
    Traceback (most recent call last):
      File "/Users/pijain/projects/jupyter-ai/packages/jupyter-ai/jupyter_ai/default_flow/default_flow.py", line 329, in run_default_flow
        await flow.run_async({})
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 72, in run_async
        return await self._run_async(shared)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 87, in _run_async
        async def _run_async(self,shared): p=await self.prep_async(shared); o=await self._orch_async(shared); return await self.post_async(shared,p,o)
                                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 85, in _orch_async
        while curr: curr.set_params(p); last_action=await curr._run_async(shared) if isinstance(curr,AsyncNode) else curr._run(shared); curr=copy.copy(self.get_next_node(curr,last_action))
                                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 73, in _run_async
        async def _run_async(self,shared): p=await self.prep_async(shared); e=await self._exec(p); return await self.post_async(shared,p,e)
                                                                              ^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 68, in _exec
        if i==self.max_retries-1: return await self.exec_fallback_async(prep_res,e)
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 62, in exec_fallback_async
        async def exec_fallback_async(self,prep_res,exc): raise exc
                                                          ^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/pocketflow/__init__.py", line 66, in _exec
        try: return await self.exec_async(prep_res)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/projects/jupyter-ai/packages/jupyter-ai/jupyter_ai/default_flow/default_flow.py", line 171, in exec_async
        tools=self.toolkit.to_json(),
              ^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/projects/jupyter-ai/packages/jupyter-ai/jupyter_ai/tools/models.py", line 233, in to_json
        "function": function_to_dict(tool.callable),
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/litellm/utils.py", line 5108, in function_to_dict
        raise e
      File "/Users/pijain/miniforge3/envs/jaidev/lib/python3.11/site-packages/litellm/utils.py", line 5106, in function_to_dict
        from numpydoc.docscrape import NumpyDocString
    ModuleNotFoundError: No module named 'numpydoc'

I used the dev-install script and have litellm v1.77.1 installed. numpydoc seems to be an optional dependency in litellm.
https://github.com/BerriAI/litellm/blob/main/pyproject.toml#L35

3coins · 2025-09-18T20:39:58Z

The default toolkit might need some tweaking or tool calling is not working as expected, I couldn't get it to read a file in the same dir as chat.

Server logs do show calling the tool node:

[I 2025-09-18 13:31:50.300 AiExtension] Running RootNode.post_async()
[I 2025-09-18 13:31:50.300 AiExtension] Running ToolExecutorNode.prep_async()
[I 2025-09-18 13:31:50.300 AiExtension] Running ToolExecutorNode.exec_async()
[I 2025-09-18 13:31:50.300 AiExtension] Running ToolExecutorNode.post_async()
[I 2025-09-18 13:31:50.301 AiExtension] Running RootNode.exec_async()
[I 2025-09-18 13:31:51.307 ServerApp] Saving the content from room text:chat:61726951-270e-41f9-afc6-4f01b6404258
[I 2025-09-18 13:31:51.308 YDocExtension] Saving file: agent-test.chat
[I 2025-09-18 13:31:52.424 ServerApp] Saving the content from room text:chat:61726951-270e-41f9-afc6-4f01b6404258
[I 2025-09-18 13:31:52.428 YDocExtension] Saving file: agent-test.chat
[I 2025-09-18 13:31:52.606 AiExtension] Running RootNode.post_async()
[I 2025-09-18 13:31:52.607 AiExtension] Running ToolExecutorNode.prep_async()
[I 2025-09-18 13:31:52.607 AiExtension] Running ToolExecutorNode.exec_async()
[I 2025-09-18 13:31:52.622 AiExtension] Running ToolExecutorNode.post_async()

dlqqq · 2025-09-22T17:28:57Z

@3coins Thanks for the review!

Yes, I forgot that numpydoc is required for one of the litellm.utils functions to work (specifically the one that turns a function w/ docstring => tool dictionary). I'll add that to the dependencies.

I've also noticed that the tools don't always work. I'm having a hard time distinguishing between the LLM calling the tool incorrectly v.s. the tool not working at all. I think we can address this by adding more logging + improving the system prompt.

Finally, make sure to run jlpm && jlpm build to see the new tool call component. It should appear in the chat even if the tool call failed.

theo-yoon · 2025-09-29T07:01:09Z

packages/jupyter-ai/src/web-components/jai-tool-call.tsx

+  if (!props.id || !props.type || !props.function_name) {
+    return null;
+  }


Hello. Overall, everything seems to be working well for me.

However, I noticed that the id prop is undefined in my environment when using the JaiToolCall component. Interestingly, it functions correctly after renaming all instances of id to tool_id in the files toolcall_list.py, web-components-plugin.ts, and jai-tool-call.tsx.

Am I understanding this correctly?

dlqqq force-pushed the make-jupyternaut-agentic branch from 4e4a3ae to f09caf5 Compare September 8, 2025 17:48

dlqqq added the enhancement New feature or request label Sep 8, 2025

pr-triage-board bot added this to PR triage (experimental) Sep 8, 2025

dlqqq added 8 commits September 15, 2025 10:10

add 'litellm_utils' module

3165e72

WIP: first working copy of Jupyternaut as an agent

d504b34

clean up tool calling flow & show in chat

5aa46bf

add temporary default toolkit

7ba285d

update tool calling APIs

4059489

improve bash tool reliability, drop jupyter_ai_tools for now

056676e

add jai-tool-call web component to show tool calls & outputs

1311989

move litellm_utils => litellm_lib

5223fa6

dlqqq force-pushed the make-jupyternaut-agentic branch from f09caf5 to 5223fa6 Compare September 15, 2025 17:11

dlqqq added 3 commits September 18, 2025 10:22

migrate all agent logic to module using pocketflow

ea8a92a

add pocketflow as a dependency

fd950bd

update jai-tool-call web component API to do less JSON parsing

793fb4e

dlqqq added 2 commits September 18, 2025 10:41

remove debug logs

29f075d

show writing indicator while processing request

95d2615

dlqqq marked this pull request as ready for review September 18, 2025 18:08

3coins reviewed Sep 18, 2025

View reviewed changes

theo-yoon reviewed Sep 29, 2025

View reviewed changes

joadoumie mentioned this pull request Sep 30, 2025

Add a new text command for a more ergonomic way to attach text to the chat interface #1476

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Jupyternaut V3: Agentic tool-calling #1479

Jupyternaut V3: Agentic tool-calling #1479

Uh oh!

dlqqq commented Sep 8, 2025 •

edited

Loading

Uh oh!

dlqqq commented Sep 10, 2025

Uh oh!

dlqqq commented Sep 18, 2025

Uh oh!

dlqqq commented Sep 18, 2025

Uh oh!

3coins Sep 18, 2025

Uh oh!

3coins Sep 18, 2025

Uh oh!

3coins commented Sep 18, 2025 •

edited

Loading

Uh oh!

3coins commented Sep 18, 2025

Uh oh!

dlqqq commented Sep 22, 2025

Uh oh!

theo-yoon Sep 29, 2025

Uh oh!

Uh oh!

Uh oh!

Jupyternaut V3: Agentic tool-calling #1479

Are you sure you want to change the base?

Jupyternaut V3: Agentic tool-calling #1479

Uh oh!

Conversation

dlqqq commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Demo

Code changes

Frontend

Backend

User-facing changes

Backwards-incompatible changes

Follow-up items

Uh oh!

dlqqq commented Sep 10, 2025

Uh oh!

dlqqq commented Sep 18, 2025

Uh oh!

dlqqq commented Sep 18, 2025

Uh oh!

3coins Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

3coins Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

3coins commented Sep 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

3coins commented Sep 18, 2025

Uh oh!

dlqqq commented Sep 22, 2025

Uh oh!

theo-yoon Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dlqqq commented Sep 8, 2025 •

edited

Loading

3coins commented Sep 18, 2025 •

edited

Loading