feat(client): add header OpenAI-Project (#1320)

openai · Apr 16, 2024 · 3408b5d · 3408b5d
1 parent 1a83130
commit 3408b5d
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 0 deletions.
diff --git a/src/openai/__init__.py b/src/openai/__init__.py
@@ -108,6 +108,8 @@
 
 organization: str | None = None
 
+project: str | None = None
+
 base_url: str | _httpx.URL | None = None
 
 timeout: float | Timeout | None = DEFAULT_TIMEOUT
@@ -159,6 +161,17 @@ def organization(self, value: str | None) -> None:  # type: ignore
 
         organization = value
 
+    @property  # type: ignore
+    @override
+    def project(self) -> str | None:
+        return project
+
+    @project.setter  # type: ignore
+    def project(self, value: str | None) -> None:  # type: ignore
+        global project
+
+        project = value
+
     @property
     @override
     def base_url(self) -> _httpx.URL:
@@ -310,6 +323,7 @@ def _load_client() -> OpenAI:  # type: ignore[reportUnusedFunction]
         _client = _ModuleClient(
             api_key=api_key,
             organization=organization,
+            project=project,
             base_url=base_url,
             timeout=timeout,
             max_retries=max_retries,

diff --git a/src/openai/_client.py b/src/openai/_client.py
@@ -64,12 +64,14 @@ class OpenAI(SyncAPIClient):
     # client options
     api_key: str
     organization: str | None
+    project: str | None
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
@@ -94,6 +96,7 @@ def __init__(
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
         - `api_key` from `OPENAI_API_KEY`
         - `organization` from `OPENAI_ORG_ID`
+        - `project` from `OPENAI_PROJECT_ID`
         """
         if api_key is None:
             api_key = os.environ.get("OPENAI_API_KEY")
@@ -107,6 +110,10 @@ def __init__(
             organization = os.environ.get("OPENAI_ORG_ID")
         self.organization = organization
 
+        if project is None:
+            project = os.environ.get("OPENAI_PROJECT_ID")
+        self.project = project
+
         if base_url is None:
             base_url = os.environ.get("OPENAI_BASE_URL")
         if base_url is None:
@@ -157,6 +164,7 @@ def default_headers(self) -> dict[str, str | Omit]:
             **super().default_headers,
             "X-Stainless-Async": "false",
             "OpenAI-Organization": self.organization if self.organization is not None else Omit(),
+            "OpenAI-Project": self.project if self.project is not None else Omit(),
             **self._custom_headers,
         }
 
@@ -165,6 +173,7 @@ def copy(
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.Client | None = None,
@@ -200,6 +209,7 @@ def copy(
         return self.__class__(
             api_key=api_key or self.api_key,
             organization=organization or self.organization,
+            project=project or self.project,
             base_url=base_url or self.base_url,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,
@@ -266,12 +276,14 @@ class AsyncOpenAI(AsyncAPIClient):
     # client options
     api_key: str
     organization: str | None
+    project: str | None
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
@@ -296,6 +308,7 @@ def __init__(
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
         - `api_key` from `OPENAI_API_KEY`
         - `organization` from `OPENAI_ORG_ID`
+        - `project` from `OPENAI_PROJECT_ID`
         """
         if api_key is None:
             api_key = os.environ.get("OPENAI_API_KEY")
@@ -309,6 +322,10 @@ def __init__(
             organization = os.environ.get("OPENAI_ORG_ID")
         self.organization = organization
 
+        if project is None:
+            project = os.environ.get("OPENAI_PROJECT_ID")
+        self.project = project
+
         if base_url is None:
             base_url = os.environ.get("OPENAI_BASE_URL")
         if base_url is None:
@@ -359,6 +376,7 @@ def default_headers(self) -> dict[str, str | Omit]:
             **super().default_headers,
             "X-Stainless-Async": f"async:{get_async_library()}",
             "OpenAI-Organization": self.organization if self.organization is not None else Omit(),
+            "OpenAI-Project": self.project if self.project is not None else Omit(),
             **self._custom_headers,
         }
 
@@ -367,6 +385,7 @@ def copy(
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.AsyncClient | None = None,
@@ -402,6 +421,7 @@ def copy(
         return self.__class__(
             api_key=api_key or self.api_key,
             organization=organization or self.organization,
+            project=project or self.project,
             base_url=base_url or self.base_url,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,

diff --git a/src/openai/lib/azure.py b/src/openai/lib/azure.py
@@ -130,6 +130,7 @@ def __init__(
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AzureADTokenProvider | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
@@ -143,6 +144,7 @@ def __init__(
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
         - `api_key` from `AZURE_OPENAI_API_KEY`
         - `organization` from `OPENAI_ORG_ID`
+        - `project` from `OPENAI_PROJECT_ID`
         - `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
         - `api_version` from `OPENAI_API_VERSION`
         - `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`
@@ -205,6 +207,7 @@ def __init__(
         super().__init__(
             api_key=api_key,
             organization=organization,
+            project=project,
             base_url=base_url,
             timeout=timeout,
             max_retries=max_retries,
@@ -223,6 +226,7 @@ def copy(
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         api_version: str | None = None,
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AzureADTokenProvider | None = None,
@@ -242,6 +246,7 @@ def copy(
         return super().copy(
             api_key=api_key,
             organization=organization,
+            project=project,
             base_url=base_url,
             timeout=timeout,
             http_client=http_client,
@@ -306,6 +311,7 @@ def __init__(
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
         organization: str | None = None,
+        project: str | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -325,6 +331,7 @@ def __init__(
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
         organization: str | None = None,
+        project: str | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -344,6 +351,7 @@ def __init__(
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
         organization: str | None = None,
+        project: str | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -363,6 +371,7 @@ def __init__(
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
         organization: str | None = None,
+        project: str | None = None,
         base_url: str | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
@@ -376,6 +385,7 @@ def __init__(
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
         - `api_key` from `AZURE_OPENAI_API_KEY`
         - `organization` from `OPENAI_ORG_ID`
+        - `project` from `OPENAI_PROJECT_ID`
         - `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
         - `api_version` from `OPENAI_API_VERSION`
         - `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`
@@ -438,6 +448,7 @@ def __init__(
         super().__init__(
             api_key=api_key,
             organization=organization,
+            project=project,
             base_url=base_url,
             timeout=timeout,
             max_retries=max_retries,
@@ -456,6 +467,7 @@ def copy(
         *,
         api_key: str | None = None,
         organization: str | None = None,
+        project: str | None = None,
         api_version: str | None = None,
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
@@ -475,6 +487,7 @@ def copy(
         return super().copy(
             api_key=api_key,
             organization=organization,
+            project=project,
             base_url=base_url,
             timeout=timeout,
             http_client=http_client,

diff --git a/tests/test_module_client.py b/tests/test_module_client.py
@@ -16,6 +16,7 @@ def reset_state() -> None:
     openai._reset_client()
     openai.api_key = None or "My API Key"
     openai.organization = None
+    openai.project = None
     openai.base_url = None
     openai.timeout = DEFAULT_TIMEOUT
     openai.max_retries = DEFAULT_MAX_RETRIES