Harden Azure ML url validation (#88)

This PR hardens the URL validation for the AzureML LLM endpoint. 

It also migrates the validator method to the Pydantic V2 style.

Author: tonybaloney

libs/community/langchain_community/llms/azureml_endpoint.py

@@ -4,12 +4,20 @@
 from abc import abstractmethod
 from enum import Enum
 from typing import Any, Dict, List, Mapping, Optional
+from urllib.parse import urlparse
 
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.llms import BaseLLM
 from langchain_core.outputs import Generation, LLMResult
 from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
-from pydantic import BaseModel, ConfigDict, SecretStr, model_validator, validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    SecretStr,
+    field_validator,
+    model_validator,
+    validator,
+)
 
 DEFAULT_TIMEOUT = 50
 
@@ -431,43 +439,42 @@ def validate_content_formatter(
             )
         return field_value
 
-    @validator("endpoint_url")
-    def validate_endpoint_url(cls, field_value: Any) -> str:
+    @field_validator("endpoint_url", mode="after")
+    @classmethod
+    def validate_endpoint_url(cls, value: str) -> str:
         """Validate that endpoint url is complete."""
-        if field_value.endswith("/"):
-            field_value = field_value[:-1]
-        if field_value.endswith("inference.ml.azure.com"):
+        if value.endswith("/"):  # trim trailing slash
+            value = value[:-1]
+        url = urlparse(value)
+        if not url.path or url.path == "/":
             raise ValueError(
                 "`endpoint_url` should contain the full invocation URL including "
                 "`/score` for `endpoint_api_type='dedicated'` or `/completions` "
                 "or `/models/chat/completions` "
                 "for `endpoint_api_type='serverless'`"
             )
-        return field_value
+        return value
 
     @validator("endpoint_api_type")
     def validate_endpoint_api_type(
         cls, field_value: Any, values: Dict
     ) -> AzureMLEndpointApiType:
         """Validate that endpoint api type is compatible with the URL format."""
-        endpoint_url = values.get("endpoint_url")
+        endpoint_url = urlparse(values.get("endpoint_url"))
         if (
-            (
-                field_value == AzureMLEndpointApiType.dedicated
-                or field_value == AzureMLEndpointApiType.realtime
-            )
-            and not endpoint_url.endswith("/score")  # type: ignore[union-attr]
-        ):
+            field_value == AzureMLEndpointApiType.dedicated
+            or field_value == AzureMLEndpointApiType.realtime
+        ) and not endpoint_url.path == "/score":
             raise ValueError(
                 "Endpoints of type `dedicated` should follow the format "
                 "`https://<your-endpoint>.<your_region>.inference.ml.azure.com/score`."
                 " If your endpoint URL ends with `/completions` or"
                 "`/models/chat/completions`,"
                 "use `endpoint_api_type='serverless'` instead."
             )
-        if field_value == AzureMLEndpointApiType.serverless and not (
-            endpoint_url.endswith("/completions")  # type: ignore[union-attr]
-            or endpoint_url.endswith("/models/chat/completions")  # type: ignore[union-attr]
+        if (
+            field_value == AzureMLEndpointApiType.serverless
+            and endpoint_url.path not in ["/completions", "/models/chat/completions"]
         ):
             raise ValueError(
                 "Endpoints of type `serverless` should follow the format "

libs/community/tests/integration_tests/llms/test_azureml_endpoint.py

@@ -127,12 +127,16 @@ def format_response_payload(self, output: bytes) -> str:  # type: ignore[overrid
         llm.invoke("Foo")
 
 
-def test_incorrect_url() -> None:
+@pytest.mark.parametrize(
+    "endpoint_url",
+    ["https://endpoint.inference.com", "https://endpoint.inference.com/"],
+)
+def test_incorrect_url(endpoint_url: str) -> None:
     """Testing AzureML Endpoint for an incorrect URL"""
     with pytest.raises(ValidationError):
         llm = AzureMLOnlineEndpoint(
             endpoint_api_key=os.getenv("OSS_ENDPOINT_API_KEY"),  # type: ignore[arg-type]
-            endpoint_url="https://endpoint.inference.com",
+            endpoint_url=endpoint_url,
             deployment_name=os.getenv("OSS_DEPLOYMENT_NAME"),  # type: ignore[arg-type]
             content_formatter=OSSContentFormatter(),
         )