address pr comments

NumberAI · Jul 30, 2024 · bea7a4a · bea7a4a
1 parent cac15c9
commit bea7a4a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 12 deletions.
diff --git 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
@@ -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")