fix: fixed custom hooks incorrectly handle URLs

Author: pawel-kmiecik

_test_unstructured_client/unit/test_request_utils.py

@@ -2,7 +2,11 @@
 import httpx
 import pytest
 
-from unstructured_client._hooks.custom.request_utils import create_pdf_chunk_request_params, get_multipart_stream_fields
+from unstructured_client._hooks.custom.request_utils import (
+    create_pdf_chunk_request_params,
+    get_base_url,
+    get_multipart_stream_fields,
+)
 from unstructured_client.models import shared
 
 
@@ -70,3 +74,16 @@ def test_multipart_stream_fields_raises_value_error_when_filename_is_not_set():
 def test_create_pdf_chunk_request_params(input_form_data, page_number, expected_form_data):
     form_data = create_pdf_chunk_request_params(input_form_data, page_number)
     assert form_data == expected_form_data
+
+
+@pytest.mark.parametrize(
+    ("url", "expected_base_url"),
+    [
+        ("https://api.unstructuredapp.io/general/v0/general", "https://api.unstructuredapp.io"),
+        ("https://api.unstructuredapp.io/general/v0/general?some_param=23", "https://api.unstructuredapp.io"),
+        ("http://localhost:3000/general/v0/general", "http://localhost:3000"),
+    ],
+)
+def test_get_base_url(url: str, expected_base_url: str):
+    assert get_base_url(url) == expected_base_url
+

src/unstructured_client/_hooks/custom/request_utils.py

@@ -5,8 +5,10 @@
 import json
 import logging
 from typing import Tuple, Any, BinaryIO
+from urllib.parse import urlparse
 
 import httpx
+from httpx import URL
 from httpx._multipart import DataField, FileField
 
 from unstructured_client._hooks.custom.common import UNSTRUCTURED_CLIENT_LOGGER_NAME
@@ -224,3 +226,15 @@ def create_response(elements: list) -> httpx.Response:
     response.headers.update({"Content-Length": content_length})
     setattr(response, "_content", content)
     return response
+
+def get_base_url(url: str | URL) -> str:
+    """Extracts the base URL from the given URL.
+
+    Args:
+        url: The URL.
+
+    Returns:
+        The base URL.
+    """
+    parsed_url = urlparse(str(url))
+    return f"{parsed_url.scheme}://{parsed_url.netloc}"

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -31,6 +31,7 @@
     PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
     PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
 )
+from unstructured_client._hooks.custom.request_utils import get_base_url
 from unstructured_client._hooks.types import (
     AfterErrorContext,
     AfterErrorHook,
@@ -156,7 +157,8 @@ class SplitPdfHook(SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorH
 
     def __init__(self) -> None:
         self.client: Optional[HttpClient] = None
-        self.base_url: Optional[str] = None
+        self.partition_base_url: Optional[str] = None
+        self.is_partition_request: bool = False
         self.async_client: Optional[AsyncHttpClient] = None
         self.coroutines_to_execute: dict[
             str, list[partial[Coroutine[Any, Any, httpx.Response]]]
@@ -212,7 +214,9 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
         #         return await self.base_transport.handle_async_request(request)
 
         # Instead, save the base url so we can use it for our dummy request
-        self.base_url = base_url
+        # As this can be overwritten with Platform API URL, we need to get it again in
+        # `before_request` hook from the request object as the real URL is not available here.
+        self.partition_base_url = base_url
 
         # Explicit cast to httpx.Client to avoid a typing error
         httpx_client = cast(httpx.Client, client)
@@ -246,6 +250,16 @@ def before_request(
             Union[httpx.PreparedRequest, Exception]: If `splitPdfPage` is set to `true`,
             the last page request; otherwise, the original request.
         """
+
+        # Actually the general.partition operation overwrites the default client's base url (as
+        # the platform operations do). Here we need to get the base url from the request object.
+        if hook_ctx.operation_id == "partition":
+            self.partition_base_url = get_base_url(request.url)
+            self.is_partition_request = True
+        else:
+            self.is_partition_request = False
+            return request
+
         if self.client is None:
             logger.warning("HTTP client not accessible! Continuing without splitting.")
             return request
@@ -391,7 +405,7 @@ def before_request(
         # dummy_request = httpx.Request("GET",  "http://no-op")
         return httpx.Request(
             "GET",
-            f"{self.base_url}/general/docs",
+            f"{self.partition_base_url}/general/docs",
             headers={"operation_id": operation_id},
         )
 
@@ -644,6 +658,9 @@ def after_success(
             combined response object; otherwise, the original response. Can return
             exception if it ocurred during the execution.
         """
+        if not self.is_partition_request:
+            return response
+
         # Grab the correct id out of the dummy request
         operation_id = response.request.headers.get("operation_id")