From 6a5cf7f83acf1bad34f4e13ad5179b4b57e5afb6 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 5 Jul 2024 11:40:59 -0700 Subject: [PATCH] Revert "(fix) fixed bug with the watsonx embedding endpoint" --- litellm/llms/watsonx.py | 397 ++++++++++++++++++---------------------- 1 file changed, 175 insertions(+), 222 deletions(-) diff --git a/litellm/llms/watsonx.py b/litellm/llms/watsonx.py index fbee32e6eb1..34176a23a45 100644 --- a/litellm/llms/watsonx.py +++ b/litellm/llms/watsonx.py @@ -1,7 +1,5 @@ -import json, types, time # noqa: E401 -import asyncio -from datetime import datetime from enum import Enum +import json, types, time # noqa: E401 from contextlib import asynccontextmanager, contextmanager from typing import ( Callable, @@ -14,6 +12,8 @@ from typing import ( Any, Union, List, + ContextManager, + AsyncContextManager, ) import httpx # type: ignore @@ -285,10 +285,7 @@ class IBMWatsonXAI(BaseLLM): ) def _get_api_params( - self, - params: dict, - print_verbose: Optional[Callable] = None, - generate_token: Optional[bool] = True, + self, params: dict, print_verbose: Optional[Callable] = None ) -> dict: """ Find watsonx.ai credentials in the params or environment variables and return the headers for authentication. @@ -368,7 +365,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 and generate_token: + if token is None and api_key is not None: # generate the auth token if print_verbose is not None: print_verbose("Generating IAM token for Watsonx.ai") @@ -396,41 +393,12 @@ class IBMWatsonXAI(BaseLLM): "api_version": api_version, } - def _process_text_gen_response( - self, json_resp: dict, model_response = 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, messages: list, custom_prompt_dict: dict, - model_response, + model_response: ModelResponse, print_verbose: Callable, encoding, logging_obj, @@ -458,7 +426,27 @@ class IBMWatsonXAI(BaseLLM): prompt = convert_messages_to_prompt( model, messages, provider, custom_prompt_dict ) - model_response["model"] = model + + 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 def process_stream_response( stream_resp: Union[Iterator[str], AsyncIterator], @@ -482,7 +470,7 @@ class IBMWatsonXAI(BaseLLM): ) as resp: json_resp = resp.json() - return self._process_text_gen_response(json_resp, model_response) + return process_text_gen_response(json_resp) async def handle_text_request_async(request_params: dict) -> ModelResponse: async with self.request_manager.async_request( @@ -491,7 +479,7 @@ class IBMWatsonXAI(BaseLLM): timeout=timeout, ) as resp: json_resp = resp.json() - return self._process_text_gen_response(json_resp, model_response) + return process_text_gen_response(json_resp) def handle_stream_request(request_params: dict) -> litellm.CustomStreamWrapper: # stream the response - generated chunks will be handled @@ -505,9 +493,7 @@ 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( @@ -534,7 +520,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: @@ -544,29 +530,6 @@ 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=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, @@ -577,7 +540,6 @@ class IBMWatsonXAI(BaseLLM): model_response=None, optional_params=None, encoding=None, - print_verbose=None, aembedding=None, ): """ @@ -591,8 +553,6 @@ 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] @@ -624,33 +584,43 @@ 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() - 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) + return process_embedding_response(json_resp) 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() - 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) + return process_embedding_response(json_resp) try: if aembedding is True: - return handle_aembedding(req_params) - else: return handle_embedding(req_params) + else: + return handle_aembedding(req_params) except WatsonXAIError as e: raise e except Exception as e: @@ -694,144 +664,127 @@ 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 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: - retries = 0 - while retries < 3: - resp = requests.request(**request_params) - 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}", - ) - yield resp - except Exception as e: - raise WatsonXAIError(status_code=500, message=str(e)) - if not stream: - self.post_call(resp, request_params) + 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. - @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 - ), + 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: + ... + ``` + """ + + 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")" ) - # 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: + 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: + 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 + ), + ) + # async_handler.client.verify = False + if "json" in request_params: + request_params["data"] = json.dumps(request_params.pop("json", {})) + method = request_params.pop("method") 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 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) + 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