pydantic-ai-slim 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

pydantic_ai/models/ollama.py

@@ -0,0 +1,116 @@
+from __future__ import annotations as _annotations
+
+from dataclasses import dataclass
+from typing import Literal, Union
+
+from httpx import AsyncClient as AsyncHTTPClient
+
+from ..tools import ToolDefinition
+from . import (
+    AgentModel,
+    Model,
+    cached_async_http_client,
+)
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as e:
+    raise ImportError(
+        'Please install `openai` to use the OpenAI model, '
+        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+    ) from e
+
+
+from .openai import OpenAIModel
+
+CommonOllamaModelNames = Literal[
+    'codellama',
+    'gemma',
+    'gemma2',
+    'llama3',
+    'llama3.1',
+    'llama3.2',
+    'llama3.2-vision',
+    'llama3.3',
+    'mistral',
+    'mistral-nemo',
+    'mixtral',
+    'phi3',
+    'qwq',
+    'qwen',
+    'qwen2',
+    'qwen2.5',
+    'starcoder2',
+]
+"""This contains just the most common ollama models.
+
+For a full list see [ollama.com/library](https://ollama.com/library).
+"""
+OllamaModelName = Union[CommonOllamaModelNames, str]
+"""Possible ollama models.
+
+Since Ollama supports hundreds of models, we explicitly list the most models but
+allow any name in the type hints.
+"""
+
+
+@dataclass(init=False)
+class OllamaModel(Model):
+    """A model that implements Ollama using the OpenAI API.
+
+    Internally, this uses the [OpenAI Python client](https://github.com/openai/openai-python) to interact with the Ollama server.
+
+    Apart from `__init__`, all methods are private or match those of the base class.
+    """
+
+    model_name: OllamaModelName
+    openai_model: OpenAIModel
+
+    def __init__(
+        self,
+        model_name: OllamaModelName,
+        *,
+        base_url: str | None = 'http://localhost:11434/v1/',
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ):
+        """Initialize an Ollama model.
+
+        Ollama has built-in compatability for the OpenAI chat completions API ([source](https://ollama.com/blog/openai-compatibility)), so we reuse the
+        [`OpenAIModel`][pydantic_ai.models.openai.OpenAIModel] here.
+
+        Args:
+            model_name: The name of the Ollama model to use. List of models available [here](https://ollama.com/library)
+                You must first download the model (`ollama pull <MODEL-NAME>`) in order to use the model
+            base_url: The base url for the ollama requests. The default value is the ollama default
+            openai_client: An existing
+                [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
+                client to use, if provided, `base_url` and `http_client` must be `None`.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        self.model_name = model_name
+        if openai_client is not None:
+            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
+            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            self.openai_model = OpenAIModel(model_name=model_name, openai_client=openai_client)
+        else:
+            # API key is not required for ollama but a value is required to create the client
+            http_client_ = http_client or cached_async_http_client()
+            oai_client = AsyncOpenAI(base_url=base_url, api_key='ollama', http_client=http_client_)
+            self.openai_model = OpenAIModel(model_name=model_name, openai_client=oai_client)
+
+    async def agent_model(
+        self,
+        *,
+        function_tools: list[ToolDefinition],
+        allow_text_result: bool,
+        result_tools: list[ToolDefinition],
+    ) -> AgentModel:
+        return await self.openai_model.agent_model(
+            function_tools=function_tools,
+            allow_text_result=allow_text_result,
+            result_tools=result_tools,
+        )
+
+    def name(self) -> str:
+        return f'ollama:{self.model_name}'

pydantic_ai/models/openai.py

@@ -1,28 +1,34 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Literal, overload
+from itertools import chain
+from typing import Literal, Union, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
 from .. import UnexpectedModelBehavior, _utils, result
+from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
     ArgsJson,
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
 )
 from ..result import Cost
+from ..settings import ModelSettings
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -37,11 +43,17 @@ try:
     from openai.types import ChatModel, chat
     from openai.types.chat import ChatCompletionChunk
     from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
-except ImportError as e:
+except ImportError as _import_error:
     raise ImportError(
         'Please install `openai` to use the OpenAI model, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai[openai]'`"
-    ) from e
+        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+    ) from _import_error
+
+OpenAIModelName = Union[ChatModel, str]
+"""
+Using this more broad type for the model name instead of the ChatModel definition
+allows this model to be used more easily with other model types (ie, Ollama)
+"""
 
 
 @dataclass(init=False)
@@ -53,13 +65,14 @@ class OpenAIModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    model_name: ChatModel
+    model_name: OpenAIModelName
     client: AsyncOpenAI = field(repr=False)
 
     def __init__(
         self,
-        model_name: ChatModel,
+        model_name: OpenAIModelName,
         *,
+        base_url: str | None = None,
         api_key: str | None = None,
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
@@ -70,32 +83,36 @@ class OpenAIModel(Model):
             model_name: The name of the OpenAI model to use. List of model names available
                 [here](https://github.com/openai/openai-python/blob/v1.54.3/src/openai/types/chat_model.py#L7)
                 (Unfortunately, despite being ask to do so, OpenAI do not provide `.inv` files for their API).
+            base_url: The base url for the OpenAI requests. If not provided, the `OPENAI_BASE_URL` environment variable
+                will be used if available. Otherwise, defaults to OpenAI's base url.
             api_key: The API key to use for authentication, if not provided, the `OPENAI_API_KEY` environment variable
                 will be used if available.
             openai_client: An existing
                 [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
-                client to use, if provided, `api_key` and `http_client` must be `None`.
+                client to use. If provided, `base_url`, `api_key`, and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
-        self.model_name: ChatModel = model_name
+        self.model_name: OpenAIModelName = model_name
         if openai_client is not None:
             assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
             assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
             self.client = openai_client
         elif http_client is not None:
-            self.client = AsyncOpenAI(api_key=api_key, http_client=http_client)
+            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
         else:
-            self.client = AsyncOpenAI(api_key=api_key, http_client=cached_async_http_client())
+            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
         check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools.values()]
-        if result_tools is not None:
+        tools = [self._map_tool_definition(r) for r in function_tools]
+        if result_tools:
             tools += [self._map_tool_definition(r) for r in result_tools]
         return OpenAIAgentModel(
             self.client,
@@ -108,13 +125,13 @@ class OpenAIModel(Model):
         return f'openai:{self.model_name}'
 
     @staticmethod
-    def _map_tool_definition(f: AbstractToolDefinition) -> chat.ChatCompletionToolParam:
+    def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam:
         return {
             'type': 'function',
             'function': {
                 'name': f.name,
                 'description': f.description,
-                'parameters': f.json_schema,
+                'parameters': f.parameters_json_schema,
             },
         }
 
@@ -124,32 +141,38 @@ class OpenAIAgentModel(AgentModel):
     """Implementation of `AgentModel` for OpenAI models."""
 
     client: AsyncOpenAI
-    model_name: ChatModel
+    model_name: OpenAIModelName
     allow_text_result: bool
     tools: list[chat.ChatCompletionToolParam]
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, result.Cost]:
-        response = await self._completions_create(messages, False)
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Cost]:
+        response = await self._completions_create(messages, False, model_settings)
         return self._process_response(response), _map_cost(response)
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        response = await self._completions_create(messages, True)
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        response = await self._completions_create(messages, True, model_settings)
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _completions_create(
-        self, messages: list[Message], stream: Literal[True]
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
-    async def _completions_create(self, messages: list[Message], stream: Literal[False]) -> chat.ChatCompletion:
+    async def _completions_create(
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+    ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[Message], stream: bool
+        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
         # standalone function to make it easier to override
         if not self.tools:
@@ -159,7 +182,10 @@ class OpenAIAgentModel(AgentModel):
         else:
             tool_choice = 'auto'
 
-        openai_messages = [self._map_message(m) for m in messages]
+        openai_messages = list(chain(*(self._map_message(m) for m in messages)))
+
+        model_settings = model_settings or {}
+
         return await self.client.chat.completions.create(
             model=self.model_name,
             messages=openai_messages,
@@ -169,93 +195,104 @@ class OpenAIAgentModel(AgentModel):
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
             stream_options={'include_usage': True} if stream else NOT_GIVEN,
+            max_tokens=model_settings.get('max_tokens', NOT_GIVEN),
+            temperature=model_settings.get('temperature', NOT_GIVEN),
+            top_p=model_settings.get('top_p', NOT_GIVEN),
+            timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
     @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelAnyResponse:
+    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
+        items: list[ModelResponsePart] = []
+        if choice.message.content is not None:
+            items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
-            return ModelStructuredResponse(
-                [ToolCall.from_json(c.function.name, c.function.arguments, c.id) for c in choice.message.tool_calls],
-                timestamp=timestamp,
-            )
-        else:
-            assert choice.message.content is not None, choice
-            return ModelTextResponse(choice.message.content, timestamp=timestamp)
+            for c in choice.message.tool_calls:
+                items.append(ToolCallPart.from_json(c.function.name, c.function.arguments, c.id))
+        return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
     async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
-        try:
-            first_chunk = await response.__anext__()
-        except StopAsyncIteration as e:  # pragma: no cover
-            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
-        timestamp = datetime.fromtimestamp(first_chunk.created, tz=timezone.utc)
-        delta = first_chunk.choices[0].delta
-        start_cost = _map_cost(first_chunk)
-
-        # the first chunk may only contain `role`, so we iterate until we get either `tool_calls` or `content`
-        while delta.tool_calls is None and delta.content is None:
+        timestamp: datetime | None = None
+        start_cost = Cost()
+        # the first chunk may contain enough information so we iterate until we get either `tool_calls` or `content`
+        while True:
             try:
-                next_chunk = await response.__anext__()
+                chunk = await response.__anext__()
             except StopAsyncIteration as e:
                 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
-            delta = next_chunk.choices[0].delta
-            start_cost += _map_cost(next_chunk)
 
-        if delta.content is not None:
-            return OpenAIStreamTextResponse(delta.content, response, timestamp, start_cost)
+            timestamp = timestamp or datetime.fromtimestamp(chunk.created, tz=timezone.utc)
+            start_cost += _map_cost(chunk)
+
+            if chunk.choices:
+                delta = chunk.choices[0].delta
+
+                if delta.content is not None:
+                    return OpenAIStreamTextResponse(delta.content, response, timestamp, start_cost)
+                elif delta.tool_calls is not None:
+                    return OpenAIStreamStructuredResponse(
+                        response,
+                        {c.index: c for c in delta.tool_calls},
+                        timestamp,
+                        start_cost,
+                    )
+                # else continue until we get either delta.content or delta.tool_calls
+
+    @classmethod
+    def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
+        """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
+        if isinstance(message, ModelRequest):
+            yield from cls._map_user_message(message)
+        elif isinstance(message, ModelResponse):
+            texts: list[str] = []
+            tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
+            for item in message.parts:
+                if isinstance(item, TextPart):
+                    texts.append(item.content)
+                elif isinstance(item, ToolCallPart):
+                    tool_calls.append(_map_tool_call(item))
+                else:
+                    assert_never(item)
+            message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
+            if texts:
+                # Note: model responses from this model should only have one text item, so the following
+                # shouldn't merge multiple texts into one unless you switch models between runs:
+                message_param['content'] = '\n\n'.join(texts)
+            if tool_calls:
+                message_param['tool_calls'] = tool_calls
+            yield message_param
         else:
-            assert delta.tool_calls is not None, f'Expected delta with tool_calls, got {delta}'
-            return OpenAIStreamStructuredResponse(
-                response,
-                {c.index: c for c in delta.tool_calls},
-                timestamp,
-                start_cost,
-            )
+            assert_never(message)
 
-    @staticmethod
-    def _map_message(message: Message) -> chat.ChatCompletionMessageParam:
-        """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
-        if message.role == 'system':
-            # SystemPrompt ->
-            return chat.ChatCompletionSystemMessageParam(role='system', content=message.content)
-        elif message.role == 'user':
-            # UserPrompt ->
-            return chat.ChatCompletionUserMessageParam(role='user', content=message.content)
-        elif message.role == 'tool-return':
-            # ToolReturn ->
-            return chat.ChatCompletionToolMessageParam(
-                role='tool',
-                tool_call_id=_guard_tool_id(message),
-                content=message.model_response_str(),
-            )
-        elif message.role == 'retry-prompt':
-            # RetryPrompt ->
-            if message.tool_name is None:
-                return chat.ChatCompletionUserMessageParam(role='user', content=message.model_response())
-            else:
-                return chat.ChatCompletionToolMessageParam(
+    @classmethod
+    def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
+        for part in message.parts:
+            if isinstance(part, SystemPromptPart):
+                yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
+            elif isinstance(part, UserPromptPart):
+                yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+            elif isinstance(part, ToolReturnPart):
+                yield chat.ChatCompletionToolMessageParam(
                     role='tool',
-                    tool_call_id=_guard_tool_id(message),
-                    content=message.model_response(),
+                    tool_call_id=_guard_tool_call_id(t=part, model_source='OpenAI'),
+                    content=part.model_response_str(),
                 )
-        elif message.role == 'model-text-response':
-            # ModelTextResponse ->
-            return chat.ChatCompletionAssistantMessageParam(role='assistant', content=message.content)
-        elif message.role == 'model-structured-response':
-            assert (
-                message.role == 'model-structured-response'
-            ), f'Expected role to be "llm-tool-calls", got {message.role}'
-            # ModelStructuredResponse ->
-            return chat.ChatCompletionAssistantMessageParam(
-                role='assistant',
-                tool_calls=[_map_tool_call(t) for t in message.calls],
-            )
-        else:
-            assert_never(message)
+            elif isinstance(part, RetryPromptPart):
+                if part.tool_name is None:
+                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response())
+                else:
+                    yield chat.ChatCompletionToolMessageParam(
+                        role='tool',
+                        tool_call_id=_guard_tool_call_id(t=part, model_source='OpenAI'),
+                        content=part.model_response(),
+                    )
+            else:
+                assert_never(part)
 
 
 @dataclass
@@ -330,14 +367,14 @@ class OpenAIStreamStructuredResponse(StreamStructuredResponse):
             else:
                 self._delta_tool_calls[new.index] = new
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        calls: list[ToolCall] = []
+    def get(self, *, final: bool = False) -> ModelResponse:
+        items: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    calls.append(ToolCall.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_json(f.name, f.arguments, c.id))
 
-        return ModelStructuredResponse(calls, timestamp=self._timestamp)
+        return ModelResponse(items, timestamp=self._timestamp)
 
     def cost(self) -> Cost:
         return self._cost
@@ -346,16 +383,10 @@ class OpenAIStreamStructuredResponse(StreamStructuredResponse):
         return self._timestamp
 
 
-def _guard_tool_id(t: ToolCall | ToolReturn | RetryPrompt) -> str:
-    """Type guard that checks a `tool_id` is not None both for static typing and runtime."""
-    assert t.tool_id is not None, f'OpenAI requires `tool_id` to be set: {t}'
-    return t.tool_id
-
-
-def _map_tool_call(t: ToolCall) -> chat.ChatCompletionMessageToolCallParam:
+def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
     assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
     return chat.ChatCompletionMessageToolCallParam(
-        id=_guard_tool_id(t),
+        id=_guard_tool_call_id(t=t, model_source='OpenAI'),
         type='function',
         function={'name': t.tool_name, 'arguments': t.args.args_json},
     )