type stubs for some third-party libraries

pydantic_ai_slim/pydantic_ai/models/outlines.py

@@ -60,7 +60,7 @@
     )
     from outlines.models.vllm_offline import (
         VLLMOffline,
-        from_vllm_offline,  # pyright: ignore[reportUnknownVariableType]
+        from_vllm_offline,
     )
     from outlines.types.dsl import JsonSchema
     from PIL import Image as PILImage
@@ -393,7 +393,7 @@ def _format_vllm_offline_inference_kwargs(  # pragma: no cover
         self, model_settings: dict[str, Any]
     ) -> dict[str, Any]:
         """Select the model settings supported by the vLLMOffline model."""
-        from vllm.sampling_params import SamplingParams  # pyright: ignore
+        from vllm.sampling_params import SamplingParams
 
         supported_args = [
             'max_tokens',

pydantic_ai_slim/pyproject.toml

@@ -78,7 +78,7 @@ bedrock = ["boto3>=1.40.14"]
 huggingface = ["huggingface-hub[inference]>=0.33.5"]
 outlines-transformers = ["outlines[transformers]>=1.0.0, <1.3.0; (sys_platform != 'darwin' or platform_machine != 'x86_64')", "transformers>=4.0.0", "pillow", "torch; (sys_platform != 'darwin' or platform_machine != 'x86_64')"]
 outlines-llamacpp = ["outlines[llamacpp]>=1.0.0, <1.3.0"]
-outlines-mlxlm = ["outlines[mlxlm]>=1.0.0, <1.3.0; (sys_platform != 'darwin' or platform_machine != 'x86_64')"]
+outlines-mlxlm = ["outlines[mlxlm]>=1.0.0, <1.3.0; platform_system == 'Darwin' and platform_machine == 'arm64'"]
 outlines-sglang = ["outlines[sglang]>=1.0.0, <1.3.0", "pillow"]
 outlines-vllm-offline = ["vllm; python_version < '3.12' and (sys_platform != 'darwin' or platform_machine != 'x86_64')", "torch; (sys_platform != 'darwin' or platform_machine != 'x86_64')", "outlines>=1.0.0, <1.3.0"]
 # Tools

pyproject.toml

@@ -56,7 +56,7 @@ dbos = ["pydantic-ai-slim[dbos]=={{ version }}"]
 prefect = ["pydantic-ai-slim[prefect]=={{ version }}"]
 outlines-transformers = ["pydantic-ai-slim[outlines-transformers]=={{ version }}"]
 outlines-llamacpp = ["pydantic-ai-slim[outlines-llamacpp]=={{ version }}"]
-outlines-mlxlm = ["pydantic-ai-slim[outlines-mlxlm]=={{ version }}"]
+outlines-mlxlm = ["pydantic-ai-slim[outlines-mlxlm]=={{ version }}; platform_system == 'Darwin' and platform_machine == 'arm64'"]
 outlines-sglang = ["pydantic-ai-slim[outlines-sglang]=={{ version }}"]
 outlines-vllm-offline = ["pydantic-ai-slim[outlines-vllm-offline]=={{ version }}"]
 
@@ -142,6 +142,7 @@ include = [
     "clai/**/*.py",
     "tests/**/*.py",
     "docs/**/*.py",
+    "stubs/**/*.pyi",
 ]
 
 [tool.ruff.lint]
@@ -186,8 +187,10 @@ quote-style = "single"
 "examples/**/*.py" = ["D101", "D103"]
 "tests/**/*.py" = ["D"]
 "docs/**/*.py" = ["D"]
+"stubs/**/*.pyi" = ["F401", "PYI044", "PYI035", "ANN401"]
 
 [tool.pyright]
+stubPath = "stubs"
 pythonVersion = "3.12"
 typeCheckingMode = "strict"
 reportMissingTypeStubs = false
@@ -217,6 +220,7 @@ exclude = [
 [tool.mypy]
 files = "tests/typed_agent.py,tests/typed_graph.py"
 strict = true
+mypy_path = "stubs"
 
 [tool.pytest.ini_options]
 testpaths = ["tests", "docs/.hooks"]

stubs/llama_cpp.pyi

@@ -0,0 +1,19 @@
+from collections.abc import Sequence
+from os import PathLike
+from typing import Any, Literal
+
+from typing_extensions import Self
+
+class Llama:
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        filename: str | None = None,
+        additional_files: Sequence[str] | None = None,
+        local_dir: str | PathLike[str] | None = None,
+        local_dir_use_symlinks: bool | Literal['auto'] = 'auto',
+        cache_dir: str | PathLike[str] | None = None,
+        **kwargs: Any,
+    ) -> Self: ...

stubs/mlx/__init__.pyi

@@ -0,0 +1,4 @@
+from typing import Any
+
+# mlx is imported as a package, primarily for mlx.nn
+__all__: list[str] = []

stubs/mlx/nn.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+class Module: ...

stubs/mlx_lm.pyi

@@ -0,0 +1,7 @@
+from typing import Any
+
+from mlx.nn import Module
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+def load(model_path: str) -> tuple[Module, PreTrainedTokenizer]: ...
+def generate_step(*args: Any, **kwargs: Any) -> Any: ...

stubs/outlines/__init__.pyi

@@ -0,0 +1,3 @@
+from . import models
+
+__all__: list[str] = []

stubs/outlines/models/__init__.pyi

@@ -0,0 +1,3 @@
+from . import base, llamacpp, mlxlm, sglang, transformers, vllm_offline
+
+__all__: list[str] = []

stubs/outlines/models/base.pyi

@@ -0,0 +1,10 @@
+from collections.abc import AsyncIterable, Iterable
+from typing import Any
+
+class Model:
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def stream(self, *args: Any, **kwargs: Any) -> Iterable[Any]: ...
+
+class AsyncModel(Model):
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def stream(self, *args: Any, **kwargs: Any) -> AsyncIterable[Any]: ...

stubs/outlines/models/llamacpp.pyi

@@ -0,0 +1,10 @@
+from typing import TYPE_CHECKING
+
+from outlines.models.base import Model
+
+if TYPE_CHECKING:
+    from llama_cpp import Llama
+
+class LlamaCpp(Model): ...
+
+def from_llamacpp(model: Llama) -> LlamaCpp: ...

stubs/outlines/models/mlxlm.pyi

@@ -0,0 +1,10 @@
+from typing import Any
+
+from mlx.nn import Module
+from outlines.models.base import Model
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+class MLXLM(Model):
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+
+def from_mlxlm(model: Module, tokenizer: PreTrainedTokenizer) -> MLXLM: ...

stubs/outlines/models/sglang.pyi

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING, Any, Union
+
+from outlines.models.base import AsyncModel, Model
+
+if TYPE_CHECKING:
+    from openai import AsyncOpenAI, OpenAI
+
+class SGLang(Model): ...
+class AsyncSGLang(AsyncModel): ...
+
+def from_sglang(client: OpenAI | AsyncOpenAI, *args: Any, **kwargs: Any) -> SGLang | AsyncSGLang: ...

stubs/outlines/models/transformers.pyi

@@ -0,0 +1,17 @@
+from typing import Any
+
+from outlines.models.base import Model
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, LlavaForConditionalGeneration
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+class Transformers(Model): ...
+class TransformersMultiModal(Model): ...
+
+def from_transformers(
+    model: PreTrainedModel | AutoModelForCausalLM | LlavaForConditionalGeneration,
+    tokenizer_or_processor: PreTrainedTokenizer | ProcessorMixin | AutoTokenizer | AutoProcessor,
+    *,
+    device_dtype: Any = None,
+) -> Transformers | TransformersMultiModal: ...

stubs/outlines/models/vllm_offline.pyi

@@ -0,0 +1,10 @@
+from typing import TYPE_CHECKING
+
+from outlines.models.base import Model
+
+if TYPE_CHECKING:
+    from vllm import LLM
+
+class VLLMOffline(Model): ...
+
+def from_vllm_offline(model: LLM) -> VLLMOffline: ...

stubs/transformers/__init__.pyi

@@ -0,0 +1,26 @@
+from typing import Any
+
+from typing_extensions import Self
+
+from . import modeling_utils, processing_utils, tokenization_utils
+from .modeling_utils import PreTrainedModel
+from .processing_utils import ProcessorMixin
+from .tokenization_utils import PreTrainedTokenizer
+
+class AutoModelForCausalLM(PreTrainedModel):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+class AutoTokenizer(PreTrainedTokenizer):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+class AutoProcessor(ProcessorMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+class LlavaForConditionalGeneration(PreTrainedModel):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+def from_pretrained(*args: Any, **kwargs: Any) -> Any: ...

stubs/transformers/modeling_utils.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+class PreTrainedModel: ...

stubs/transformers/processing_utils.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+class ProcessorMixin: ...

stubs/transformers/tokenization_utils.pyi

@@ -0,0 +1,8 @@
+from typing import Any
+
+class PreTrainedTokenizer:
+    chat_template: str | None
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+
+class ProcessorMixin: ...

stubs/vllm/__init__.pyi

@@ -0,0 +1,4 @@
+from typing import Any
+
+class LLM:
+    def __init__(self, model: str, *args: Any, **kwargs: Any) -> None: ...

stubs/vllm/sampling_params.pyi

@@ -0,0 +1,25 @@
+from typing import Any
+
+class SamplingParams:
+    max_tokens: int | None
+    temperature: float | None
+    top_p: float | None
+    seed: int | None
+    presence_penalty: float | None
+    frequency_penalty: float | None
+    logit_bias: dict[int, float] | None
+    extra_body: dict[str, Any] | None
+
+    def __init__(
+        self,
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        seed: int | None = None,
+        presence_penalty: float | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[int, float] | None = None,
+        extra_body: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None: ...