diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index 62787a0b5fd..4347a731134 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -141,227 +141,233 @@ def completion( litellm_params=None, logger_fn=None, ): - headers = validate_environment(api_key, headers) - task = get_hf_task_for_model(model) - print_verbose(f"{model}, {task}") - completion_url = "" - input_text = None - if "https" in model: - completion_url = model - elif api_base: - completion_url = api_base - elif "HF_API_BASE" in os.environ: - completion_url = os.getenv("HF_API_BASE", "") - elif "HUGGINGFACE_API_BASE" in os.environ: - completion_url = os.getenv("HUGGINGFACE_API_BASE", "") - else: - completion_url = f"https://api-inference.huggingface.co/models/{model}" - - ## Load Config - config=litellm.HuggingfaceConfig.get_config() - for k, v in config.items(): - if k not in optional_params: # completion(top_k=3) > huggingfaceConfig(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ### MAP INPUT PARAMS - if task == "conversational": - inference_params = copy.deepcopy(optional_params) - inference_params.pop("details") - inference_params.pop("return_full_text") - past_user_inputs = [] - generated_responses = [] - text = "" - for message in messages: - if message["role"] == "user": - if text != "": - past_user_inputs.append(text) - text = message["content"] - elif message["role"] == "assistant" or message["role"] == "system": - generated_responses.append(message["content"]) - data = { - "inputs": { - "text": text, - "past_user_inputs": past_user_inputs, - "generated_responses": generated_responses - }, - "parameters": inference_params - } - input_text = "".join(message["content"] for message in messages) - elif task == "text-generation-inference": - # always send "details" and "return_full_text" as params - if model in custom_prompt_dict: - # check if the model has a registered custom prompt - model_prompt_details = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages - ) + try: + headers = validate_environment(api_key, headers) + task = get_hf_task_for_model(model) + print_verbose(f"{model}, {task}") + completion_url = "" + input_text = None + if "https" in model: + completion_url = model + elif api_base: + completion_url = api_base + elif "HF_API_BASE" in os.environ: + completion_url = os.getenv("HF_API_BASE", "") + elif "HUGGINGFACE_API_BASE" in os.environ: + completion_url = os.getenv("HUGGINGFACE_API_BASE", "") else: - prompt = prompt_factory(model=model, messages=messages) - data = { - "inputs": prompt, - "parameters": optional_params, - "stream": True if "stream" in optional_params and optional_params["stream"] == True else False, - } - input_text = prompt - else: - # Non TGI and Conversational llms - # We need this branch, it removes 'details' and 'return_full_text' from params - if model in custom_prompt_dict: - # check if the model has a registered custom prompt - model_prompt_details = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", {}), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - bos_token=model_prompt_details.get("bos_token", ""), - eos_token=model_prompt_details.get("eos_token", ""), - messages=messages, - ) - else: - prompt = prompt_factory(model=model, messages=messages) - inference_params = copy.deepcopy(optional_params) - inference_params.pop("details") - inference_params.pop("return_full_text") - data = { - "inputs": prompt, - "parameters": inference_params, - "stream": True if "stream" in optional_params and optional_params["stream"] == True else False, - } - input_text = prompt - ## LOGGING - logging_obj.pre_call( - input=input_text, - api_key=api_key, - additional_args={"complete_input_dict": data, "task": task, "headers": headers}, - ) - ## COMPLETION CALL - if "stream" in optional_params and optional_params["stream"] == True: - response = requests.post( - completion_url, - headers=headers, - data=json.dumps(data), - stream=optional_params["stream"] - ) - return response.iter_lines() - else: - response = requests.post( - completion_url, - headers=headers, - data=json.dumps(data) - ) + completion_url = f"https://api-inference.huggingface.co/models/{model}" - ## Some servers might return streaming responses even though stream was not set to true. (e.g. Baseten) - is_streamed = False - if response.__dict__['headers'].get("Content-Type", "") == "text/event-stream": - is_streamed = True - - # iterate over the complete streamed response, and return the final answer - if is_streamed: - streamed_response = CustomStreamWrapper(completion_stream=response.iter_lines(), model=model, custom_llm_provider="huggingface", logging_obj=logging_obj) - content = "" - for chunk in streamed_response: - content += chunk["choices"][0]["delta"]["content"] - completion_response: List[Dict[str, Any]] = [{"generated_text": content}] - ## LOGGING - logging_obj.post_call( - input=input_text, - api_key=api_key, - original_response=completion_response, - additional_args={"complete_input_dict": data, "task": task}, - ) - else: - ## LOGGING - logging_obj.post_call( - input=input_text, - api_key=api_key, - original_response=response.text, - additional_args={"complete_input_dict": data, "task": task}, - ) - ## RESPONSE OBJECT - try: - completion_response = response.json() - except: - raise HuggingfaceError( - message=response.text, status_code=response.status_code + ## Load Config + config=litellm.HuggingfaceConfig.get_config() + for k, v in config.items(): + if k not in optional_params: # completion(top_k=3) > huggingfaceConfig(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ### MAP INPUT PARAMS + if task == "conversational": + inference_params = copy.deepcopy(optional_params) + inference_params.pop("details") + inference_params.pop("return_full_text") + past_user_inputs = [] + generated_responses = [] + text = "" + for message in messages: + if message["role"] == "user": + if text != "": + past_user_inputs.append(text) + text = message["content"] + elif message["role"] == "assistant" or message["role"] == "system": + generated_responses.append(message["content"]) + data = { + "inputs": { + "text": text, + "past_user_inputs": past_user_inputs, + "generated_responses": generated_responses + }, + "parameters": inference_params + } + input_text = "".join(message["content"] for message in messages) + elif task == "text-generation-inference": + # always send "details" and "return_full_text" as params + if model in custom_prompt_dict: + # check if the model has a registered custom prompt + model_prompt_details = custom_prompt_dict[model] + prompt = custom_prompt( + role_dict=model_prompt_details.get("roles", None), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), + final_prompt_value=model_prompt_details.get("final_prompt_value", ""), + messages=messages ) - print_verbose(f"response: {completion_response}") - if isinstance(completion_response, dict) and "error" in completion_response: - print_verbose(f"completion error: {completion_response['error']}") - print_verbose(f"response.status_code: {response.status_code}") - raise HuggingfaceError( - message=completion_response["error"], - status_code=response.status_code, - ) - else: - if task == "conversational": - if len(completion_response["generated_text"]) > 0: # type: ignore - model_response["choices"][0]["message"][ - "content" - ] = completion_response["generated_text"] # type: ignore - elif task == "text-generation-inference": - if len(completion_response[0]["generated_text"]) > 0: - model_response["choices"][0]["message"][ - "content" - ] = completion_response[0]["generated_text"] - ## GETTING LOGPROBS + FINISH REASON - if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]: - model_response.choices[0].finish_reason = completion_response[0]["details"]["finish_reason"] - sum_logprob = 0 - for token in completion_response[0]["details"]["tokens"]: - sum_logprob += token["logprob"] - model_response["choices"][0]["message"]._logprob = sum_logprob - if "best_of" in optional_params and optional_params["best_of"] > 1: - if "details" in completion_response[0] and "best_of_sequences" in completion_response[0]["details"]: - choices_list = [] - for idx, item in enumerate(completion_response[0]["details"]["best_of_sequences"]): - sum_logprob = 0 - for token in item["tokens"]: - sum_logprob += token["logprob"] - if len(item["generated_text"]) > 0: - message_obj = Message(content=item["generated_text"], logprobs=sum_logprob) - else: - message_obj = Message(content=None) - choice_obj = Choices(finish_reason=item["finish_reason"], index=idx+1, message=message_obj) - choices_list.append(choice_obj) - model_response["choices"].extend(choices_list) else: - if len(completion_response[0]["generated_text"]) > 0: - model_response["choices"][0]["message"][ - "content" - ] = completion_response[0]["generated_text"] - ## CALCULATING USAGE - prompt_tokens = 0 - try: - prompt_tokens = len( - encoding.encode(input_text) - ) ##[TODO] use the llama2 tokenizer here - except: - # this should remain non blocking we should not block a response returning if calculating usage fails - pass - print_verbose(f'output: {model_response["choices"][0]["message"]}') - output_text = model_response["choices"][0]["message"].get("content", "") - if output_text is not None and len(output_text) > 0: - completion_tokens = 0 + prompt = prompt_factory(model=model, messages=messages) + data = { + "inputs": prompt, + "parameters": optional_params, + "stream": True if "stream" in optional_params and optional_params["stream"] == True else False, + } + input_text = prompt + else: + # Non TGI and Conversational llms + # We need this branch, it removes 'details' and 'return_full_text' from params + if model in custom_prompt_dict: + # check if the model has a registered custom prompt + model_prompt_details = custom_prompt_dict[model] + prompt = custom_prompt( + role_dict=model_prompt_details.get("roles", {}), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), + final_prompt_value=model_prompt_details.get("final_prompt_value", ""), + bos_token=model_prompt_details.get("bos_token", ""), + eos_token=model_prompt_details.get("eos_token", ""), + messages=messages, + ) + else: + prompt = prompt_factory(model=model, messages=messages) + inference_params = copy.deepcopy(optional_params) + inference_params.pop("details") + inference_params.pop("return_full_text") + data = { + "inputs": prompt, + "parameters": inference_params, + "stream": True if "stream" in optional_params and optional_params["stream"] == True else False, + } + input_text = prompt + ## LOGGING + logging_obj.pre_call( + input=input_text, + api_key=api_key, + additional_args={"complete_input_dict": data, "task": task, "headers": headers}, + ) + ## COMPLETION CALL + if "stream" in optional_params and optional_params["stream"] == True: + response = requests.post( + completion_url, + headers=headers, + data=json.dumps(data), + stream=optional_params["stream"] + ) + return response.iter_lines() + else: + response = requests.post( + completion_url, + headers=headers, + data=json.dumps(data) + ) + + ## Some servers might return streaming responses even though stream was not set to true. (e.g. Baseten) + is_streamed = False + if response.__dict__['headers'].get("Content-Type", "") == "text/event-stream": + is_streamed = True + + # iterate over the complete streamed response, and return the final answer + if is_streamed: + streamed_response = CustomStreamWrapper(completion_stream=response.iter_lines(), model=model, custom_llm_provider="huggingface", logging_obj=logging_obj) + content = "" + for chunk in streamed_response: + content += chunk["choices"][0]["delta"]["content"] + completion_response: List[Dict[str, Any]] = [{"generated_text": content}] + ## LOGGING + logging_obj.post_call( + input=input_text, + api_key=api_key, + original_response=completion_response, + additional_args={"complete_input_dict": data, "task": task}, + ) + else: + ## LOGGING + logging_obj.post_call( + input=input_text, + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data, "task": task}, + ) + ## RESPONSE OBJECT + try: + completion_response = response.json() + except: + raise HuggingfaceError( + message=f"Original Response received: {response.text}; Stacktrace: {traceback.format_exc()}", status_code=response.status_code + ) + print_verbose(f"response: {completion_response}") + if isinstance(completion_response, dict) and "error" in completion_response: + print_verbose(f"completion error: {completion_response['error']}") + print_verbose(f"response.status_code: {response.status_code}") + raise HuggingfaceError( + message=completion_response["error"], + status_code=response.status_code, + ) + else: + if task == "conversational": + if len(completion_response["generated_text"]) > 0: # type: ignore + model_response["choices"][0]["message"][ + "content" + ] = completion_response["generated_text"] # type: ignore + elif task == "text-generation-inference": + if len(completion_response[0]["generated_text"]) > 0: + model_response["choices"][0]["message"][ + "content" + ] = completion_response[0]["generated_text"] + ## GETTING LOGPROBS + FINISH REASON + if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]: + model_response.choices[0].finish_reason = completion_response[0]["details"]["finish_reason"] + sum_logprob = 0 + for token in completion_response[0]["details"]["tokens"]: + sum_logprob += token["logprob"] + model_response["choices"][0]["message"]._logprob = sum_logprob + if "best_of" in optional_params and optional_params["best_of"] > 1: + if "details" in completion_response[0] and "best_of_sequences" in completion_response[0]["details"]: + choices_list = [] + for idx, item in enumerate(completion_response[0]["details"]["best_of_sequences"]): + sum_logprob = 0 + for token in item["tokens"]: + sum_logprob += token["logprob"] + if len(item["generated_text"]) > 0: + message_obj = Message(content=item["generated_text"], logprobs=sum_logprob) + else: + message_obj = Message(content=None) + choice_obj = Choices(finish_reason=item["finish_reason"], index=idx+1, message=message_obj) + choices_list.append(choice_obj) + model_response["choices"].extend(choices_list) + else: + if len(completion_response[0]["generated_text"]) > 0: + model_response["choices"][0]["message"][ + "content" + ] = completion_response[0]["generated_text"] + ## CALCULATING USAGE + prompt_tokens = 0 try: - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) + prompt_tokens = len( + encoding.encode(input_text) ) ##[TODO] use the llama2 tokenizer here except: # this should remain non blocking we should not block a response returning if calculating usage fails pass - else: - completion_tokens = 0 + print_verbose(f'output: {model_response["choices"][0]["message"]}') + output_text = model_response["choices"][0]["message"].get("content", "") + if output_text is not None and len(output_text) > 0: + completion_tokens = 0 + try: + completion_tokens = len( + encoding.encode(model_response["choices"][0]["message"].get("content", "")) + ) ##[TODO] use the llama2 tokenizer here + except: + # this should remain non blocking we should not block a response returning if calculating usage fails + pass + else: + completion_tokens = 0 - model_response["created"] = time.time() - model_response["model"] = model - model_response.usage.completion_tokens = completion_tokens - model_response.usage.prompt_tokens = prompt_tokens - model_response.usage.total_tokens = prompt_tokens + completion_tokens - model_response._hidden_params["original_response"] = completion_response - return model_response + model_response["created"] = time.time() + model_response["model"] = model + model_response.usage.completion_tokens = completion_tokens + model_response.usage.prompt_tokens = prompt_tokens + model_response.usage.total_tokens = prompt_tokens + completion_tokens + model_response._hidden_params["original_response"] = completion_response + return model_response + except HuggingfaceError as e: + raise e + except Exception as e: + import traceback + raise HuggingfaceError(status_code=500, message=traceback.format_exc()) def embedding( diff --git a/litellm/main.py b/litellm/main.py index 1d56a16009b..dc2310ae80c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1961,8 +1961,7 @@ def moderation(input: str, api_key: Optional[str]=None): ## Set verbose to true -> ```litellm.set_verbose = True``` def print_verbose(print_statement): if litellm.set_verbose: - import logging - logging.info(f"LiteLLM: {print_statement}") + print(print_statement) # noqa def config_completion(**kwargs): if litellm.config_path != None: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index bce43b25e1f..fa7e895fc01 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -52,6 +52,7 @@ def is_port_in_use(port): @click.command() @click.option('--host', default='0.0.0.0', help='Host for the server to listen on.') @click.option('--port', default=8000, help='Port to bind the server to.') +@click.option('--num_workers', default=1, help='Number of uvicorn workers to spin up') @click.option('--api_base', default=None, help='API base URL.') @click.option('--api_version', default="2023-07-01-preview", help='For azure - pass in the api version.') @click.option('--model', '-m', default=None, help='The model name to pass to litellm expects') @@ -74,17 +75,17 @@ def is_port_in_use(port): @click.option('--test', flag_value=True, help='proxy chat completions url to make a test request to') @click.option('--local', is_flag=True, default=False, help='for local debugging') @click.option('--cost', is_flag=True, default=False, help='for viewing cost logs') -def run_server(host, port, api_base, api_version, model, alias, add_key, headers, save, debug, temperature, max_tokens, request_timeout, drop_params, create_proxy, add_function_to_prompt, config, file, max_budget, telemetry, logs, test, local, cost): +def run_server(host, port, api_base, api_version, model, alias, add_key, headers, save, debug, temperature, max_tokens, request_timeout, drop_params, create_proxy, add_function_to_prompt, config, file, max_budget, telemetry, logs, test, local, cost, num_workers): global feature_telemetry args = locals() if local: - from proxy_server import app, initialize, print_cost_logs, usage_telemetry, add_keys_to_config + from proxy_server import app, save_worker_config, print_cost_logs, usage_telemetry, add_keys_to_config debug = True else: try: - from .proxy_server import app, initialize, print_cost_logs, usage_telemetry, add_keys_to_config + from .proxy_server import app, save_worker_config, print_cost_logs, usage_telemetry, add_keys_to_config except ImportError as e: - from proxy_server import app, initialize, print_cost_logs, usage_telemetry, add_keys_to_config + from proxy_server import app, save_worker_config, print_cost_logs, usage_telemetry, add_keys_to_config feature_telemetry = usage_telemetry if create_proxy == True: repo_url = 'https://github.com/BerriAI/litellm' @@ -163,7 +164,7 @@ def run_server(host, port, api_base, api_version, model, alias, add_key, headers else: if headers: headers = json.loads(headers) - initialize(model=model, alias=alias, api_base=api_base, api_version=api_version, debug=debug, temperature=temperature, max_tokens=max_tokens, request_timeout=request_timeout, max_budget=max_budget, telemetry=telemetry, drop_params=drop_params, add_function_to_prompt=add_function_to_prompt, headers=headers, save=save, config=config) + save_worker_config(model=model, alias=alias, api_base=api_base, api_version=api_version, debug=debug, temperature=temperature, max_tokens=max_tokens, request_timeout=request_timeout, max_budget=max_budget, telemetry=telemetry, drop_params=drop_params, add_function_to_prompt=add_function_to_prompt, headers=headers, save=save, config=config) try: import uvicorn except: @@ -174,7 +175,7 @@ def run_server(host, port, api_base, api_version, model, alias, add_key, headers if port == 8000 and is_port_in_use(port): port = random.randint(1024, 49152) - uvicorn.run(app, host=host, port=port) + uvicorn.run("proxy_server:app", host=host, port=port, workers=num_workers) if __name__ == "__main__": diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 88c5e475645..462eed034c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -129,11 +129,12 @@ llm_router: Optional[litellm.Router] = None llm_model_list: Optional[list] = None server_settings: dict = {} log_file = "api_log.json" - +worker_config = None #### HELPER FUNCTIONS #### def print_verbose(print_statement): global user_debug + print(f"user debug value: {user_debug}") if user_debug: print(print_statement) @@ -337,6 +338,9 @@ def load_config(): except: pass +def save_worker_config(**data): + import json + os.environ["WORKER_CONFIG"] = json.dumps(data) def initialize( model, @@ -532,6 +536,7 @@ def litellm_completion(*args, **kwargs): for key, value in m["litellm_params"].items(): kwargs[key] = value break + print(f"litellm set verbose pre-call: {litellm.set_verbose}") if call_type == "chat_completion": response = litellm.completion(*args, **kwargs) elif call_type == "text_completion": @@ -540,6 +545,14 @@ def litellm_completion(*args, **kwargs): return StreamingResponse(data_generator(response), media_type='text/event-stream') return response + +@app.on_event("startup") +def startup_event(): + import json + worker_config = json.loads(os.getenv("WORKER_CONFIG")) + initialize(**worker_config) + print(f"\033[32mWorker Initialized\033[0m\n") + #### API ENDPOINTS #### @router.get("/v1/models") @router.get("/models") # if project requires model list diff --git a/litellm/utils.py b/litellm/utils.py index 5507d9a3a3e..673622d5cb1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -285,8 +285,7 @@ class TextCompletionResponse(OpenAIObject): ############################################################ def print_verbose(print_statement): if litellm.set_verbose: - import logging - logging.info(f"LiteLLM: {print_statement}") + print(print_statement) # noqa ####### LOGGING ################### from enum import Enum