From 574de1d7fd7eeaa8b47600c0042ff03888d3f251 Mon Sep 17 00:00:00 2001 From: Simon Sanchez Viloria Date: Thu, 4 Jul 2024 14:12:14 +0200 Subject: [PATCH 1/2] (fix) fixed watsonx embedding returning a coroutine --- litellm/llms/watsonx.py | 393 ++++++++++++++++++++++------------------ 1 file changed, 220 insertions(+), 173 deletions(-) diff --git a/litellm/llms/watsonx.py b/litellm/llms/watsonx.py index 34176a23a45..0ac9b4dca88 100644 --- a/litellm/llms/watsonx.py +++ b/litellm/llms/watsonx.py @@ -1,5 +1,7 @@ -from enum import Enum import json, types, time # noqa: E401 +import asyncio +from datetime import datetime +from enum import Enum from contextlib import asynccontextmanager, contextmanager from typing import ( Callable, @@ -12,8 +14,6 @@ from typing import ( Any, Union, List, - ContextManager, - AsyncContextManager, ) import httpx # type: ignore @@ -285,7 +285,10 @@ class IBMWatsonXAI(BaseLLM): ) def _get_api_params( - self, params: dict, print_verbose: Optional[Callable] = None + self, + params: dict, + print_verbose: Optional[Callable] = None, + generate_token: Optional[bool] = True, ) -> dict: """ Find watsonx.ai credentials in the params or environment variables and return the headers for authentication. @@ -365,7 +368,7 @@ class IBMWatsonXAI(BaseLLM): 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: + if token is None and api_key is not None and generate_token: # generate the auth token if print_verbose is not None: print_verbose("Generating IAM token for Watsonx.ai") @@ -393,6 +396,35 @@ class IBMWatsonXAI(BaseLLM): "api_version": api_version, } + def _process_text_gen_response( + self, json_resp: dict, model_response: Union[ModelResponse, None] = None + ) -> ModelResponse: + if "results" not in json_resp: + raise WatsonXAIError( + status_code=500, + message=f"Error: Invalid response from Watsonx.ai API: {json_resp}", + ) + if model_response is None: + model_response = ModelResponse(model=json_resp.get("model_id", None)) + 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"] + if json_resp.get("created_at"): + model_response["created"] = datetime.fromisoformat( + json_resp["created_at"] + ).timestamp() + else: + model_response["created"] = int(time.time()) + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + setattr(model_response, "usage", usage) + return model_response + def completion( self, model: str, @@ -426,27 +458,7 @@ class IBMWatsonXAI(BaseLLM): prompt = convert_messages_to_prompt( model, messages, provider, custom_prompt_dict ) - - def process_text_gen_response(json_resp: dict) -> ModelResponse: - if "results" not in json_resp: - raise WatsonXAIError( - status_code=500, - message=f"Error: Invalid response from Watsonx.ai API: {json_resp}", - ) - 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 - usage = Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) - setattr(model_response, "usage", usage) - return model_response + model_response["model"] = model def process_stream_response( stream_resp: Union[Iterator[str], AsyncIterator], @@ -470,7 +482,7 @@ class IBMWatsonXAI(BaseLLM): ) as resp: json_resp = resp.json() - return process_text_gen_response(json_resp) + return self._process_text_gen_response(json_resp, model_response) async def handle_text_request_async(request_params: dict) -> ModelResponse: async with self.request_manager.async_request( @@ -479,7 +491,7 @@ class IBMWatsonXAI(BaseLLM): timeout=timeout, ) as resp: json_resp = resp.json() - return process_text_gen_response(json_resp) + return self._process_text_gen_response(json_resp, model_response) def handle_stream_request(request_params: dict) -> litellm.CustomStreamWrapper: # stream the response - generated chunks will be handled @@ -493,7 +505,9 @@ class IBMWatsonXAI(BaseLLM): streamwrapper = process_stream_response(resp.iter_lines()) return streamwrapper - async def handle_stream_request_async(request_params: dict) -> litellm.CustomStreamWrapper: + async def handle_stream_request_async( + request_params: dict, + ) -> litellm.CustomStreamWrapper: # stream the response - generated chunks will be handled # by litellm.utils.CustomStreamWrapper.handle_watsonx_stream async with self.request_manager.async_request( @@ -520,7 +534,7 @@ class IBMWatsonXAI(BaseLLM): elif stream: # streaming text generation return handle_stream_request(req_params) - elif (acompletion is True): + elif acompletion is True: # async text generation return handle_text_request_async(req_params) else: @@ -530,6 +544,29 @@ class IBMWatsonXAI(BaseLLM): raise e except Exception as e: raise WatsonXAIError(status_code=500, message=str(e)) + + def _process_embedding_response(self, json_resp: dict, model_response:Union[ModelResponse,None]=None) -> ModelResponse: + if model_response is None: + model_response = ModelResponse(model=json_resp.get("model_id", None)) + 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 + 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 embedding( self, @@ -540,6 +577,7 @@ class IBMWatsonXAI(BaseLLM): model_response=None, optional_params=None, encoding=None, + print_verbose=None, aembedding=None, ): """ @@ -553,6 +591,8 @@ class IBMWatsonXAI(BaseLLM): if k not in optional_params: optional_params[k] = v + model_response['model'] = model + # Load auth variables from environment variables if isinstance(input, str): input = [input] @@ -584,43 +624,33 @@ class IBMWatsonXAI(BaseLLM): } request_manager = RequestManager(logging_obj) - def process_embedding_response(json_resp: dict) -> ModelResponse: - 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 handle_embedding(request_params: dict) -> ModelResponse: with request_manager.request(request_params, input=input) as resp: json_resp = resp.json() - return process_embedding_response(json_resp) + logging_obj.post_call( + original_response=json.dumps(json_resp), + input=input, + api_key=api_params.get("api_key"), + ) + return self._process_embedding_response(json_resp, model_response) async def handle_aembedding(request_params: dict) -> ModelResponse: - async with request_manager.async_request(request_params, input=input) as resp: + async with request_manager.async_request( + request_params, input=input + ) as resp: json_resp = resp.json() - return process_embedding_response(json_resp) + logging_obj.post_call( + original_response=json.dumps(json_resp), + input=input, + api_key=api_params.get("api_key"), + ) + return self._process_embedding_response(json_resp, model_response) try: if aembedding is True: - return handle_embedding(req_params) - else: return handle_aembedding(req_params) + else: + return handle_embedding(req_params) except WatsonXAIError as e: raise e except Exception as e: @@ -664,127 +694,144 @@ class IBMWatsonXAI(BaseLLM): return [res["model_id"] for res in json_resp["resources"]] class RequestManager: + """ + Returns a context manager that manages the response from the request. + if async_ is True, returns an async context manager, otherwise returns a regular context manager. + + Usage: + ```python + request_params = dict(method="POST", url="https://api.example.com", headers={"Authorization" : "Bearer token"}, json={"key": "value"}) + request_manager = RequestManager(logging_obj=logging_obj) + with request_manager.request(request_params) as resp: + ... + # or + async with request_manager.async_request(request_params) as resp: + ... + ``` + """ + + def __init__(self, logging_obj=None): + self.logging_obj = logging_obj + + def pre_call( + self, + request_params: dict, + input: Optional[Any] = None, + ): + if self.logging_obj is None: + return + request_str = ( + f"response = {request_params['method']}(\n" + f"\turl={request_params['url']},\n" + f"\tjson={request_params.get('json')},\n" + f")" + ) + self.logging_obj.pre_call( + input=input, + api_key=request_params["headers"].get("Authorization"), + additional_args={ + "complete_input_dict": request_params.get("json"), + "request_str": request_str, + }, + ) + + def post_call(self, resp, request_params): + if self.logging_obj is None: + return + self.logging_obj.post_call( + input=input, + api_key=request_params["headers"].get("Authorization"), + original_response=json.dumps(resp.json()), + additional_args={ + "status_code": resp.status_code, + "complete_input_dict": request_params.get( + "data", request_params.get("json") + ), + }, + ) + + @contextmanager + def request( + self, + request_params: dict, + stream: bool = False, + input: Optional[Any] = None, + timeout=None, + ) -> Generator[requests.Response, None, None]: """ - Returns a context manager that manages the response from the request. - if async_ is True, returns an async context manager, otherwise returns a regular context manager. - - Usage: - ```python - request_params = dict(method="POST", url="https://api.example.com", headers={"Authorization" : "Bearer token"}, json={"key": "value"}) - request_manager = RequestManager(logging_obj=logging_obj) - async with request_manager.request(request_params) as resp: - ... - # or - with request_manager.async_request(request_params) as resp: - ... - ``` + Returns a context manager that yields the response from the request. """ - - def __init__(self, logging_obj=None): - self.logging_obj = logging_obj - - def pre_call( - self, - request_params: dict, - input: Optional[Any] = None, - ): - if self.logging_obj is None: - return - request_str = ( - f"response = {request_params['method']}(\n" - f"\turl={request_params['url']},\n" - f"\tjson={request_params.get('json')},\n" - f")" - ) - self.logging_obj.pre_call( - input=input, - api_key=request_params["headers"].get("Authorization"), - additional_args={ - "complete_input_dict": request_params.get("json"), - "request_str": request_str, - }, - ) - - def post_call(self, resp, request_params): - if self.logging_obj is None: - return - self.logging_obj.post_call( - input=input, - api_key=request_params["headers"].get("Authorization"), - original_response=json.dumps(resp.json()), - additional_args={ - "status_code": resp.status_code, - "complete_input_dict": request_params.get( - "data", request_params.get("json") - ), - }, - ) - - @contextmanager - def request( - self, - request_params: dict, - stream: bool = False, - input: Optional[Any] = None, - timeout=None, - ) -> Generator[requests.Response, None, None]: - """ - Returns a context manager that yields the response from the request. - """ - self.pre_call(request_params, input) - if timeout: - request_params["timeout"] = timeout - if stream: - request_params["stream"] = stream - try: + self.pre_call(request_params, input) + if timeout: + request_params["timeout"] = timeout + if stream: + request_params["stream"] = stream + try: + retries = 0 + while retries < 3: resp = requests.request(**request_params) - if not resp.ok: - raise WatsonXAIError( - status_code=resp.status_code, - message=f"Error {resp.status_code} ({resp.reason}): {resp.text}", - ) - yield resp - except Exception as e: - raise WatsonXAIError(status_code=500, message=str(e)) - if not stream: - self.post_call(resp, request_params) - - @asynccontextmanager - async def async_request( - self, - request_params: dict, - stream: bool = False, - input: Optional[Any] = None, - timeout=None, - ) -> AsyncGenerator[httpx.Response, None]: - self.pre_call(request_params, input) - if timeout: - request_params["timeout"] = timeout - if stream: - request_params["stream"] = stream - try: - # async with AsyncHTTPHandler(timeout=timeout) as client: - self.async_handler = AsyncHTTPHandler( - timeout=httpx.Timeout( - timeout=request_params.pop("timeout", 600.0), connect=5.0 - ), + if resp.status_code in [429, 503, 504, 520]: + # to handle rate limiting and service unavailable errors + # see: ibm_watsonx_ai.foundation_models.inference.base_model_inference.BaseModelInference._send_inference_payload + time.sleep(2**retries) + retries += 1 + else: + break + if not resp.ok: + raise WatsonXAIError( + status_code=resp.status_code, + message=f"Error {resp.status_code} ({resp.reason}): {resp.text}", ) - # async_handler.client.verify = False - if "json" in request_params: - request_params["data"] = json.dumps(request_params.pop("json", {})) - method = request_params.pop("method") + yield resp + except Exception as e: + raise WatsonXAIError(status_code=500, message=str(e)) + if not stream: + self.post_call(resp, request_params) + + @asynccontextmanager + async def async_request( + self, + request_params: dict, + stream: bool = False, + input: Optional[Any] = None, + timeout=None, + ) -> AsyncGenerator[httpx.Response, None]: + self.pre_call(request_params, input) + if timeout: + request_params["timeout"] = timeout + if stream: + request_params["stream"] = stream + try: + self.async_handler = AsyncHTTPHandler( + timeout=httpx.Timeout( + timeout=request_params.pop("timeout", 600.0), connect=5.0 + ), + ) + # async_handler.client.verify = False + if "json" in request_params: + request_params["data"] = json.dumps(request_params.pop("json", {})) + method = request_params.pop("method") + retries = 0 + while retries < 3: if method.upper() == "POST": resp = await self.async_handler.post(**request_params) else: resp = await self.async_handler.get(**request_params) - if resp.status_code not in [200, 201]: - raise WatsonXAIError( - status_code=resp.status_code, - message=f"Error {resp.status_code} ({resp.reason}): {resp.text}", - ) - yield resp - # await async_handler.close() - except Exception as e: - raise WatsonXAIError(status_code=500, message=str(e)) - if not stream: - self.post_call(resp, request_params) \ No newline at end of file + if resp.status_code in [429, 503, 504, 520]: + # to handle rate limiting and service unavailable errors + # see: ibm_watsonx_ai.foundation_models.inference.base_model_inference.BaseModelInference._send_inference_payload + await asyncio.sleep(2**retries) + retries += 1 + else: + break + if resp.is_error: + raise WatsonXAIError( + status_code=resp.status_code, + message=f"Error {resp.status_code} ({resp.reason}): {resp.text}", + ) + yield resp + # await async_handler.close() + except Exception as e: + raise WatsonXAIError(status_code=500, message=str(e)) + if not stream: + self.post_call(resp, request_params) From e4014d3e0e0e929da4ca69ee94ded4297f309aa9 Mon Sep 17 00:00:00 2001 From: Simon Sanchez Viloria Date: Thu, 4 Jul 2024 14:26:17 +0200 Subject: [PATCH 2/2] fix linting error --- litellm/llms/watsonx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/watsonx.py b/litellm/llms/watsonx.py index 0ac9b4dca88..fbee32e6eb1 100644 --- a/litellm/llms/watsonx.py +++ b/litellm/llms/watsonx.py @@ -397,7 +397,7 @@ class IBMWatsonXAI(BaseLLM): } def _process_text_gen_response( - self, json_resp: dict, model_response: Union[ModelResponse, None] = None + self, json_resp: dict, model_response = None ) -> ModelResponse: if "results" not in json_resp: raise WatsonXAIError( @@ -430,7 +430,7 @@ class IBMWatsonXAI(BaseLLM): model: str, messages: list, custom_prompt_dict: dict, - model_response: ModelResponse, + model_response, print_verbose: Callable, encoding, logging_obj, @@ -545,7 +545,7 @@ class IBMWatsonXAI(BaseLLM): except Exception as e: raise WatsonXAIError(status_code=500, message=str(e)) - def _process_embedding_response(self, json_resp: dict, model_response:Union[ModelResponse,None]=None) -> ModelResponse: + def _process_embedding_response(self, json_resp: dict, model_response=None) -> ModelResponse: if model_response is None: model_response = ModelResponse(model=json_resp.get("model_id", None)) results = json_resp.get("results", [])