forked from phoenix/litellm-mirror
566 lines
21 KiB
Python
566 lines
21 KiB
Python
import json, types, time # noqa: E401
|
|
from contextlib import contextmanager
|
|
from typing import Callable, Dict, Optional, Any, Union, List
|
|
|
|
import httpx
|
|
import requests
|
|
import litellm
|
|
from litellm.utils import ModelResponse, get_secret, Usage
|
|
|
|
from .base import BaseLLM
|
|
from .prompt_templates import factory as ptf
|
|
|
|
|
|
class WatsonXAIError(Exception):
|
|
def __init__(self, status_code, message, url: str = None):
|
|
self.status_code = status_code
|
|
self.message = message
|
|
url = url or "https://https://us-south.ml.cloud.ibm.com"
|
|
self.request = httpx.Request(method="POST", url=url)
|
|
self.response = httpx.Response(status_code=status_code, request=self.request)
|
|
super().__init__(
|
|
self.message
|
|
) # Call the base class constructor with the parameters it needs
|
|
|
|
|
|
class IBMWatsonXAIConfig:
|
|
"""
|
|
Reference: https://cloud.ibm.com/apidocs/watsonx-ai#text-generation
|
|
(See ibm_watsonx_ai.metanames.GenTextParamsMetaNames for a list of all available params)
|
|
|
|
Supported params for all available watsonx.ai foundational models.
|
|
|
|
- `decoding_method` (str): One of "greedy" or "sample"
|
|
|
|
- `temperature` (float): Sets the model temperature for sampling - not available when decoding_method='greedy'.
|
|
|
|
- `max_new_tokens` (integer): Maximum length of the generated tokens.
|
|
|
|
- `min_new_tokens` (integer): Maximum length of input tokens. Any more than this will be truncated.
|
|
|
|
- `length_penalty` (dict): A dictionary with keys "decay_factor" and "start_index".
|
|
|
|
- `stop_sequences` (string[]): list of strings to use as stop sequences.
|
|
|
|
- `top_k` (integer): top k for sampling - not available when decoding_method='greedy'.
|
|
|
|
- `top_p` (integer): top p for sampling - not available when decoding_method='greedy'.
|
|
|
|
- `repetition_penalty` (float): token repetition penalty during text generation.
|
|
|
|
- `truncate_input_tokens` (integer): Truncate input tokens to this length.
|
|
|
|
- `include_stop_sequences` (bool): If True, the stop sequence will be included at the end of the generated text in the case of a match.
|
|
|
|
- `return_options` (dict): A dictionary of options to return. Options include "input_text", "generated_tokens", "input_tokens", "token_ranks". Values are boolean.
|
|
|
|
- `random_seed` (integer): Random seed for text generation.
|
|
|
|
- `moderations` (dict): Dictionary of properties that control the moderations, for usages such as Hate and profanity (HAP) and PII filtering.
|
|
|
|
- `stream` (bool): If True, the model will return a stream of responses.
|
|
"""
|
|
|
|
decoding_method: Optional[str] = "sample"
|
|
temperature: Optional[float] = None
|
|
max_new_tokens: Optional[int] = None # litellm.max_tokens
|
|
min_new_tokens: Optional[int] = None
|
|
length_penalty: Optional[dict] = None # e.g {"decay_factor": 2.5, "start_index": 5}
|
|
stop_sequences: Optional[List[str]] = None # e.g ["}", ")", "."]
|
|
top_k: Optional[int] = None
|
|
top_p: Optional[float] = None
|
|
repetition_penalty: Optional[float] = None
|
|
truncate_input_tokens: Optional[int] = None
|
|
include_stop_sequences: Optional[bool] = False
|
|
return_options: Optional[dict] = None
|
|
return_options: Optional[Dict[str, bool]] = None
|
|
random_seed: Optional[int] = None # e.g 42
|
|
moderations: Optional[dict] = None
|
|
stream: Optional[bool] = False
|
|
|
|
def __init__(
|
|
self,
|
|
decoding_method: Optional[str] = None,
|
|
temperature: Optional[float] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
min_new_tokens: Optional[int] = None,
|
|
length_penalty: Optional[dict] = None,
|
|
stop_sequences: Optional[List[str]] = None,
|
|
top_k: Optional[int] = None,
|
|
top_p: Optional[float] = None,
|
|
repetition_penalty: Optional[float] = None,
|
|
truncate_input_tokens: Optional[int] = None,
|
|
include_stop_sequences: Optional[bool] = None,
|
|
return_options: Optional[dict] = None,
|
|
random_seed: Optional[int] = None,
|
|
moderations: Optional[dict] = None,
|
|
stream: Optional[bool] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
locals_ = locals()
|
|
for key, value in locals_.items():
|
|
if key != "self" and value is not None:
|
|
setattr(self.__class__, key, value)
|
|
|
|
@classmethod
|
|
def get_config(cls):
|
|
return {
|
|
k: v
|
|
for k, v in cls.__dict__.items()
|
|
if not k.startswith("__")
|
|
and not isinstance(
|
|
v,
|
|
(
|
|
types.FunctionType,
|
|
types.BuiltinFunctionType,
|
|
classmethod,
|
|
staticmethod,
|
|
),
|
|
)
|
|
and v is not None
|
|
}
|
|
|
|
def get_supported_openai_params(self):
|
|
return [
|
|
"temperature", # equivalent to temperature
|
|
"max_tokens", # equivalent to max_new_tokens
|
|
"top_p", # equivalent to top_p
|
|
"frequency_penalty", # equivalent to repetition_penalty
|
|
"stop", # equivalent to stop_sequences
|
|
"seed", # equivalent to random_seed
|
|
"stream", # equivalent to stream
|
|
]
|
|
|
|
|
|
def convert_messages_to_prompt(model, messages, provider, custom_prompt_dict):
|
|
# handle anthropic prompts and amazon titan prompts
|
|
if model in custom_prompt_dict:
|
|
# check if the model has a registered custom prompt
|
|
model_prompt_dict = custom_prompt_dict[model]
|
|
prompt = ptf.custom_prompt(
|
|
messages=messages,
|
|
role_dict=model_prompt_dict.get(
|
|
"role_dict", model_prompt_dict.get("roles")
|
|
),
|
|
initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""),
|
|
final_prompt_value=model_prompt_dict.get("final_prompt_value", ""),
|
|
bos_token=model_prompt_dict.get("bos_token", ""),
|
|
eos_token=model_prompt_dict.get("eos_token", ""),
|
|
)
|
|
return prompt
|
|
elif provider == "ibm":
|
|
prompt = ptf.prompt_factory(
|
|
model=model, messages=messages, custom_llm_provider="watsonx"
|
|
)
|
|
elif provider == "ibm-mistralai":
|
|
prompt = ptf.mistral_instruct_pt(messages=messages)
|
|
else:
|
|
prompt = ptf.prompt_factory(
|
|
model=model, messages=messages, custom_llm_provider="watsonx"
|
|
)
|
|
return prompt
|
|
|
|
|
|
class IBMWatsonXAI(BaseLLM):
|
|
"""
|
|
Class to interface with IBM Watsonx.ai API for text generation and embeddings.
|
|
|
|
Reference: https://cloud.ibm.com/apidocs/watsonx-ai
|
|
"""
|
|
|
|
api_version = "2024-03-13"
|
|
_text_gen_endpoint = "/ml/v1/text/generation"
|
|
_text_gen_stream_endpoint = "/ml/v1/text/generation_stream"
|
|
_deployment_text_gen_endpoint = "/ml/v1/deployments/{deployment_id}/text/generation"
|
|
_deployment_text_gen_stream_endpoint = (
|
|
"/ml/v1/deployments/{deployment_id}/text/generation_stream"
|
|
)
|
|
_embeddings_endpoint = "/ml/v1/text/embeddings"
|
|
_prompts_endpoint = "/ml/v1/prompts"
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def _prepare_text_generation_req(
|
|
self,
|
|
model_id: str,
|
|
prompt: str,
|
|
stream: bool,
|
|
optional_params: dict,
|
|
print_verbose: Callable = None,
|
|
) -> httpx.Request:
|
|
"""
|
|
Get the request parameters for text generation.
|
|
"""
|
|
api_params = self._get_api_params(optional_params, print_verbose=print_verbose)
|
|
# build auth headers
|
|
api_token = api_params.get("token")
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_token}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
}
|
|
extra_body_params = optional_params.pop("extra_body", {})
|
|
optional_params.update(extra_body_params)
|
|
# init the payload to the text generation call
|
|
payload = {
|
|
"input": prompt,
|
|
"moderations": optional_params.pop("moderations", {}),
|
|
"parameters": optional_params,
|
|
}
|
|
request_params = dict(version=api_params["api_version"])
|
|
# text generation endpoint deployment or model / stream or not
|
|
if model_id.startswith("deployment/"):
|
|
# deployment models are passed in as 'deployment/<deployment_id>'
|
|
if api_params.get("space_id") is None:
|
|
raise WatsonXAIError(
|
|
status_code=401,
|
|
url=api_params["url"],
|
|
message="Error: space_id is required for models called using the 'deployment/' endpoint. Pass in the space_id as a parameter or set it in the WX_SPACE_ID environment variable.",
|
|
)
|
|
deployment_id = "/".join(model_id.split("/")[1:])
|
|
endpoint = (
|
|
self._deployment_text_gen_stream_endpoint
|
|
if stream
|
|
else self._deployment_text_gen_endpoint
|
|
)
|
|
endpoint = endpoint.format(deployment_id=deployment_id)
|
|
else:
|
|
payload["model_id"] = model_id
|
|
payload["project_id"] = api_params["project_id"]
|
|
endpoint = (
|
|
self._text_gen_stream_endpoint if stream else self._text_gen_endpoint
|
|
)
|
|
url = api_params["url"].rstrip("/") + endpoint
|
|
return httpx.Request(
|
|
"POST", url, headers=headers, json=payload, params=request_params
|
|
)
|
|
|
|
def _get_api_params(self, params: dict, print_verbose: Callable = None) -> dict:
|
|
"""
|
|
Find watsonx.ai credentials in the params or environment variables and return the headers for authentication.
|
|
"""
|
|
# Load auth variables from params
|
|
url = params.pop("url", None)
|
|
api_key = params.pop("apikey", None)
|
|
token = params.pop("token", None)
|
|
project_id = params.pop("project_id", None) # watsonx.ai project_id
|
|
space_id = params.pop("space_id", None) # watsonx.ai deployment space_id
|
|
region_name = params.pop("region_name", params.pop("region", None))
|
|
wx_credentials = params.pop("wx_credentials", None)
|
|
api_version = params.pop("api_version", IBMWatsonXAI.api_version)
|
|
# Load auth variables from environment variables
|
|
if url is None:
|
|
url = (
|
|
get_secret("WATSONX_URL")
|
|
or get_secret("WX_URL")
|
|
or get_secret("WML_URL")
|
|
)
|
|
if api_key is None:
|
|
api_key = (
|
|
get_secret("WATSONX_APIKEY")
|
|
or get_secret("WATSONX_API_KEY")
|
|
or get_secret("WX_API_KEY")
|
|
)
|
|
if token is None:
|
|
token = get_secret("WATSONX_TOKEN") or get_secret("WX_TOKEN")
|
|
if project_id is None:
|
|
project_id = (
|
|
get_secret("WATSONX_PROJECT_ID")
|
|
or get_secret("WX_PROJECT_ID")
|
|
or get_secret("PROJECT_ID")
|
|
)
|
|
if region_name is None:
|
|
region_name = (
|
|
get_secret("WATSONX_REGION")
|
|
or get_secret("WX_REGION")
|
|
or get_secret("REGION")
|
|
)
|
|
if space_id is None:
|
|
space_id = (
|
|
get_secret("WATSONX_DEPLOYMENT_SPACE_ID")
|
|
or get_secret("WATSONX_SPACE_ID")
|
|
or get_secret("WX_SPACE_ID")
|
|
or get_secret("SPACE_ID")
|
|
)
|
|
|
|
# credentials parsing
|
|
if wx_credentials is not None:
|
|
url = wx_credentials.get("url", url)
|
|
api_key = wx_credentials.get(
|
|
"apikey", wx_credentials.get("api_key", api_key)
|
|
)
|
|
token = wx_credentials.get("token", token)
|
|
|
|
# verify that all required credentials are present
|
|
if url is None:
|
|
raise WatsonXAIError(
|
|
status_code=401,
|
|
message="Error: Watsonx URL not set. Set WX_URL in environment variables or pass in as a parameter.",
|
|
)
|
|
if token is None and api_key is not None:
|
|
# generate the auth token
|
|
if print_verbose:
|
|
print_verbose("Generating IAM token for Watsonx.ai")
|
|
token = self.generate_iam_token(api_key)
|
|
elif token is None and api_key is None:
|
|
raise WatsonXAIError(
|
|
status_code=401,
|
|
url=url,
|
|
message="Error: API key or token not found. Set WX_API_KEY or WX_TOKEN in environment variables or pass in as a parameter.",
|
|
)
|
|
if project_id is None:
|
|
raise WatsonXAIError(
|
|
status_code=401,
|
|
url=url,
|
|
message="Error: Watsonx project_id not set. Set WX_PROJECT_ID in environment variables or pass in as a parameter.",
|
|
)
|
|
|
|
return {
|
|
"url": url,
|
|
"api_key": api_key,
|
|
"token": token,
|
|
"project_id": project_id,
|
|
"space_id": space_id,
|
|
"region_name": region_name,
|
|
"api_version": api_version,
|
|
}
|
|
|
|
def completion(
|
|
self,
|
|
model: str,
|
|
messages: list,
|
|
custom_prompt_dict: dict,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
logging_obj,
|
|
optional_params: Optional[dict] = None,
|
|
litellm_params: Optional[dict] = None,
|
|
logger_fn=None,
|
|
timeout: float = None,
|
|
):
|
|
"""
|
|
Send a text generation request to the IBM Watsonx.ai API.
|
|
Reference: https://cloud.ibm.com/apidocs/watsonx-ai#text-generation
|
|
"""
|
|
stream = optional_params.pop("stream", False)
|
|
|
|
# Load default configs
|
|
config = IBMWatsonXAIConfig.get_config()
|
|
for k, v in config.items():
|
|
if k not in optional_params:
|
|
optional_params[k] = v
|
|
|
|
# Make prompt to send to model
|
|
provider = model.split("/")[0]
|
|
# model_name = "/".join(model.split("/")[1:])
|
|
prompt = convert_messages_to_prompt(
|
|
model, messages, provider, custom_prompt_dict
|
|
)
|
|
|
|
def process_text_request(request: httpx.Request) -> ModelResponse:
|
|
with self._manage_response(
|
|
request, logging_obj=logging_obj, input=prompt, timeout=timeout
|
|
) as resp:
|
|
json_resp = resp.json()
|
|
|
|
generated_text = json_resp["results"][0]["generated_text"]
|
|
prompt_tokens = json_resp["results"][0]["input_token_count"]
|
|
completion_tokens = json_resp["results"][0]["generated_token_count"]
|
|
model_response["choices"][0]["message"]["content"] = generated_text
|
|
model_response["finish_reason"] = json_resp["results"][0]["stop_reason"]
|
|
model_response["created"] = int(time.time())
|
|
model_response["model"] = model
|
|
model_response.usage = Usage(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
)
|
|
return model_response
|
|
|
|
def process_stream_request(
|
|
request: httpx.Request,
|
|
) -> litellm.CustomStreamWrapper:
|
|
# stream the response - generated chunks will be handled
|
|
# by litellm.utils.CustomStreamWrapper.handle_watsonx_stream
|
|
with self._manage_response(
|
|
request,
|
|
logging_obj=logging_obj,
|
|
stream=True,
|
|
input=prompt,
|
|
timeout=timeout,
|
|
) as resp:
|
|
response = litellm.CustomStreamWrapper(
|
|
resp.iter_lines(),
|
|
model=model,
|
|
custom_llm_provider="watsonx",
|
|
logging_obj=logging_obj,
|
|
)
|
|
return response
|
|
|
|
try:
|
|
## Get the response from the model
|
|
request = self._prepare_text_generation_req(
|
|
model_id=model,
|
|
prompt=prompt,
|
|
stream=stream,
|
|
optional_params=optional_params,
|
|
print_verbose=print_verbose,
|
|
)
|
|
if stream:
|
|
return process_stream_request(request)
|
|
else:
|
|
return process_text_request(request)
|
|
except WatsonXAIError as e:
|
|
raise e
|
|
except Exception as e:
|
|
raise WatsonXAIError(status_code=500, message=str(e))
|
|
|
|
def embedding(
|
|
self,
|
|
model: str,
|
|
input: Union[list, str],
|
|
api_key: Optional[str] = None,
|
|
logging_obj=None,
|
|
model_response=None,
|
|
optional_params=None,
|
|
encoding=None,
|
|
):
|
|
"""
|
|
Send a text embedding request to the IBM Watsonx.ai API.
|
|
"""
|
|
if optional_params is None:
|
|
optional_params = {}
|
|
# Load default configs
|
|
config = IBMWatsonXAIConfig.get_config()
|
|
for k, v in config.items():
|
|
if k not in optional_params:
|
|
optional_params[k] = v
|
|
|
|
# Load auth variables from environment variables
|
|
if isinstance(input, str):
|
|
input = [input]
|
|
if api_key is not None:
|
|
optional_params["api_key"] = api_key
|
|
api_params = self._get_api_params(optional_params)
|
|
# build auth headers
|
|
api_token = api_params.get("token")
|
|
headers = {
|
|
"Authorization": f"Bearer {api_token}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
}
|
|
# init the payload to the text generation call
|
|
payload = {
|
|
"inputs": input,
|
|
"model_id": model,
|
|
"project_id": api_params["project_id"],
|
|
"parameters": optional_params,
|
|
}
|
|
request_params = dict(version=api_params["api_version"])
|
|
url = api_params["url"].rstrip("/") + self._embeddings_endpoint
|
|
request = httpx.Request(
|
|
"POST", url, headers=headers, json=payload, params=request_params
|
|
)
|
|
with self._manage_response(
|
|
request, logging_obj=logging_obj, input=input
|
|
) as resp:
|
|
json_resp = resp.json()
|
|
|
|
results = json_resp.get("results", [])
|
|
embedding_response = []
|
|
for idx, result in enumerate(results):
|
|
embedding_response.append(
|
|
{"object": "embedding", "index": idx, "embedding": result["embedding"]}
|
|
)
|
|
model_response["object"] = "list"
|
|
model_response["data"] = embedding_response
|
|
model_response["model"] = model
|
|
input_tokens = json_resp.get("input_token_count", 0)
|
|
model_response.usage = Usage(
|
|
prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
|
|
)
|
|
return model_response
|
|
|
|
def generate_iam_token(self, api_key=None, **params):
|
|
headers = {}
|
|
headers["Content-Type"] = "application/x-www-form-urlencoded"
|
|
if api_key is None:
|
|
api_key = get_secret("WX_API_KEY") or get_secret("WATSONX_API_KEY")
|
|
if api_key is None:
|
|
raise ValueError("API key is required")
|
|
headers["Accept"] = "application/json"
|
|
data = {
|
|
"grant_type": "urn:ibm:params:oauth:grant-type:apikey",
|
|
"apikey": api_key,
|
|
}
|
|
response = httpx.post(
|
|
"https://iam.cloud.ibm.com/identity/token", data=data, headers=headers
|
|
)
|
|
response.raise_for_status()
|
|
json_data = response.json()
|
|
iam_access_token = json_data["access_token"]
|
|
self.token = iam_access_token
|
|
return iam_access_token
|
|
|
|
@contextmanager
|
|
def _manage_response(
|
|
self,
|
|
request: httpx.Request,
|
|
logging_obj: Any,
|
|
stream: bool = False,
|
|
input: Optional[Any] = None,
|
|
timeout: float = None,
|
|
):
|
|
request_str = (
|
|
f"response = {request.method}(\n"
|
|
f"\turl={request.url},\n"
|
|
f"\tjson={request.content.decode()},\n"
|
|
f")"
|
|
)
|
|
json_input = json.loads(request.content.decode())
|
|
headers = dict(request.headers)
|
|
logging_obj.pre_call(
|
|
input=input,
|
|
api_key=request.headers.get("Authorization"),
|
|
additional_args={
|
|
"complete_input_dict": json_input,
|
|
"request_str": request_str,
|
|
},
|
|
)
|
|
try:
|
|
if stream:
|
|
resp = requests.request(
|
|
method=request.method,
|
|
url=str(request.url),
|
|
headers=headers,
|
|
json=json_input,
|
|
stream=True,
|
|
timeout=timeout,
|
|
)
|
|
# resp.raise_for_status()
|
|
yield resp
|
|
else:
|
|
resp = requests.request(
|
|
method=request.method,
|
|
url=str(request.url),
|
|
headers=headers,
|
|
json=json_input,
|
|
timeout=timeout,
|
|
)
|
|
resp.raise_for_status()
|
|
yield resp
|
|
except Exception as e:
|
|
raise WatsonXAIError(status_code=500, message=str(e))
|
|
if not stream:
|
|
logging_obj.post_call(
|
|
input=input,
|
|
api_key=request.headers.get("Authorization"),
|
|
original_response=json.dumps(resp.json()),
|
|
additional_args={
|
|
"status_code": resp.status_code,
|
|
"complete_input_dict": request,
|
|
},
|
|
)
|