diff --git a/python_modules/libraries/dagster-openai/dagster_openai/resources.py b/python_modules/libraries/dagster-openai/dagster_openai/resources.py index 4f5fa8c61f03e..8e6c2e883f174 100644 --- a/python_modules/libraries/dagster-openai/dagster_openai/resources.py +++ b/python_modules/libraries/dagster-openai/dagster_openai/resources.py @@ -180,9 +180,9 @@ def openai_asset(context: AssetExecutionContext, openai: OpenAIResource): """ api_key: str = Field(description=("OpenAI API key. See https://platform.openai.com/api-keys")) - organization: Optional[str] = None - project: Optional[str] = None - base_url: Optional[str] = None + organization: Optional[str] = Field(default=None) + project: Optional[str] = Field(default=None) + base_url: Optional[str] = Field(default=None) _client: Client = PrivateAttr() @@ -215,14 +215,12 @@ def _wrap_with_usage_metadata( def setup_for_execution(self, context: InitResourceContext) -> None: # Set up an OpenAI client based on the API key. - kwargs = {} - if self.organization: - kwargs["organization"] = self.organization - if self.project: - kwargs["project"] = self.project - if self.base_url: - kwargs["base_url"] = self.base_url - self._client = Client(api_key=self.api_key, **kwargs) + self._client = Client( + api_key=self.api_key, + organization=self.organization, + project=self.project, + base_url=self.base_url, + ) @public @contextmanager diff --git a/python_modules/libraries/dagster-openai/dagster_openai_tests/test_resources.py b/python_modules/libraries/dagster-openai/dagster_openai_tests/test_resources.py index e6ec5a969e957..13bf90b25f08f 100644 --- a/python_modules/libraries/dagster-openai/dagster_openai_tests/test_resources.py +++ b/python_modules/libraries/dagster-openai/dagster_openai_tests/test_resources.py @@ -28,7 +28,12 @@ def test_openai_client(mock_client) -> None: mock_context = MagicMock() with openai_resource.get_client(mock_context): - mock_client.assert_called_once_with(api_key="xoxp-1234123412341234-12341234-1234") + mock_client.assert_called_once_with( + api_key="xoxp-1234123412341234-12341234-1234", + organization=None, + project=None, + base_url=None, + ) @patch("dagster_openai.resources.Client")