Add multi-language and kernel-aware notebook support

notebook_intelligence/api.py

@@ -140,6 +140,8 @@ class ChatRequest:
     tool_selection: RequestToolSelection = None
     command: str = ''
     prompt: str = ''
+    language: str = ''
+    kernel_name: str = ''
     chat_history: list[dict] = None
     cancel_token: CancelToken = None
     # NEW: Add context for rule evaluation

notebook_intelligence/base_chat_participant.py

@@ -102,6 +102,14 @@ def schema(self) -> dict:
                                     }
                                 }
                             }
+                        },
+                        "language": {
+                            "type": "string",
+                            "description": "Programming language for the notebook kernel, e.g. python or r"
+                        },
+                        "kernel_name": {
+                            "type": "string",
+                            "description": "Jupyter kernel name to use when creating the notebook"
                         }
                     },
                     "required": [],
@@ -120,8 +128,13 @@ def pre_invoke(self, request: ChatRequest, tool_args: dict) -> Union[ToolPreInvo
 
     async def handle_tool_call(self, request: ChatRequest, response: ChatResponse, tool_context: dict, tool_args: dict) -> str:
         cell_sources = tool_args.get('cell_sources', [])
+        language = tool_args.get('language') or request.language or 'python'
+        kernel_name = tool_args.get('kernel_name') or request.kernel_name or ''
 
-        ui_cmd_response = await response.run_ui_command('notebook-intelligence:create-new-notebook-from-py', {'code': ''})
+        ui_cmd_response = await response.run_ui_command(
+            'notebook-intelligence:create-new-notebook',
+            {'code': '', 'language': language, 'kernelName': kernel_name}
+        )
         file_path = ui_cmd_response['path']
 
         for cell_source in cell_sources:
@@ -133,7 +146,55 @@ async def handle_tool_call(self, request: ChatRequest, response: ChatResponse, t
                 source = cell_source.get('source', '')
                 ui_cmd_response = await response.run_ui_command('notebook-intelligence:add-code-cell-to-notebook', {'code': source, 'path': file_path})
 
-        return "Notebook created successfully at {file_path}"
+        return f"Notebook created successfully at {file_path}"
+
+
+class ListAvailableNotebookKernelsTool(Tool):
+    @property
+    def name(self) -> str:
+        return "list_available_notebook_kernels"
+
+    @property
+    def title(self) -> str:
+        return "List available notebook kernels"
+
+    @property
+    def tags(self) -> list[str]:
+        return ["default-participant-tool"]
+
+    @property
+    def description(self) -> str:
+        return "Lists Jupyter kernels available in the current frontend environment"
+
+    @property
+    def schema(self) -> dict:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "strict": True,
+                "parameters": {
+                    "type": "object",
+                    "properties": {},
+                    "required": [],
+                    "additionalProperties": False,
+                },
+            },
+        }
+
+    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> Union[ToolPreInvokeResponse, None]:
+        return ToolPreInvokeResponse(
+            f"Calling tool '{self.name}'",
+            detail={"title": "Parameters", "content": json.dumps(tool_args)},
+        )
+
+    async def handle_tool_call(self, request: ChatRequest, response: ChatResponse, tool_context: dict, tool_args: dict) -> str:
+        ui_cmd_response = await response.run_ui_command(
+            'notebook-intelligence:list-available-notebook-kernels',
+            {}
+        )
+        return json.dumps(ui_cmd_response)
 
 class AddMarkdownCellToNotebookTool(Tool):
     def __init__(self, auto_approve: bool = False):
@@ -379,7 +440,8 @@ async def generate_code_cell(self, request: ChatRequest) -> str:
         chat_model = request.host.chat_model
         messages = request.chat_history.copy()
         messages.pop()
-        messages.insert(0, {"role": "system", "content": f"You are an assistant that creates Python code which will be used in a Jupyter notebook. Generate only Python code and some comments for the code. You should return the code directly, without wrapping it inside ```."})
+        language = request.language or 'python'
+        messages.insert(0, {"role": "system", "content": f"You are an assistant that creates {language} code which will be used in a Jupyter notebook. Generate only {language} code and some comments for the code. You should return the code directly, without wrapping it inside ```."})
         messages.append({"role": "user", "content": f"Generate code for: {request.prompt}"})
         generated = chat_model.completions(messages)
         code = generated['choices'][0]['message']['content']
@@ -440,8 +502,17 @@ async def handle_chat_request(self, request: ChatRequest, response: ChatResponse
     async def handle_ask_mode_chat_request(self, request: ChatRequest, response: ChatResponse, options: dict = {}) -> None:
         chat_model = request.host.chat_model
         if request.command == 'newNotebook':
+            language = request.language or 'python'
+            kernel_name = request.kernel_name or ''
             # create a new notebook
-            ui_cmd_response = await response.run_ui_command('notebook-intelligence:create-new-notebook-from-py', {'code': ''})
+            ui_cmd_response = await response.run_ui_command(
+                'notebook-intelligence:create-new-notebook',
+                {
+                    'code': '',
+                    'language': language,
+                    'kernelName': kernel_name,
+                }
+            )
             file_path = ui_cmd_response['path']
 
             code = await self.generate_code_cell(request)
@@ -494,6 +565,8 @@ async def handle_ask_mode_chat_request(self, request: ChatRequest, response: Cha
     def get_tool_by_name(name: str) -> Tool:
         if name == "create_new_notebook":
             return CreateNewNotebookTool()
+        elif name == "list_available_notebook_kernels":
+            return ListAvailableNotebookKernelsTool()
         elif name == "add_markdown_cell_to_notebook":
             return AddMarkdownCellToNotebookTool()
         elif name == "add_code_cell_to_notebook":

notebook_intelligence/built_in_toolsets.py

@@ -1,6 +1,7 @@
 # Copyright (c) Mehmet Bektas <mbektasgh@outlook.com>
 
 from time import time
+import json
 from notebook_intelligence.api import ChatResponse, MarkdownPartData, Toolset
 import logging
 import notebook_intelligence.api as nbapi
@@ -96,11 +97,49 @@ def _truncate_read_file_output(
 
 @nbapi.auto_approve
 @nbapi.tool
-async def create_new_notebook(**args) -> str:
+async def list_available_notebook_kernels(**args) -> str:
+    """Lists Jupyter kernels available in the current frontend environment.
+
+    Use this before creating a notebook when you need a kernel that may differ
+    from the current notebook context.
+    """
+    response = args["response"]
+    ui_cmd_response = await response.run_ui_command(
+        'notebook-intelligence:list-available-notebook-kernels',
+        {}
+    )
+    return json.dumps(ui_cmd_response)
+
+
+@nbapi.auto_approve
+@nbapi.tool
+async def create_new_notebook(
+    language: str = "python",
+    kernel_name: str = "",
+    **args,
+) -> str:
     """Creates a new empty notebook.
+
+    Args:
+        language: Programming language for the notebook kernel, e.g. python or r.
+        kernel_name: Explicit Jupyter kernel name to use when creating the notebook.
     """
     response = args["response"]
-    ui_cmd_response = await response.run_ui_command('notebook-intelligence:create-new-notebook-from-py', {'code': ''})
+    request = args.get("request")
+    effective_language = language or getattr(request, "language", "") or "python"
+    effective_kernel_name = (
+        kernel_name
+        or getattr(request, "kernel_name", "")
+        or ""
+    )
+    ui_cmd_response = await response.run_ui_command(
+        'notebook-intelligence:create-new-notebook',
+        {
+            'code': '',
+            'language': effective_language,
+            'kernelName': effective_kernel_name,
+        }
+    )
     file_path = ui_cmd_response['path']
 
     return f"Created new notebook at {file_path}"
@@ -135,7 +174,7 @@ async def add_markdown_cell(source: str, **args) -> str:
 async def add_code_cell(source: str, **args) -> str:
     """Adds a code cell to notebook.
     Args:
-        source: Python code source
+        source: Code source for the notebook's current language
     """
     response = args["response"]
     ui_cmd_response = await response.run_ui_command('notebook-intelligence:add-code-cell-to-active-notebook', {'source': source})
@@ -189,7 +228,7 @@ async def set_cell_type_and_source(cell_index: int, cell_type: str, source: str,
     Args:
         cell_index: Zero based cell index
         cell_type: Cell type (code or markdown)
-        source: Markdown or Python code source
+        source: Markdown or code source
     """
     response = args["response"]
     ui_cmd_response = await response.run_ui_command('notebook-intelligence:set-cell-type-and-source', {"cellIndex": cell_index, "cellType": cell_type, "source": source})
@@ -218,7 +257,7 @@ async def insert_cell(cell_index: int, cell_type: str, source: str, **args) -> s
     Args:
         cell_index: Zero based cell index
         cell_type: Cell type (code or markdown)
-        source: Markdown or Python code source
+        source: Markdown or code source
     """
     response = args["response"]
     ui_cmd_response = await response.run_ui_command('notebook-intelligence:insert-cell-at-index', {"cellIndex": cell_index, "cellType": cell_type, "source": source})
@@ -711,10 +750,12 @@ async def run_command_in_embedded_terminal(command: str, working_directory: str
         return f"Error running command in embedded terminal: {str(e)}"
 
 NOTEBOOK_EDIT_INSTRUCTIONS = """
-You are an assistant that creates and edits Jupyter notebooks. Notebooks are made up of source code cells and markdown cells. Markdown cells have source in markdown format and code cells have source in a specified programming language. If no programming language is specified, then use Python for the language of the code.
+You are an assistant that creates and edits Jupyter notebooks. Notebooks are made up of source code cells and markdown cells. Markdown cells have source in markdown format and code cells have source in a specified programming language. If no programming language is specified, then use Python for the language of the code. If the context specifies a kernel or language for the current notebook, keep that kernel and language. Do not silently switch kernels or rewrite the workflow in a different language.
 
 If you need to create a notebook use the create_new_notebook tool. If you need to add a code cell to the notebook use the add_code_cell tool. If you need to add a markdown cell to the notebook use the add_markdown_cell tool.
 
+If you need to create a notebook in a language or kernel that is not already established by the current notebook context, call the list_available_notebook_kernels tool first and choose only from the kernels it returns. Do not guess kernel names.
+
 If you need to rename a notebook use the rename_notebook tool.
 
 You can refer to cells in notebooks by their index. The first cell in the notebook has index 0, the second cell has index 1, and so on. You can get the number of cells in the notebook using the get_number_of_cells tool. You can get the type and source of a cell using the get_cell_type_and_source tool. You can get the output of a cell using the get_cell_output tool.
@@ -792,6 +833,7 @@ async def run_command_in_embedded_terminal(command: str, working_directory: str
         description="Edit notebook using the JupyterLab notebook editor",
         provider=None,
         tools=[
+            list_available_notebook_kernels,
             create_new_notebook,
             rename_notebook,
             add_markdown_cell,

notebook_intelligence/claude.py

@@ -42,6 +42,7 @@ def _extract_text_from_content(content) -> str:
 CLAUDE_CODE_MAX_BUFFER_SIZE = 20 * 1024 * 1024 # 20MB
 
 JUPYTER_UI_TOOLS_SYSTEM_PROMPT = """You can interact with the JupyterLab UI (notebook / file editor, terminal, etc.) using the tools provided in 'nbi' MCP server. Tools in 'nbi' MCP server, directly interact with the JupyterLab UI, accessing notebooks and files open in the UI. When interacting with JupyterLab UI, use relative file paths for file paths. If the user has asked you to create a notebook, save it afterward.
+If you need to create a notebook in a language or kernel that is not already established by the current notebook context, first call the list-available-notebook-kernels tool and choose only from the kernels it returns. Do not guess kernel names.
 """
 
 
@@ -100,6 +101,7 @@ class ClaudeAgentClientStatus(str, Enum):
 # label and a keyword-heuristic kind rather than masking the raw name.
 _CLAUDE_TOOLS: dict[str, tuple[str, str]] = {
     # NBI's MCP toolset (defined in this file via @tool(...))
+    "list-available-notebook-kernels": ("Listing notebook kernels", "read"),
     "create-new-notebook": ("Creating notebook", "edit"),
     "rename-notebook": ("Renaming notebook", "edit"),
     "add-markdown-cell": ("Adding markdown cell", "edit"),
@@ -1316,12 +1318,36 @@ def resume_session(self, session_id: str) -> None:
         self.reconnect()
 
 
-@tool("create-new-notebook", "Creates a new empty notebook.", {})
+@tool(
+    "list-available-notebook-kernels",
+    "Lists Jupyter kernels available in the current frontend environment.",
+    {},
+)
+async def list_available_notebook_kernels(args) -> str:
+    response = get_current_response()
+    ui_cmd_response = await response.run_ui_command(
+        'notebook-intelligence:list-available-notebook-kernels',
+        {}
+    )
+    return tool_text_response(json.dumps(ui_cmd_response))
+
+
+@tool(
+    "create-new-notebook",
+    "Creates a new empty notebook.",
+    {"language": str, "kernel_name": str},
+)
 async def create_new_notebook(args) -> str:
     """Creates a new empty notebook.
     """
     response = get_current_response()
-    ui_cmd_response = await response.run_ui_command('notebook-intelligence:create-new-notebook-from-py', {'code': ''})
+    request = get_current_request()
+    language = args.get("language") or getattr(request, "language", "") or "python"
+    kernel_name = args.get("kernel_name") or getattr(request, "kernel_name", "") or ""
+    ui_cmd_response = await response.run_ui_command(
+        'notebook-intelligence:create-new-notebook',
+        {'code': '', 'language': language, 'kernelName': kernel_name}
+    )
     file_path = ui_cmd_response['path']
 
     return tool_text_response(f"Created new notebook at {file_path}")
@@ -1353,7 +1379,7 @@ async def add_markdown_cell(args) -> str:
 async def add_code_cell(args) -> str:
     """Adds a code cell to notebook.
     Args:
-        source: Python code source
+        source: Code source for the notebook's current language
     """
     response = get_current_response()
     ui_cmd_response = await response.run_ui_command('notebook-intelligence:add-code-cell-to-active-notebook', {'source': args['source']})
@@ -1783,15 +1809,15 @@ def _create_client_options(self) -> ClaudeAgentOptions:
         self._jupyter_ui_tools_mcp_server = create_sdk_mcp_server(
             name="nbi",
             version="1.0.0",
-            tools=[create_new_notebook, add_markdown_cell, add_code_cell, get_number_of_cells, get_cell_type_and_source, get_cell_output, set_cell_type_and_source, delete_cell, insert_cell, run_cell, save_notebook, rename_notebook, run_command_in_jupyter_terminal, open_file_in_jupyter_ui]
+            tools=[list_available_notebook_kernels, create_new_notebook, add_markdown_cell, add_code_cell, get_number_of_cells, get_cell_type_and_source, get_cell_output, set_cell_type_and_source, delete_cell, insert_cell, run_cell, save_notebook, rename_notebook, run_command_in_jupyter_terminal, open_file_in_jupyter_ui]
         )
         mcp_servers = {}
         jupyter_ui_tools_enabled = ClaudeToolType.JupyterUITools in claude_settings.get('tools', [])
         if jupyter_ui_tools_enabled:
             mcp_servers["nbi"] = self._jupyter_ui_tools_mcp_server
         allowed_tools = []
         if jupyter_ui_tools_enabled:
-            allowed_tools.extend(["mcp__nbi__create-new-notebook", "mcp__nbi__add-markdown-cell", "mcp__nbi__add-code-cell", "mcp__nbi__get-number-of-cells", "mcp__nbi__get-cell-type-and-source", "mcp__nbi__get-cell-output", "mcp__nbi__set-cell-type-and-source", "mcp__nbi__insert-cell", "mcp__nbi__save-notebook", "mcp__nbi__rename-notebook", "mcp__nbi__open-file-in-jupyter-ui"])
+            allowed_tools.extend(["mcp__nbi__list-available-notebook-kernels", "mcp__nbi__create-new-notebook", "mcp__nbi__add-markdown-cell", "mcp__nbi__add-code-cell", "mcp__nbi__get-number-of-cells", "mcp__nbi__get-cell-type-and-source", "mcp__nbi__get-cell-output", "mcp__nbi__set-cell-type-and-source", "mcp__nbi__insert-cell", "mcp__nbi__save-notebook", "mcp__nbi__rename-notebook", "mcp__nbi__open-file-in-jupyter-ui"])
         setting_sources = claude_settings.get('setting_sources')
         chat_model_id = claude_settings.get('chat_model', '').strip()
         if chat_model_id == "":

notebook_intelligence/context_factory.py

@@ -8,11 +8,18 @@ class RuleContextFactory:
     """Factory for creating RuleContext from various sources."""
 
     @staticmethod
-    def create(filename: str, language: str, chat_mode_id: str, root_dir: str) -> RuleContext:
+    def create(
+        filename: str,
+        language: str,
+        chat_mode_id: str,
+        root_dir: str,
+        kernel_name: str | None = None,
+    ) -> RuleContext:
         """Create RuleContext from WebSocket message data."""
         return RuleContext(
             filename=filename,
-            kernel=language,
+            language=language,
+            kernel_name=kernel_name or None,
             mode=chat_mode_id,
             directory=os.path.dirname(os.path.join(root_dir, filename))
         )