|
7 | 7 | from unstructured_client._hooks import HookContext |
8 | 8 | from unstructured_client.models import errors, operations, shared |
9 | 9 | from unstructured_client.types import BaseModel, OptionalNullable, UNSET |
| 10 | +from unstructured_client._hooks.custom.clean_server_url_hook import clean_server_url |
10 | 11 |
|
11 | 12 |
|
12 | 13 | class PartitionAcceptEnum(str, Enum): |
@@ -51,7 +52,8 @@ def partition( |
51 | 52 | # Note(austin): Add a custom check to handle the default server URL |
52 | 53 | # The SDK globally defaults to the platform URL. |
53 | 54 | # If that hasn't changed, we need to switch to the partition url here. |
54 | | - if "https://platform.unstructuredapp.io" in base_url: |
| 55 | + base_url = clean_server_url(base_url) |
| 56 | + if base_url == "https://platform.unstructuredapp.io": |
55 | 57 | base_url = "https://api.unstructuredapp.io" |
56 | 58 |
|
57 | 59 | if not isinstance(request, BaseModel): |
@@ -184,7 +186,8 @@ async def partition_async( |
184 | 186 | # Note(austin): Add a custom check to handle the default server URL |
185 | 187 | # The SDK globally defaults to the platform URL. |
186 | 188 | # If that hasn't changed, we need to switch to the partition url here. |
187 | | - if "https://platform.unstructuredapp.io" in base_url: |
| 189 | + base_url = clean_server_url(base_url) |
| 190 | + if base_url == "https://platform.unstructuredapp.io": |
188 | 191 | base_url = "https://api.unstructuredapp.io" |
189 | 192 |
|
190 | 193 | if not isinstance(request, BaseModel): |
|
0 commit comments