mirror of https://github.com/hwchase17/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
552 lines
20 KiB
Python
552 lines
20 KiB
Python
import json
|
|
import urllib.request
|
|
import warnings
|
|
from abc import abstractmethod
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Mapping, Optional
|
|
|
|
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.pydantic_v1 import BaseModel, SecretStr, root_validator, validator
|
|
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
|
|
|
|
DEFAULT_TIMEOUT = 50
|
|
|
|
|
|
class AzureMLEndpointClient(object):
|
|
"""AzureML Managed Endpoint client."""
|
|
|
|
def __init__(
|
|
self,
|
|
endpoint_url: str,
|
|
endpoint_api_key: str,
|
|
deployment_name: str = "",
|
|
timeout: int = DEFAULT_TIMEOUT,
|
|
) -> None:
|
|
"""Initialize the class."""
|
|
if not endpoint_api_key or not endpoint_url:
|
|
raise ValueError(
|
|
"""A key/token and REST endpoint should
|
|
be provided to invoke the endpoint"""
|
|
)
|
|
self.endpoint_url = endpoint_url
|
|
self.endpoint_api_key = endpoint_api_key
|
|
self.deployment_name = deployment_name
|
|
self.timeout = timeout
|
|
|
|
def call(
|
|
self,
|
|
body: bytes,
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
**kwargs: Any,
|
|
) -> bytes:
|
|
"""call."""
|
|
|
|
# The azureml-model-deployment header will force the request to go to a
|
|
# specific deployment. Remove this header to have the request observe the
|
|
# endpoint traffic rules.
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": ("Bearer " + self.endpoint_api_key),
|
|
}
|
|
if self.deployment_name != "":
|
|
headers["azureml-model-deployment"] = self.deployment_name
|
|
|
|
req = urllib.request.Request(self.endpoint_url, body, headers)
|
|
response = urllib.request.urlopen(
|
|
req, timeout=kwargs.get("timeout", self.timeout)
|
|
)
|
|
result = response.read()
|
|
return result
|
|
|
|
|
|
class AzureMLEndpointApiType(str, Enum):
|
|
"""Azure ML endpoints API types. Use `dedicated` for models deployed in hosted
|
|
infrastructure (also known as Online Endpoints in Azure Machine Learning),
|
|
or `serverless` for models deployed as a service with a
|
|
pay-as-you-go billing or PTU.
|
|
"""
|
|
|
|
dedicated = "dedicated"
|
|
realtime = "realtime" #: Deprecated
|
|
serverless = "serverless"
|
|
|
|
|
|
class ContentFormatterBase:
|
|
"""Transform request and response of AzureML endpoint to match with
|
|
required schema.
|
|
"""
|
|
|
|
"""
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
class ContentFormatter(ContentFormatterBase):
|
|
content_type = "application/json"
|
|
accepts = "application/json"
|
|
|
|
def format_request_payload(
|
|
self,
|
|
prompt: str,
|
|
model_kwargs: Dict,
|
|
api_type: AzureMLEndpointApiType,
|
|
) -> bytes:
|
|
input_str = json.dumps(
|
|
{
|
|
"inputs": {"input_string": [prompt]},
|
|
"parameters": model_kwargs,
|
|
}
|
|
)
|
|
return str.encode(input_str)
|
|
|
|
def format_response_payload(
|
|
self, output: str, api_type: AzureMLEndpointApiType
|
|
) -> str:
|
|
response_json = json.loads(output)
|
|
return response_json[0]["0"]
|
|
"""
|
|
content_type: Optional[str] = "application/json"
|
|
"""The MIME type of the input data passed to the endpoint"""
|
|
|
|
accepts: Optional[str] = "application/json"
|
|
"""The MIME type of the response data returned from the endpoint"""
|
|
|
|
format_error_msg: str = (
|
|
"Error while formatting response payload for chat model of type "
|
|
" `{api_type}`. Are you using the right formatter for the deployed "
|
|
" model and endpoint type?"
|
|
)
|
|
|
|
@staticmethod
|
|
def escape_special_characters(prompt: str) -> str:
|
|
"""Escapes any special characters in `prompt`"""
|
|
escape_map = {
|
|
"\\": "\\\\",
|
|
'"': '\\"',
|
|
"\b": "\\b",
|
|
"\f": "\\f",
|
|
"\n": "\\n",
|
|
"\r": "\\r",
|
|
"\t": "\\t",
|
|
}
|
|
|
|
# Replace each occurrence of the specified characters with escaped versions
|
|
for escape_sequence, escaped_sequence in escape_map.items():
|
|
prompt = prompt.replace(escape_sequence, escaped_sequence)
|
|
|
|
return prompt
|
|
|
|
@property
|
|
def supported_api_types(self) -> List[AzureMLEndpointApiType]:
|
|
"""Supported APIs for the given formatter. Azure ML supports
|
|
deploying models using different hosting methods. Each method may have
|
|
a different API structure."""
|
|
|
|
return [AzureMLEndpointApiType.dedicated]
|
|
|
|
def format_request_payload(
|
|
self,
|
|
prompt: str,
|
|
model_kwargs: Dict,
|
|
api_type: AzureMLEndpointApiType = AzureMLEndpointApiType.dedicated,
|
|
) -> Any:
|
|
"""Formats the request body according to the input schema of
|
|
the model. Returns bytes or seekable file like object in the
|
|
format specified in the content_type request header.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
def format_response_payload(
|
|
self,
|
|
output: bytes,
|
|
api_type: AzureMLEndpointApiType = AzureMLEndpointApiType.dedicated,
|
|
) -> Generation:
|
|
"""Formats the response body according to the output
|
|
schema of the model. Returns the data type that is
|
|
received from the response.
|
|
"""
|
|
|
|
|
|
class GPT2ContentFormatter(ContentFormatterBase):
|
|
"""Content handler for GPT2"""
|
|
|
|
@property
|
|
def supported_api_types(self) -> List[AzureMLEndpointApiType]:
|
|
return [AzureMLEndpointApiType.dedicated]
|
|
|
|
def format_request_payload( # type: ignore[override]
|
|
self, prompt: str, model_kwargs: Dict, api_type: AzureMLEndpointApiType
|
|
) -> bytes:
|
|
prompt = ContentFormatterBase.escape_special_characters(prompt)
|
|
request_payload = json.dumps(
|
|
{"inputs": {"input_string": [f'"{prompt}"']}, "parameters": model_kwargs}
|
|
)
|
|
return str.encode(request_payload)
|
|
|
|
def format_response_payload( # type: ignore[override]
|
|
self, output: bytes, api_type: AzureMLEndpointApiType
|
|
) -> Generation:
|
|
try:
|
|
choice = json.loads(output)[0]["0"]
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
raise ValueError(self.format_error_msg.format(api_type=api_type)) from e # type: ignore[union-attr]
|
|
return Generation(text=choice)
|
|
|
|
|
|
class OSSContentFormatter(GPT2ContentFormatter):
|
|
"""Deprecated: Kept for backwards compatibility
|
|
|
|
Content handler for LLMs from the OSS catalog."""
|
|
|
|
content_formatter: Any = None
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
warnings.warn(
|
|
"""`OSSContentFormatter` will be deprecated in the future.
|
|
Please use `GPT2ContentFormatter` instead.
|
|
"""
|
|
)
|
|
|
|
|
|
class HFContentFormatter(ContentFormatterBase):
|
|
"""Content handler for LLMs from the HuggingFace catalog."""
|
|
|
|
@property
|
|
def supported_api_types(self) -> List[AzureMLEndpointApiType]:
|
|
return [AzureMLEndpointApiType.dedicated]
|
|
|
|
def format_request_payload( # type: ignore[override]
|
|
self, prompt: str, model_kwargs: Dict, api_type: AzureMLEndpointApiType
|
|
) -> bytes:
|
|
ContentFormatterBase.escape_special_characters(prompt)
|
|
request_payload = json.dumps(
|
|
{"inputs": [f'"{prompt}"'], "parameters": model_kwargs}
|
|
)
|
|
return str.encode(request_payload)
|
|
|
|
def format_response_payload( # type: ignore[override]
|
|
self, output: bytes, api_type: AzureMLEndpointApiType
|
|
) -> Generation:
|
|
try:
|
|
choice = json.loads(output)[0]["0"]["generated_text"]
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
raise ValueError(self.format_error_msg.format(api_type=api_type)) from e # type: ignore[union-attr]
|
|
return Generation(text=choice)
|
|
|
|
|
|
class DollyContentFormatter(ContentFormatterBase):
|
|
"""Content handler for the Dolly-v2-12b model"""
|
|
|
|
@property
|
|
def supported_api_types(self) -> List[AzureMLEndpointApiType]:
|
|
return [AzureMLEndpointApiType.dedicated]
|
|
|
|
def format_request_payload( # type: ignore[override]
|
|
self, prompt: str, model_kwargs: Dict, api_type: AzureMLEndpointApiType
|
|
) -> bytes:
|
|
prompt = ContentFormatterBase.escape_special_characters(prompt)
|
|
request_payload = json.dumps(
|
|
{
|
|
"input_data": {"input_string": [f'"{prompt}"']},
|
|
"parameters": model_kwargs,
|
|
}
|
|
)
|
|
return str.encode(request_payload)
|
|
|
|
def format_response_payload( # type: ignore[override]
|
|
self, output: bytes, api_type: AzureMLEndpointApiType
|
|
) -> Generation:
|
|
try:
|
|
choice = json.loads(output)[0]
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
raise ValueError(self.format_error_msg.format(api_type=api_type)) from e # type: ignore[union-attr]
|
|
return Generation(text=choice)
|
|
|
|
|
|
class CustomOpenAIContentFormatter(ContentFormatterBase):
|
|
"""Content formatter for models that use the OpenAI like API scheme."""
|
|
|
|
@property
|
|
def supported_api_types(self) -> List[AzureMLEndpointApiType]:
|
|
return [AzureMLEndpointApiType.dedicated, AzureMLEndpointApiType.serverless]
|
|
|
|
def format_request_payload( # type: ignore[override]
|
|
self, prompt: str, model_kwargs: Dict, api_type: AzureMLEndpointApiType
|
|
) -> bytes:
|
|
"""Formats the request according to the chosen api"""
|
|
prompt = ContentFormatterBase.escape_special_characters(prompt)
|
|
if api_type in [
|
|
AzureMLEndpointApiType.dedicated,
|
|
AzureMLEndpointApiType.realtime,
|
|
]:
|
|
request_payload = json.dumps(
|
|
{
|
|
"input_data": {
|
|
"input_string": [f'"{prompt}"'],
|
|
"parameters": model_kwargs,
|
|
}
|
|
}
|
|
)
|
|
elif api_type == AzureMLEndpointApiType.serverless:
|
|
request_payload = json.dumps({"prompt": prompt, **model_kwargs})
|
|
else:
|
|
raise ValueError(
|
|
f"`api_type` {api_type} is not supported by this formatter"
|
|
)
|
|
return str.encode(request_payload)
|
|
|
|
def format_response_payload( # type: ignore[override]
|
|
self, output: bytes, api_type: AzureMLEndpointApiType
|
|
) -> Generation:
|
|
"""Formats response"""
|
|
if api_type in [
|
|
AzureMLEndpointApiType.dedicated,
|
|
AzureMLEndpointApiType.realtime,
|
|
]:
|
|
try:
|
|
choice = json.loads(output)[0]["0"]
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
raise ValueError(self.format_error_msg.format(api_type=api_type)) from e # type: ignore[union-attr]
|
|
return Generation(text=choice)
|
|
if api_type == AzureMLEndpointApiType.serverless:
|
|
try:
|
|
choice = json.loads(output)["choices"][0]
|
|
if not isinstance(choice, dict):
|
|
raise TypeError(
|
|
"Endpoint response is not well formed for a chat "
|
|
"model. Expected `dict` but `{type(choice)}` was "
|
|
"received."
|
|
)
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
raise ValueError(self.format_error_msg.format(api_type=api_type)) from e # type: ignore[union-attr]
|
|
return Generation(
|
|
text=choice["text"].strip(),
|
|
generation_info=dict(
|
|
finish_reason=choice.get("finish_reason"),
|
|
logprobs=choice.get("logprobs"),
|
|
),
|
|
)
|
|
raise ValueError(f"`api_type` {api_type} is not supported by this formatter")
|
|
|
|
|
|
class LlamaContentFormatter(CustomOpenAIContentFormatter):
|
|
"""Deprecated: Kept for backwards compatibility
|
|
|
|
Content formatter for Llama."""
|
|
|
|
content_formatter: Any = None
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
warnings.warn(
|
|
"""`LlamaContentFormatter` will be deprecated in the future.
|
|
Please use `CustomOpenAIContentFormatter` instead.
|
|
"""
|
|
)
|
|
|
|
|
|
class AzureMLBaseEndpoint(BaseModel):
|
|
"""Azure ML Online Endpoint models."""
|
|
|
|
endpoint_url: str = ""
|
|
"""URL of pre-existing Endpoint. Should be passed to constructor or specified as
|
|
env var `AZUREML_ENDPOINT_URL`."""
|
|
|
|
endpoint_api_type: AzureMLEndpointApiType = AzureMLEndpointApiType.dedicated
|
|
"""Type of the endpoint being consumed. Possible values are `serverless` for
|
|
pay-as-you-go and `dedicated` for dedicated endpoints. """
|
|
|
|
endpoint_api_key: SecretStr = convert_to_secret_str("")
|
|
"""Authentication Key for Endpoint. Should be passed to constructor or specified as
|
|
env var `AZUREML_ENDPOINT_API_KEY`."""
|
|
|
|
deployment_name: str = ""
|
|
"""Deployment Name for Endpoint. NOT REQUIRED to call endpoint. Should be passed
|
|
to constructor or specified as env var `AZUREML_DEPLOYMENT_NAME`."""
|
|
|
|
timeout: int = DEFAULT_TIMEOUT
|
|
"""Request timeout for calls to the endpoint"""
|
|
|
|
http_client: Any = None #: :meta private:
|
|
|
|
max_retries: int = 1
|
|
|
|
content_formatter: Any = None
|
|
"""The content formatter that provides an input and output
|
|
transform function to handle formats between the LLM and
|
|
the endpoint"""
|
|
|
|
model_kwargs: Optional[dict] = None
|
|
"""Keyword arguments to pass to the model."""
|
|
|
|
@root_validator(pre=True)
|
|
def validate_environ(cls, values: Dict) -> Dict:
|
|
values["endpoint_api_key"] = convert_to_secret_str(
|
|
get_from_dict_or_env(values, "endpoint_api_key", "AZUREML_ENDPOINT_API_KEY")
|
|
)
|
|
values["endpoint_url"] = get_from_dict_or_env(
|
|
values, "endpoint_url", "AZUREML_ENDPOINT_URL"
|
|
)
|
|
values["deployment_name"] = get_from_dict_or_env(
|
|
values, "deployment_name", "AZUREML_DEPLOYMENT_NAME", ""
|
|
)
|
|
values["endpoint_api_type"] = get_from_dict_or_env(
|
|
values,
|
|
"endpoint_api_type",
|
|
"AZUREML_ENDPOINT_API_TYPE",
|
|
AzureMLEndpointApiType.dedicated,
|
|
)
|
|
values["timeout"] = get_from_dict_or_env(
|
|
values,
|
|
"timeout",
|
|
"AZUREML_TIMEOUT",
|
|
str(DEFAULT_TIMEOUT),
|
|
)
|
|
|
|
return values
|
|
|
|
@validator("content_formatter")
|
|
def validate_content_formatter(
|
|
cls, field_value: Any, values: Dict
|
|
) -> ContentFormatterBase:
|
|
"""Validate that content formatter is supported by endpoint type."""
|
|
endpoint_api_type = values.get("endpoint_api_type")
|
|
if endpoint_api_type not in field_value.supported_api_types:
|
|
raise ValueError(
|
|
f"Content formatter f{type(field_value)} is not supported by this "
|
|
f"endpoint. Supported types are {field_value.supported_api_types} "
|
|
f"but endpoint is {endpoint_api_type}."
|
|
)
|
|
return field_value
|
|
|
|
@validator("endpoint_url")
|
|
def validate_endpoint_url(cls, field_value: Any) -> str:
|
|
"""Validate that endpoint url is complete."""
|
|
if field_value.endswith("/"):
|
|
field_value = field_value[:-1]
|
|
if field_value.endswith("inference.ml.azure.com"):
|
|
raise ValueError(
|
|
"`endpoint_url` should contain the full invocation URL including "
|
|
"`/score` for `endpoint_api_type='dedicated'` or `/v1/completions` "
|
|
"or `/v1/chat/completions` for `endpoint_api_type='serverless'`"
|
|
)
|
|
return field_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")
|
|
if (
|
|
(
|
|
field_value == AzureMLEndpointApiType.dedicated
|
|
or field_value == AzureMLEndpointApiType.realtime
|
|
)
|
|
and not endpoint_url.endswith("/score") # type: ignore[union-attr]
|
|
):
|
|
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 `/v1/completions` or"
|
|
"`/v1/chat/completions`, use `endpoint_api_type='serverless'` instead."
|
|
)
|
|
if field_value == AzureMLEndpointApiType.serverless and not (
|
|
endpoint_url.endswith("/v1/completions") # type: ignore[union-attr]
|
|
or endpoint_url.endswith("/v1/chat/completions") # type: ignore[union-attr]
|
|
):
|
|
raise ValueError(
|
|
"Endpoints of type `serverless` should follow the format "
|
|
"`https://<your-endpoint>.<your_region>.inference.ml.azure.com/v1/chat/completions`"
|
|
" or `https://<your-endpoint>.<your_region>.inference.ml.azure.com/v1/chat/completions`"
|
|
)
|
|
|
|
return field_value
|
|
|
|
@validator("http_client", always=True)
|
|
def validate_client(cls, field_value: Any, values: Dict) -> AzureMLEndpointClient:
|
|
"""Validate that api key and python package exists in environment."""
|
|
endpoint_url = values.get("endpoint_url")
|
|
endpoint_key = values.get("endpoint_api_key")
|
|
deployment_name = values.get("deployment_name")
|
|
timeout = values.get("timeout", DEFAULT_TIMEOUT)
|
|
|
|
http_client = AzureMLEndpointClient(
|
|
endpoint_url, # type: ignore
|
|
endpoint_key.get_secret_value(), # type: ignore
|
|
deployment_name, # type: ignore
|
|
timeout, # type: ignore
|
|
)
|
|
|
|
return http_client
|
|
|
|
|
|
class AzureMLOnlineEndpoint(BaseLLM, AzureMLBaseEndpoint):
|
|
"""Azure ML Online Endpoint models.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
azure_llm = AzureMLOnlineEndpoint(
|
|
endpoint_url="https://<your-endpoint>.<your_region>.inference.ml.azure.com/score",
|
|
endpoint_api_type=AzureMLApiType.dedicated,
|
|
endpoint_api_key="my-api-key",
|
|
timeout=120,
|
|
content_formatter=content_formatter,
|
|
)
|
|
"""
|
|
|
|
@property
|
|
def _identifying_params(self) -> Mapping[str, Any]:
|
|
"""Get the identifying parameters."""
|
|
_model_kwargs = self.model_kwargs or {}
|
|
return {
|
|
**{"deployment_name": self.deployment_name},
|
|
**{"model_kwargs": _model_kwargs},
|
|
}
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
"""Return type of llm."""
|
|
return "azureml_endpoint"
|
|
|
|
def _generate(
|
|
self,
|
|
prompts: List[str],
|
|
stop: Optional[List[str]] = None,
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
**kwargs: Any,
|
|
) -> LLMResult:
|
|
"""Run the LLM on the given prompts.
|
|
|
|
Args:
|
|
prompts: The prompt to pass into the model.
|
|
stop: Optional list of stop words to use when generating.
|
|
Returns:
|
|
The string generated by the model.
|
|
Example:
|
|
.. code-block:: python
|
|
response = azureml_model.invoke("Tell me a joke.")
|
|
"""
|
|
_model_kwargs = self.model_kwargs or {}
|
|
_model_kwargs.update(kwargs)
|
|
if stop:
|
|
_model_kwargs["stop"] = stop
|
|
generations = []
|
|
|
|
for prompt in prompts:
|
|
request_payload = self.content_formatter.format_request_payload(
|
|
prompt, _model_kwargs, self.endpoint_api_type
|
|
)
|
|
response_payload = self.http_client.call(
|
|
body=request_payload, run_manager=run_manager
|
|
)
|
|
generated_text = self.content_formatter.format_response_payload(
|
|
response_payload, self.endpoint_api_type
|
|
)
|
|
generations.append([generated_text])
|
|
|
|
return LLMResult(generations=generations)
|