Pydantic AI Base Kernel

Bases: MetaKernel

Kernel wrapper for pydantic agents. Allows to interact with an AI agent through jupyter kernel messaging protocol.

Allows to access agent history (tool calling history). Thanks to pydantic-ai library, nearly all inference providers are supported. With metakernel magics, you can load a new config file, display agent history.

Tool validation is supported through pydantic-ai DeferredToolRequest being forwarded to frontend with jupyter input_request messages.

Most of the metakernel magics are disabled on purpose.

Source code in pydantic_ai_kernel/kernel.py

class PydanticAIBaseKernel(MetaKernel):
    """
    Kernel wrapper for pydantic agents. Allows to interact with an AI agent
    through jupyter kernel messaging protocol.

    Allows to access agent history (tool calling history). Thanks to pydantic-ai
    library, nearly all inference providers are supported.
    With metakernel magics, you can load a new config file,
    display agent history.

    Tool validation is supported through pydantic-ai DeferredToolRequest
    being forwarded to frontend with jupyter input_request messages.

    Most of the metakernel magics are disabled on purpose.
    """

    implementation = "PydanticAI Agent Kernel"
    implementation_version = "1.0"
    language = "no-op"
    help_suffix = (
        "♫"  # unusual suffix on purpose, because ? is very common in natural language
    )
    language_version = "0.1"
    language_info = {
        "name": "pydantic_ai",
        "mimetype": "text/plain",
        "file_extension": ".ai",
    }
    banner = "Pydantic AI Kernel"
    app_name = "pydantic_ai"

    def __init__(
        self,
        agent_config: AgentConfig | None = None,
        tools: list[Tool] | None = None,
        toolsets: Sequence[FunctionToolset | ToolsetFunc] | None = None,
        output_type: Type = str,
        authorized_magics_names: list[str] | None = None,
        additional_agent_kwargs: dict[str, Any] | None = None,
        **kwargs,
    ):
        """
        Parameters :
        ---
            - agent_config (AgentConfig | None = None) : the configuration object for the
                agent. Deals with system prompt, inference provider API, ...
            - tools (list[Tool] | None = None) : list of pydantic-ai tools, which can be
                given to the agent on startup. Useful for subclasses.
            - toolsets (list[FunctionToolset] | None = None) : list of toolsets given to
                the agent on startup. Useful for subclasses.
            - output_type (Type = str) : python Type for a constrained output of the agent.
                See pydantic-ai documentation.
            - authorized_magics_names (list[str] | None = None) : Useful for subclasses.
                List of class names for the magics you want to whitelist for this kernel.
                By default, all non-whitelisted magics are not loaded. You must manually
                add a class name (e.g. FileMagic, ...) to have it running on this kernel.
            - additional_agent_kwargs (dict[str, Any] | None = None) : any kwargs which
                will be given to pydantic-ai agent initialization.
        """

        if authorized_magics_names is None:
            authorized_magics_names = []
        authorized_magics_names += [
            "HelpMagic",
            "MagicMagic",
            "AgentHistoryMagic",
            "ToolsMagic",
            "ForgetMagic",
            "ConfigMagic",
            "FileMagic",
            "MCPMagic",
            "UsageMagic",
            "ThinkMagic",
            "FormatterMagic",
        ]
        self.authorized_magics_names = authorized_magics_names
        self.additional_agent_kwargs = (
            additional_agent_kwargs if additional_agent_kwargs else {}
        )

        super().__init__(**kwargs)
        if os.getenv("PYDANTIC_AI_KERNEL_LOG_LEVEL", False):
            add_custom_logger_handler(self.log)
        self.agent_config = agent_config

        self.tools = tools if tools is not None else []
        self.toolsets: list[AbstractToolset | FunctionToolset | ToolsetFunc] = (
            list(toolsets) if toolsets is not None else []
        )
        self.base_toolsets = (
            toolsets if toolsets is not None else []
        )  # initial toolsets, before loading
        # additional from config file

        self.log.info(
            f"Initialization: tools : {self.tools}, toolsets : {self.toolsets}"
        )
        self._agent = None
        self.output_type = output_type
        self.all_messages_ids = []
        self.config_has_changed = False
        self.out_md_stack = []
        self.log.info("Initialization of Pydantic AI Kernel is successful.")

    @property
    def formatter(self):
        if self.agent_config is None:
            return self.load_config().formatter
        return self.agent_config.formatter

    @property
    def use_widget(self):
        return self.agent.use_widget

    def reload_magics(self) -> None:
        """
        Override metakernel `reload_magics` to fix the magics not being
        propagated through subclasses. Subclasses of pydantic-ai-kernel
        could not by default access pydantic-ai-kernel magics.

        We first reload_magics with the official method, and then we seek
        for a magics dir **within** this file's directory.
        """
        super().reload_magics()
        try:
            this_class_magics_dir = os.path.join(
                os.path.dirname(os.path.abspath(__file__)), "magics"
            )
            sys.path.append(this_class_magics_dir)
            this_class_magic_files = glob.glob(
                os.path.join(this_class_magics_dir, "*.py")
            )
            for this_class_magic in this_class_magic_files:
                basename = os.path.basename(this_class_magic)
                if basename == "__init__.py":
                    continue
                module = __import__(os.path.splitext(basename)[0])
                importlib.reload(module)
                module.register_magics(self)
        except Exception as e:
            self.log.debug(
                f"PydanticAIBaseKernel could not register its own magics : `{e}`."
            )

    def register_magics(self, magic_klass: type[Magic]) -> None:
        """
        Override the metakernel 'register_magics', to prevent of
        loading non-whitelisted magics.

        Parameters :
        ---
            - magic_klass: The subclass of Magic
        """
        if magic_klass.__name__ in self.authorized_magics_names:
            self.log.info(f"Adding magic : {magic_klass}")
            super().register_magics(magic_klass)

    def get_config_dir(self) -> str:
        home_directory = os.path.expanduser("~")
        config_name = f"jupyter_{self.app_name}_config.yaml"
        return os.path.join(home_directory, f".jupyter/{config_name}")

    def load_config(self, path: str | None = None) -> AgentConfig:
        """
        Try to load config file at ~/.jupyter/jupyter_<kernel_name>_config.yaml.
        Returns the validated config object, or raise an Error.

        Parameters :
        ---
            - path (Optional[str]) : path to the config file. If None,
                is set to ~/.jupyter/jupyter_<app_name>_config.yaml

        Returns :
        ---
            AgentConfig pydantic BaseModel, validated
        """
        if path is None:
            dir = self.get_config_dir()
        else:
            dir = path
        try:
            with open(dir, "rt") as f:
                conf = yaml.safe_load(f)
            validated_conf = AgentConfig.model_validate(conf)
            return validated_conf
        except Exception as e:
            raise LoadConfigError(
                f"Could not load and validate config file for agent at {dir}."
            ) from e

    def create_agent(self) -> JuPydanticAgent:
        """
        Creates the pydantic-ai agent, from config file.
        """
        if self.agent_config is None:
            self.agent_config = self.load_config()
        try:
            model = self.agent_config.model.get_model
        except NotImplementedError as e:
            model = self.agent_config.model.model_name
            self.log.warning(e)

        self.config_toolsets = self.create_toolsets_from_config()
        self.log.debug(f"Adding toolsets from configuration : {self.config_toolsets}")
        self.toolsets = self.base_toolsets + self.config_toolsets

        agent = Agent(
            model,
            output_type=[self.output_type, DeferredToolRequests],
            system_prompt=self.agent_config.system_prompt,
            tools=self.tools,
            toolsets=self.toolsets,
            name=self.agent_config.agent_name,
            **self.additional_agent_kwargs,
        )

        return JuPydanticAgent(agent, self, self.agent_config)

    def create_toolsets_from_config(self) -> list[AbstractToolset]:
        """
        Uses config file to create abstract toolsets
        """
        if self.agent_config is None:
            self.agent_config = self.load_config()

        if self.agent_config.mcp_servers is None:
            self.log.info("No MCP servers loaded")
            return []

        with NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump({"mcpServers": self.agent_config.mcp_servers}, f)
            temp_path = f.name
        toolsets = load_mcp_toolsets(temp_path)
        self.log.info(
            f"Raw MCP toolsets loaded from config: {toolsets}. Boosting them."
        )
        if self.agent_config.mcp_servers_user_approval is None:
            mcp_servers_user_approval = {
                each_mcp: True for each_mcp in self.agent_config.mcp_servers
            }
        else:
            mcp_servers_user_approval = self.agent_config.mcp_servers_user_approval

        out_toolsets = []
        for each_mcp in toolsets:
            if not isinstance(each_mcp, PrefixedToolset):
                continue
            if (
                each_mcp.wrapped.id not in mcp_servers_user_approval
                or each_mcp.wrapped.id is None
            ):
                continue

            user_approval = mcp_servers_user_approval[each_mcp.wrapped.id]
            if isinstance(user_approval, bool):
                if user_approval:
                    checked_toolset = ApprovalRequiredToolset(each_mcp)
                else:
                    checked_toolset = each_mcp
            elif isinstance(user_approval, dict):
                approv = user_approval
                checked_toolset = each_mcp.approval_required(
                    lambda ctx, tool_def, tool_args: approv.get(tool_def.name, True)
                )
            else:
                continue
            out_toolsets.append(checked_toolset)

        Path(temp_path).unlink()

        return out_toolsets

    @property
    def agent(self) -> JuPydanticAgent:
        if self._agent is None:
            self._agent = self.create_agent()
        return self._agent

    def ask_user_approval(self, prompt: Optional[str] = None) -> bool:
        """
        Sends a input_request message to the frontend, asking
        for user approval.

        Returns:
        ---
            True if the user gave his approval, else False
        """
        user_approval = self.raw_input(f"{prompt} ([y]/n):")
        return user_approval in ["Y", "yes", 1, "1", "y"]

    async def do_execute_direct(self, code: str, silent=False):
        """
        Sends the code to the agent, retrieve the output from
        an async loop, and stream it to IOPub channel.

        Parameters :
        ---
            - code (str): the code (prompt) which is sent to the agent
            - silent (bool = False) : whether to execute the code silently.
                See https://jupyter-client.readthedocs.io/en/stable/messaging.html#execute
        """
        if self.config_has_changed:
            # we reload the agent and the config
            self.agent_config = self.load_config()
            self.create_agent()
        await self.agent.run(code, silent)

    async def do_is_complete(self, code: str) -> dict[str, str]:
        """
        Overwrites metakernel do_is_complete to have :
            - magic commands requires an empty line (as usual)
            - non-magic commands do not need an empty line
            - non-magic commands can be multiline if the line
                ends with ' '.
        """
        if code.startswith(self.magic_prefixes["magic"]):
            ## force requirement to end with an empty line
            if code.endswith("\n"):
                return {"status": "complete"}
            else:
                return {"status": "incomplete"}
        # otherwise, how to know is complete?
        elif code.endswith(" "):
            return {"status": "incomplete"}
        else:
            return {"status": "complete"}

    async def interrupt_request(self, stream, ident, parent):
        """
        This method is called when an interrupt_request message is received.

        By default, IPykernel restart the kernel on an interrupt request.
        We override this here, to make an early stop of the agent stream.
        """
        self.log.info("Cell interuption asked")
        if not self.session:
            return

        self.agent.is_interrupted = True  # ask the agent to stop
        content = {"status": "ok"}
        self.session.send(stream, "interrupt_reply", content, parent, ident=ident)
        return

    def get_completions(self, info: dict[str, Any]) -> list[str]:
        self.log.info(f"Completion info : {info}")
        return super().get_completions(info)

__init__(agent_config=None, tools=None, toolsets=None, output_type=str, authorized_magics_names=None, additional_agent_kwargs=None, **kwargs)

Parameters :

- agent_config (AgentConfig | None = None) : the configuration object for the
    agent. Deals with system prompt, inference provider API, ...
- tools (list[Tool] | None = None) : list of pydantic-ai tools, which can be
    given to the agent on startup. Useful for subclasses.
- toolsets (list[FunctionToolset] | None = None) : list of toolsets given to
    the agent on startup. Useful for subclasses.
- output_type (Type = str) : python Type for a constrained output of the agent.
    See pydantic-ai documentation.
- authorized_magics_names (list[str] | None = None) : Useful for subclasses.
    List of class names for the magics you want to whitelist for this kernel.
    By default, all non-whitelisted magics are not loaded. You must manually
    add a class name (e.g. FileMagic, ...) to have it running on this kernel.
- additional_agent_kwargs (dict[str, Any] | None = None) : any kwargs which
    will be given to pydantic-ai agent initialization.

Source code in pydantic_ai_kernel/kernel.py

def __init__(
    self,
    agent_config: AgentConfig | None = None,
    tools: list[Tool] | None = None,
    toolsets: Sequence[FunctionToolset | ToolsetFunc] | None = None,
    output_type: Type = str,
    authorized_magics_names: list[str] | None = None,
    additional_agent_kwargs: dict[str, Any] | None = None,
    **kwargs,
):
    """
    Parameters :
    ---
        - agent_config (AgentConfig | None = None) : the configuration object for the
            agent. Deals with system prompt, inference provider API, ...
        - tools (list[Tool] | None = None) : list of pydantic-ai tools, which can be
            given to the agent on startup. Useful for subclasses.
        - toolsets (list[FunctionToolset] | None = None) : list of toolsets given to
            the agent on startup. Useful for subclasses.
        - output_type (Type = str) : python Type for a constrained output of the agent.
            See pydantic-ai documentation.
        - authorized_magics_names (list[str] | None = None) : Useful for subclasses.
            List of class names for the magics you want to whitelist for this kernel.
            By default, all non-whitelisted magics are not loaded. You must manually
            add a class name (e.g. FileMagic, ...) to have it running on this kernel.
        - additional_agent_kwargs (dict[str, Any] | None = None) : any kwargs which
            will be given to pydantic-ai agent initialization.
    """

    if authorized_magics_names is None:
        authorized_magics_names = []
    authorized_magics_names += [
        "HelpMagic",
        "MagicMagic",
        "AgentHistoryMagic",
        "ToolsMagic",
        "ForgetMagic",
        "ConfigMagic",
        "FileMagic",
        "MCPMagic",
        "UsageMagic",
        "ThinkMagic",
        "FormatterMagic",
    ]
    self.authorized_magics_names = authorized_magics_names
    self.additional_agent_kwargs = (
        additional_agent_kwargs if additional_agent_kwargs else {}
    )

    super().__init__(**kwargs)
    if os.getenv("PYDANTIC_AI_KERNEL_LOG_LEVEL", False):
        add_custom_logger_handler(self.log)
    self.agent_config = agent_config

    self.tools = tools if tools is not None else []
    self.toolsets: list[AbstractToolset | FunctionToolset | ToolsetFunc] = (
        list(toolsets) if toolsets is not None else []
    )
    self.base_toolsets = (
        toolsets if toolsets is not None else []
    )  # initial toolsets, before loading
    # additional from config file

    self.log.info(
        f"Initialization: tools : {self.tools}, toolsets : {self.toolsets}"
    )
    self._agent = None
    self.output_type = output_type
    self.all_messages_ids = []
    self.config_has_changed = False
    self.out_md_stack = []
    self.log.info("Initialization of Pydantic AI Kernel is successful.")

ask_user_approval(prompt=None)

Sends a input_request message to the frontend, asking for user approval.

---

True if the user gave his approval, else False

Source code in pydantic_ai_kernel/kernel.py

def ask_user_approval(self, prompt: Optional[str] = None) -> bool:
    """
    Sends a input_request message to the frontend, asking
    for user approval.

    Returns:
    ---
        True if the user gave his approval, else False
    """
    user_approval = self.raw_input(f"{prompt} ([y]/n):")
    return user_approval in ["Y", "yes", 1, "1", "y"]

create_agent()

Creates the pydantic-ai agent, from config file.

Source code in pydantic_ai_kernel/kernel.py

def create_agent(self) -> JuPydanticAgent:
    """
    Creates the pydantic-ai agent, from config file.
    """
    if self.agent_config is None:
        self.agent_config = self.load_config()
    try:
        model = self.agent_config.model.get_model
    except NotImplementedError as e:
        model = self.agent_config.model.model_name
        self.log.warning(e)

    self.config_toolsets = self.create_toolsets_from_config()
    self.log.debug(f"Adding toolsets from configuration : {self.config_toolsets}")
    self.toolsets = self.base_toolsets + self.config_toolsets

    agent = Agent(
        model,
        output_type=[self.output_type, DeferredToolRequests],
        system_prompt=self.agent_config.system_prompt,
        tools=self.tools,
        toolsets=self.toolsets,
        name=self.agent_config.agent_name,
        **self.additional_agent_kwargs,
    )

    return JuPydanticAgent(agent, self, self.agent_config)

create_toolsets_from_config()

Uses config file to create abstract toolsets

Source code in pydantic_ai_kernel/kernel.py

def create_toolsets_from_config(self) -> list[AbstractToolset]:
    """
    Uses config file to create abstract toolsets
    """
    if self.agent_config is None:
        self.agent_config = self.load_config()

    if self.agent_config.mcp_servers is None:
        self.log.info("No MCP servers loaded")
        return []

    with NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        json.dump({"mcpServers": self.agent_config.mcp_servers}, f)
        temp_path = f.name
    toolsets = load_mcp_toolsets(temp_path)
    self.log.info(
        f"Raw MCP toolsets loaded from config: {toolsets}. Boosting them."
    )
    if self.agent_config.mcp_servers_user_approval is None:
        mcp_servers_user_approval = {
            each_mcp: True for each_mcp in self.agent_config.mcp_servers
        }
    else:
        mcp_servers_user_approval = self.agent_config.mcp_servers_user_approval

    out_toolsets = []
    for each_mcp in toolsets:
        if not isinstance(each_mcp, PrefixedToolset):
            continue
        if (
            each_mcp.wrapped.id not in mcp_servers_user_approval
            or each_mcp.wrapped.id is None
        ):
            continue

        user_approval = mcp_servers_user_approval[each_mcp.wrapped.id]
        if isinstance(user_approval, bool):
            if user_approval:
                checked_toolset = ApprovalRequiredToolset(each_mcp)
            else:
                checked_toolset = each_mcp
        elif isinstance(user_approval, dict):
            approv = user_approval
            checked_toolset = each_mcp.approval_required(
                lambda ctx, tool_def, tool_args: approv.get(tool_def.name, True)
            )
        else:
            continue
        out_toolsets.append(checked_toolset)

    Path(temp_path).unlink()

    return out_toolsets

do_execute_direct(code, silent=False) async

Sends the code to the agent, retrieve the output from an async loop, and stream it to IOPub channel.

Parameters :

- code (str): the code (prompt) which is sent to the agent
- silent (bool = False) : whether to execute the code silently.
    See https://jupyter-client.readthedocs.io/en/stable/messaging.html#execute

Source code in pydantic_ai_kernel/kernel.py

async def do_execute_direct(self, code: str, silent=False):
    """
    Sends the code to the agent, retrieve the output from
    an async loop, and stream it to IOPub channel.

    Parameters :
    ---
        - code (str): the code (prompt) which is sent to the agent
        - silent (bool = False) : whether to execute the code silently.
            See https://jupyter-client.readthedocs.io/en/stable/messaging.html#execute
    """
    if self.config_has_changed:
        # we reload the agent and the config
        self.agent_config = self.load_config()
        self.create_agent()
    await self.agent.run(code, silent)

do_is_complete(code) async

Overwrites metakernel do_is_complete to have

• magic commands requires an empty line (as usual)
• non-magic commands do not need an empty line
• non-magic commands can be multiline if the line ends with ' '.

Source code in pydantic_ai_kernel/kernel.py

async def do_is_complete(self, code: str) -> dict[str, str]:
    """
    Overwrites metakernel do_is_complete to have :
        - magic commands requires an empty line (as usual)
        - non-magic commands do not need an empty line
        - non-magic commands can be multiline if the line
            ends with ' '.
    """
    if code.startswith(self.magic_prefixes["magic"]):
        ## force requirement to end with an empty line
        if code.endswith("\n"):
            return {"status": "complete"}
        else:
            return {"status": "incomplete"}
    # otherwise, how to know is complete?
    elif code.endswith(" "):
        return {"status": "incomplete"}
    else:
        return {"status": "complete"}

interrupt_request(stream, ident, parent) async

This method is called when an interrupt_request message is received.

By default, IPykernel restart the kernel on an interrupt request. We override this here, to make an early stop of the agent stream.

Source code in pydantic_ai_kernel/kernel.py

async def interrupt_request(self, stream, ident, parent):
    """
    This method is called when an interrupt_request message is received.

    By default, IPykernel restart the kernel on an interrupt request.
    We override this here, to make an early stop of the agent stream.
    """
    self.log.info("Cell interuption asked")
    if not self.session:
        return

    self.agent.is_interrupted = True  # ask the agent to stop
    content = {"status": "ok"}
    self.session.send(stream, "interrupt_reply", content, parent, ident=ident)
    return

load_config(path=None)

Try to load config file at ~/.jupyter/jupyter__config.yaml. Returns the validated config object, or raise an Error.

Parameters :

- path (Optional[str]) : path to the config file. If None,
    is set to ~/.jupyter/jupyter_<app_name>_config.yaml

Returns :

AgentConfig pydantic BaseModel, validated

Source code in pydantic_ai_kernel/kernel.py

def load_config(self, path: str | None = None) -> AgentConfig:
    """
    Try to load config file at ~/.jupyter/jupyter_<kernel_name>_config.yaml.
    Returns the validated config object, or raise an Error.

    Parameters :
    ---
        - path (Optional[str]) : path to the config file. If None,
            is set to ~/.jupyter/jupyter_<app_name>_config.yaml

    Returns :
    ---
        AgentConfig pydantic BaseModel, validated
    """
    if path is None:
        dir = self.get_config_dir()
    else:
        dir = path
    try:
        with open(dir, "rt") as f:
            conf = yaml.safe_load(f)
        validated_conf = AgentConfig.model_validate(conf)
        return validated_conf
    except Exception as e:
        raise LoadConfigError(
            f"Could not load and validate config file for agent at {dir}."
        ) from e

register_magics(magic_klass)

Override the metakernel 'register_magics', to prevent of loading non-whitelisted magics.

Parameters :

- magic_klass: The subclass of Magic

Source code in pydantic_ai_kernel/kernel.py

def register_magics(self, magic_klass: type[Magic]) -> None:
    """
    Override the metakernel 'register_magics', to prevent of
    loading non-whitelisted magics.

    Parameters :
    ---
        - magic_klass: The subclass of Magic
    """
    if magic_klass.__name__ in self.authorized_magics_names:
        self.log.info(f"Adding magic : {magic_klass}")
        super().register_magics(magic_klass)

reload_magics()

Override metakernel reload_magics to fix the magics not being propagated through subclasses. Subclasses of pydantic-ai-kernel could not by default access pydantic-ai-kernel magics.

We first reload_magics with the official method, and then we seek for a magics dir within this file's directory.

Source code in pydantic_ai_kernel/kernel.py

def reload_magics(self) -> None:
    """
    Override metakernel `reload_magics` to fix the magics not being
    propagated through subclasses. Subclasses of pydantic-ai-kernel
    could not by default access pydantic-ai-kernel magics.

    We first reload_magics with the official method, and then we seek
    for a magics dir **within** this file's directory.
    """
    super().reload_magics()
    try:
        this_class_magics_dir = os.path.join(
            os.path.dirname(os.path.abspath(__file__)), "magics"
        )
        sys.path.append(this_class_magics_dir)
        this_class_magic_files = glob.glob(
            os.path.join(this_class_magics_dir, "*.py")
        )
        for this_class_magic in this_class_magic_files:
            basename = os.path.basename(this_class_magic)
            if basename == "__init__.py":
                continue
            module = __import__(os.path.splitext(basename)[0])
            importlib.reload(module)
            module.register_magics(self)
    except Exception as e:
        self.log.debug(
            f"PydanticAIBaseKernel could not register its own magics : `{e}`."
        )