tests

Author: DouweM

tests/models/mock_openai.py

@@ -9,14 +9,16 @@
 from .mock_async_stream import MockAsyncStream
 
 with try_import() as imports_successful:
-    from openai import NOT_GIVEN, AsyncOpenAI
+    from openai import AsyncOpenAI
     from openai.types import chat, responses
     from openai.types.chat.chat_completion import Choice, ChoiceLogprobs
     from openai.types.chat.chat_completion_message import ChatCompletionMessage
     from openai.types.completion_usage import CompletionUsage
     from openai.types.responses.response import ResponseUsage
     from openai.types.responses.response_output_item import ResponseOutputItem
 
+    from pydantic_ai.models.openai import NOT_GIVEN, OMIT
+
     MockChatCompletion = chat.ChatCompletion | Exception
     MockChatCompletionChunk = chat.ChatCompletionChunk | Exception
     MockResponse = responses.Response | Exception
@@ -50,7 +52,7 @@ def create_mock_stream(
     async def chat_completions_create(  # pragma: lax no cover
         self, *_args: Any, stream: bool = False, **kwargs: Any
     ) -> chat.ChatCompletion | MockAsyncStream[MockChatCompletionChunk]:
-        self.chat_completion_kwargs.append({k: v for k, v in kwargs.items() if v is not NOT_GIVEN})
+        self.chat_completion_kwargs.append({k: v for k, v in kwargs.items() if v not in (NOT_GIVEN, OMIT)})
 
         if stream:
             assert self.stream is not None, 'you can only used `stream=True` if `stream` is provided'