pydantic-ai-kernel 1.5.2__py3-none-any.whl

pydantic_ai_kernel/__init__.py

@@ -0,0 +1,6 @@
+"""A wrapper of pydantic-ai in jupyter kernel"""
+
+from .kernel import PydanticAIBaseKernel  # noqa: F401
+from .jupydantic_agent import JuPydanticAgent
+from .agent_config import AgentConfig, ModelProviderConfig, ModelConfig  # noqa: F401
+from .boosted_magic import BoostedMagic, boosted_option, complete_from_list

pydantic_ai_kernel/__main__.py

@@ -0,0 +1,4 @@
+from . import PydanticAIBaseKernel
+
+
+PydanticAIBaseKernel.run_as_main()

pydantic_ai_kernel/agent_config.py

@@ -0,0 +1,330 @@
+from pydantic_ai import UserError
+from pydantic_ai.models import infer_model, Model
+from pydantic_ai.providers import infer_provider_class, Provider
+from typing import Annotated, Any, Literal, Optional, Dict
+from pydantic import BaseModel, Field, field_validator
+import json
+import os
+
+AllModelKind = Literal[
+    "test",
+    "openai",
+    "azure",
+    "deepseek",
+    "fireworks",
+    "github",
+    "grok",
+    "heroku",
+    "moonshotai",
+    "ollama",
+    "together",
+    "vercel",
+    "litellm",
+    "nebius",
+    "ovhcloud",
+    "alibaba",
+    "openai-chat",
+    "openai-responses",
+    "google-gla",
+    "google-vertex",
+    "google",
+    "groq",
+    "cohere",
+    "mistral",
+    "openrouter",
+    "anthropic",
+    "bedrock",
+    "huggingface",
+    "cerebras",
+]
+
+
+def infer_model_type(  # noqa: C901
+    model_kind: str,
+) -> type[Model]:
+    """Infer the model type from the name.
+    Function extracted and modified from pydantic-ai `infer_model`.
+
+    Args:
+        model_kind:
+            Kind of model to instantiate. Use the string "test" to instantiate TestModel.
+    """
+    if model_kind == "test":
+        from pydantic_ai.models.test import TestModel
+
+        return TestModel
+
+    if model_kind in (
+        "openai",
+        "azure",
+        "deepseek",
+        "fireworks",
+        "github",
+        "grok",
+        "heroku",
+        "moonshotai",
+        "ollama",
+        "together",
+        "vercel",
+        "litellm",
+        "nebius",
+        "ovhcloud",
+        "alibaba",
+    ):
+        model_kind = "openai-chat"
+    elif model_kind in ("google-gla", "google-vertex"):
+        model_kind = "google"
+
+    if model_kind == "openai-chat":
+        from pydantic_ai.models.openai import OpenAIChatModel
+
+        return OpenAIChatModel
+    elif model_kind == "openai-responses":
+        from pydantic_ai.models.openai import OpenAIResponsesModel
+
+        return OpenAIResponsesModel
+    elif model_kind == "google":
+        from pydantic_ai.models.google import GoogleModel
+
+        return GoogleModel
+    elif model_kind == "groq":
+        from pydantic_ai.models.groq import GroqModel
+
+        return GroqModel
+    elif model_kind == "cohere":
+        from pydantic_ai.models.cohere import CohereModel
+
+        return CohereModel
+    elif model_kind == "mistral":
+        from pydantic_ai.models.mistral import MistralModel
+
+        return MistralModel
+    elif model_kind == "openrouter":
+        from pydantic_ai.models.openrouter import OpenRouterModel
+
+        return OpenRouterModel
+    elif model_kind == "anthropic":
+        from pydantic_ai.models.anthropic import AnthropicModel
+
+        return AnthropicModel
+    elif model_kind == "bedrock":
+        from pydantic_ai.models.bedrock import BedrockConverseModel
+
+        return BedrockConverseModel
+    elif model_kind == "huggingface":
+        from pydantic_ai.models.huggingface import HuggingFaceModel
+
+        return HuggingFaceModel
+    elif model_kind == "cerebras":
+        from pydantic_ai.models.cerebras import CerebrasModel
+
+        return CerebrasModel
+    else:
+        raise UserError(f"Unknown model: {model_kind}")  # pragma: no cover
+
+
+class ModelProviderConfig(BaseModel):
+    name: Annotated[
+        str,
+        Field(
+            description="Name of the inference provider. The complete list is the one from pydantic-ai. https://ai.pydantic.dev/api/providers/",
+            examples=[
+                "ollama",
+                "litellm",
+                "anthropic",
+                "nebius",
+                "alibaba",
+                "ovhcloud",
+            ],
+        ),
+    ]
+    params: Annotated[
+        dict,
+        Field(
+            description="The parameters that will be given to pydantic-ai Provider class. Depends on provider name. Go to : https://ai.pydantic.dev/models/overview/",
+            examples=[{"base_url": "http://localhost:11434/v1"}],
+        ),
+    ]
+
+    @property
+    def get_provider(self) -> Provider[Any]:
+        ProviderConstructor = infer_provider_class(self.name)
+        return ProviderConstructor(**self.params)
+
+
+class ModelConfig(BaseModel):
+    """
+    Describe a model and its inference provider. Is given to pydantic-ai
+    agent.
+    """
+
+    model_type: AllModelKind | None = None
+    model_name: Annotated[
+        str,
+        Field(
+            description="""Following pydantic-ai implementation, there are two ways
+ to write model name :
+    - <provider_name>:<model_name>, example 'ollama:qwen3:1.7b
+    - <model_name>
+    
+In the first one, you CAN provide additionnal informations on the provider within
+the field 'provider' below. They will **erase** the ones infered by pydantic-ai.
+A wide range of values is possible here, and you can refer to pydantic-ai documentation :
+https://ai.pydantic.dev/models/overview/
+
+In the second one, you MUST provide additionnal informations on the provider within
+the field 'provider' below.""",
+            examples=["ollama:qwen3:1.7b", "qwen3.1:1.7b"],
+        ),
+    ]
+    model_settings: Annotated[
+        dict | None,
+        Field(
+            description="Settings for the model. Depends on model and inference provider. See https://ai.pydantic.dev/api/settings/#pydantic_ai.settings.ModelSettings",
+            default=None,
+        ),
+    ]
+    model_provider: Annotated[
+        ModelProviderConfig | None,
+        Field(
+            description="Description of the inference provider. When specified, inferred information about the provider from model_name will be erased.",
+            default=None,
+        ),
+    ]
+
+    @property
+    def get_model(self) -> Model | None:
+        if self.model_type is None:
+            try:
+                return infer_model(self.model_name)
+            except Exception as e:
+                raise Exception(
+                    f"Could not access model. See error above or consider adding at least information about the model type in the configuration file. Model types are : {AllModelKind}."
+                ) from e
+        if self.model_provider is None:
+            try:
+                return infer_model(self.model_name)
+            except Exception as e:
+                raise Exception(
+                    "Could not access model. See error above or consider adding at least a model provider in configuration file"
+                ) from e
+
+        provider = self.model_provider.get_provider
+        ModelConstructor = infer_model_type(self.model_type)
+        return ModelConstructor(self.model_name, provider=provider)  # pyright: ignore
+
+
+class AgentConfig(BaseModel):
+    agent_name: Annotated[
+        str,
+        Field(
+            description="The name of the agent within the application - different from model name.",
+            examples=["my-agent", "summarizer", "image-generator"],
+        ),
+    ]
+    system_prompt: Annotated[
+        str,
+        Field(
+            description="The system prompt for this agent, or a path towards a text file containing system prompt.",
+            examples=[
+                "You are a nice agent who answers nicely to requests of the user",
+                "You are a specialist in cooking and are always willing to provide informations on new cooking recipees.",
+                "/Users/mgg/Documents/system_prompt.txt",
+                "/Users/mgg/Documents/system_prompt.md",
+            ],
+        ),
+    ]
+    model: ModelConfig
+    mcp_servers: Annotated[
+        Optional[Dict[str, Any]],
+        Field(
+            description="A dictionnary of MCP server configuration to connect to. Keys are MCP server names.",
+            examples=[
+                {
+                    "time_mcp_server": {
+                        "command": "uvx",
+                        "args": ["mcp-run-python", "stdio"],
+                    },
+                    "weather_server": {"command": "python", "args": ["mcp_server.py"]},
+                },
+                {
+                    "python-runner": {
+                        "command": "uv",
+                        "args": ["run", "mcp-run-python", "stdio"],
+                    },
+                    "weather": {"command": "python", "args": ["mcp_server.py"]},
+                    "weather-api": {"url": "http://localhost:3001/sse"},
+                    "calculator": {"url": "http://localhost:8000/mcp"},
+                },
+            ],
+        ),
+    ] = None
+    mcp_servers_user_approval: Annotated[
+        Optional[Dict[str, bool | dict]],
+        Field(
+            description="For each MCP server, information about whether the tool requires user approval or not. Can be either a boolean : True = All tools in the MCP server requires user validation (False = all tools **don't** requires validation). Or can be a dict where keys are tool names, and values boolean. When not specified, a tool needs user validation by default.",
+            examples=[
+                {
+                    "python-runner": True,
+                    "weather-api": {"get_weather": True, "authenticate_to_api": False},
+                    "calculator": False,
+                }
+            ],
+        ),
+    ] = None
+    display_thinking: Annotated[
+        bool,
+        Field(
+            description="Whether or not to display thinking text of the model. Default to true."
+        ),
+    ] = True
+    formatter: Annotated[
+        str,
+        Field(
+            description="Output formatter for magics (text, terminal, md, json).",
+            examples=["text", "terminal", "md", "json"],
+        ),
+    ] = "terminal"
+    use_widget: Annotated[
+        bool,
+        Field(
+            description="Whether to use ipywidget in frontend for enhanced displaying"
+        ),
+    ] = False
+
+    @field_validator("system_prompt", mode="after")
+    @classmethod
+    def load_system_prompt(cls, value: str):
+        """
+        If system prompt argument is a path, loads it.
+        Else, does nothing.
+        """
+        if os.path.isfile(value):
+            with open(value, "rt") as f:
+                o = f.read()
+            return o
+        return value
+
+    @field_validator("display_thinking", mode="after")
+    @classmethod
+    def validate_display_thinking(cls, value):
+        # env variable overrides config value
+        if os.getenv("PYDANTIC_AI_KERNEL_DISPLAY_THINKING", False):
+            return True
+        return value
+
+    @field_validator("formatter", mode="after")
+    @classmethod
+    def validate_formatter(cls, value):
+        # env variable overrides config value; fallback to default if not set
+        env = os.getenv("PYDANTIC_AI_KERNEL_FORMATTER", False)
+        if env:
+            return env
+        return value
+
+
+if __name__ == "__main__":
+    import json
+
+    with open("config_scheme.json", "wt") as f:
+        json.dump(AgentConfig.model_json_schema(), f)

pydantic_ai_kernel/boosted_magic.py

@@ -0,0 +1,256 @@
+from metakernel import Magic
+from metakernel.magic import (
+    _indent,
+    _split_args,
+)
+from typing import Any, TypeVar, Callable
+import inspect
+import traceback
+import argparse
+from statikomand import KomandParser
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+
+class BoostedMagic(Magic):
+    """
+    Subclass of metakernel Magics, that provides additional
+    features :
+        - completion for options names and values, based
+            on custom completer objects
+        - parsing based on argparse
+
+    To implement custom parser for magic line, the get_args
+    and call_magic had to be reimplemented.
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def get_args(self, mtype: str, name: str, code: str, args: Any) -> Any:
+        self.code = code
+        old_args = args
+        mtype = mtype.replace("sticky", "cell")
+
+        func = getattr(self, mtype + "_" + name)
+
+        if "-h" in args or "--help" in args:
+            # --help flag of argparse writes to stdout,
+            # which break kernels (do_is_complete message is
+            # unanswered, which makes the kernel to crash).
+            return ["help"], dict(), list()
+
+        try:
+            args, kwargs = _boosted_parse_args(
+                func, args, usage=self.get_help(mtype, name)
+            )
+        except Exception as e:
+            self.kernel.Error(str(e))
+            return self
+        arg_spec = inspect.getfullargspec(func)
+        fargs = arg_spec.args
+        if fargs[0] == "self":
+            fargs = fargs[1:]
+
+        fargs = [f for f in fargs if f not in kwargs.keys()]
+        if len(args) > len(fargs) and not arg_spec.varargs:
+            extra = " ".join(str(s) for s in (args[len(fargs) - 1 :]))
+            args = args[: len(fargs) - 1] + [extra]
+
+        return (args, kwargs, old_args)
+
+    def call_magic(self, mtype: str, name: str, code: str, args: Any) -> Magic:
+        self.code = code
+        old_args = args
+        mtype = mtype.replace("sticky", "cell")
+
+        func = getattr(self, mtype + "_" + name)
+        try:
+            args, kwargs = _boosted_parse_args(
+                func, args, usage=self.get_help(mtype, name)
+            )
+        except Exception as e:
+            self.kernel.Error(str(e))
+            return self
+
+        arg_spec = inspect.getfullargspec(func)
+        fargs = arg_spec.args
+        if fargs[0] == "self":
+            fargs = fargs[1:]
+
+        fargs = [f for f in fargs if f not in kwargs.keys()]
+        if len(args) > len(fargs) and not arg_spec.varargs:
+            extra = " ".join(str(s) for s in (args[len(fargs) - 1 :]))
+            args = args[: len(fargs) - 1] + [extra]
+
+        try:
+            try:
+                func(*args, **kwargs)
+            except TypeError:
+                func(old_args)
+        except Exception as exc:
+            msg = f"Error in calling magic '{name}' on {mtype}:\n    {exc!s}\n    args: {args}\n    kwargs: {kwargs}"
+            self.kernel.Error(msg)
+            self.kernel.Error(traceback.format_exc())
+            self.kernel.Error(self.get_help(mtype, name))
+            # return dummy magic to end processing:
+            return Magic(self.kernel)
+        return self
+
+    def get_completions(self, info: dict[str, Any]) -> list[str]:
+        """
+        Uses statikomand argparse subclass to complete the current magic.
+        """
+        self.kernel.log.info(f"Info boosted magic : {info}")
+        func = None
+        for each_attr in dir(self):
+            if each_attr.startswith("line"):
+                func = getattr(self, each_attr)
+                break
+            # for now, info does not contains magic type
+            # if each_attr.startswith("cell"):
+            #     func = getattr(self, each_attr)
+            #     break
+        if func is None:
+            self.kernel.log.warning("Could not find any magic to complete.")
+            return []
+        if not getattr(func, "has_boosted_options", False):
+            self.kernel.log.debug(
+                "No boosted options found for magic, hence not completing."
+            )
+            return []
+        boosted_parser: KomandParser | bool = getattr(func, "boosted_parser", False)
+        if not boosted_parser:
+            self.kernel.log.warning(
+                "Could not find boosted parser for magic, but has_boosted_options is True. Not completing."
+            )
+            return []
+        args = info["code"]
+        if isinstance(args, list):
+            args = " ".join(args)
+
+        args = _split_args(args)
+        if not isinstance(boosted_parser, KomandParser):
+            raise TypeError(
+                f"boosted_parser parameter of magic method must be a KomandParser, not {type(boosted_parser)}."
+            )
+
+        # finally calls complete method of boosted parser
+        try:
+            completed = boosted_parser.complete(args)
+        except Exception:
+            self.kernel.log.debug(
+                f"Error when completing option : {traceback.format_exc()}"
+            )
+            return []
+
+        if len(args) > 0:
+            self.kernel.log.info(f"Modifying -- and - tokens : {completed}")
+            result = []
+            last_token = args[-1]
+            self.kernel.log.info(f"last token : {last_token}")
+            for each_match in completed:
+                if each_match == last_token:
+                    continue
+                # for option tokens (starting with '-'), metakernel does
+                # not replace the full token, but adds the completion.
+                # so we need to truncate it. Somehow, tokens starting
+                # with '-' are not treated as others.
+                # For other tokens, it works as statikomand.
+                if last_token.startswith("--") and each_match.startswith("--"):
+                    each_match = each_match[2:]
+                elif last_token.startswith("-") and each_match.startswith("-"):
+                    each_match = each_match[1:]
+
+                result.append(each_match)
+        else:
+            result = completed
+        self.kernel.log.debug(f"Boosted magic completion : {result}")
+        return result
+
+
+def boosted_option(*args: Any, **kwargs: Any) -> Callable[[_F], _F]:
+    """
+    Return decorator that adds a boosted magic option to a function.
+    The boosted magic uses argparse instead of optparse, and
+    allows for tab completion among options names and values.
+    """
+
+    def decorator(func: _F) -> _F:
+        help_text = ""
+        if not getattr(func, "has_boosted_options", False):
+            func.has_boosted_options = True  # type:ignore[attr-defined]
+            func.boosted_parser = KomandParser(prog=func.__name__, add_help=False)  # type:ignore[attr-defined]
+            # boosted_parser is an instance of argparse.ArgumentParser
+            # add_help is set to False because argparse help writes to stdout, which
+            # breaks kernel machinery
+            help_text += "Options:\n-------\n"
+
+        try:
+            action = func.boosted_parser.add_argument(*args, **kwargs)  # type:ignore[attr-defined]
+        except AttributeError:
+            pass
+        except argparse.ArgumentError:
+            help_text += args[0] + "\n"
+        else:
+            help_text += _format_action(action) + "\n"
+        if func.__doc__:
+            func.__doc__ += _indent(func.__doc__, help_text)
+        else:
+            func.__doc__ = help_text
+        return func
+
+    return decorator
+
+
+def _boosted_parse_args(
+    func, args: str | list[str], usage
+) -> tuple[list[Any], dict[str, Any]]:
+
+    if isinstance(args, list):
+        args = " ".join(args)
+    args = _split_args(args)
+
+    # raise ValueError(f"dajkbz {args}")
+    if getattr(func, "has_boosted_options", False) and getattr(
+        func, "boosted_parser", False
+    ):
+        if not isinstance(func.boosted_parser, argparse.ArgumentParser):
+            return list(), dict()
+        parsed_args = func.boosted_parser.parse_args(args)
+        return parsed_args._get_args(), parsed_args.__dict__
+    return list(), dict()
+
+
+def _format_action(action: argparse.Action):
+    out = ""
+    out = ", ".join(action.option_strings) + f" : {action.help}"
+    return out
+
+
+def complete_from_list(
+    list_of_reference: list[str],
+    word: str,
+    rank: int = 0,
+):
+    """
+    Completes a word from a given list_of_reference.
+    Return all occurences of list_of_reference whose
+    beginning matches word.
+    Can be used as within a lambda function with a custom list:
+
+    ```python
+    my_list = ["bobo", "baba"]
+    completer = lambda w,r : complete_from_list(my_list, w, r)
+    ```
+
+    """
+    all_matches = []
+
+    for each_word in list_of_reference:
+        if len(each_word) < len(word):
+            continue
+        potential_match = each_word[: len(word)]
+        if potential_match == word:
+            all_matches.append(each_word)
+    return all_matches