phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix: fix linting issues

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-27 12:21:14 +05:30
parent ae4c7ce3d4
commit c88a8d71f0
2 changed files with 67 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

105
litellm/llms/oobabooga.py
View file

@ -7,6 +7,7 @@ from typing import Callable, Optional
from litellm.utils import ModelResponse, Usage from litellm.utils import ModelResponse, Usage
from .prompt_templates.factory import prompt_factory, custom_prompt from .prompt_templates.factory import prompt_factory, custom_prompt
class OobaboogaError(Exception): class OobaboogaError(Exception):
def __init__(self, status_code, message): def __init__(self, status_code, message):
self.status_code = status_code self.status_code = status_code
@ -15,6 +16,7 @@ class OobaboogaError(Exception):
self.message self.message
) # Call the base class constructor with the parameters it needs ) # Call the base class constructor with the parameters it needs
def validate_environment(api_key): def validate_environment(api_key):
headers = { headers = {
"accept": "application/json", "accept": "application/json",
@ -24,6 +26,7 @@ def validate_environment(api_key):
headers["Authorization"] = f"Token {api_key}" headers["Authorization"] = f"Token {api_key}"
return headers return headers
def completion( def completion(
model: str, model: str,
messages: list, messages: list,
@ -44,11 +47,13 @@ def completion(
completion_url = model completion_url = model
elif api_base: elif api_base:
completion_url = api_base completion_url = api_base
else: else:
raise OobaboogaError(status_code=404, message="API Base not set. Set one via completion(..,api_base='your-api-url')") raise OobaboogaError(
status_code=404,
message="API Base not set. Set one via completion(..,api_base='your-api-url')",
)
model = model model = model
completion_url = completion_url + "/v1/chat/completions" completion_url = completion_url + "/v1/chat/completions"
data = { data = {
"messages": messages, "messages": messages,
@ -57,40 +62,51 @@ def completion(
## LOGGING ## LOGGING
logging_obj.pre_call( logging_obj.pre_call(
input=messages, input=messages,
api_key=api_key, api_key=api_key,
additional_args={"complete_input_dict": data}, additional_args={"complete_input_dict": data},
) )
## COMPLETION CALL ## COMPLETION CALL
response = requests.post( response = requests.post(
completion_url, headers=headers, data=json.dumps(data), stream=optional_params["stream"] if "stream" in optional_params else False completion_url,
headers=headers,
data=json.dumps(data),
stream=optional_params["stream"] if "stream" in optional_params else False,
) )
if "stream" in optional_params and optional_params["stream"] == True: if "stream" in optional_params and optional_params["stream"] == True:
return response.iter_lines() return response.iter_lines()
else: else:
## LOGGING ## LOGGING
logging_obj.post_call( logging_obj.post_call(
input=messages, input=messages,
api_key=api_key, api_key=api_key,
original_response=response.text, original_response=response.text,
additional_args={"complete_input_dict": data}, additional_args={"complete_input_dict": data},
) )
print_verbose(f"raw model_response: {response.text}") print_verbose(f"raw model_response: {response.text}")
## RESPONSE OBJECT ## RESPONSE OBJECT
try: try:
completion_response = response.json() completion_response = response.json()
except: except:
raise OobaboogaError(message=response.text, status_code=response.status_code) raise OobaboogaError(
message=response.text, status_code=response.status_code
)
if "error" in completion_response: if "error" in completion_response:
raise OobaboogaError(message=completion_response["error"],status_code=response.status_code,) raise OobaboogaError(
message=completion_response["error"],
status_code=response.status_code,
)
else: else:
try: try:
model_response["choices"][0]["message"]["content"] = completion_response["choices"][0]["message"]["content"] model_response["choices"][0]["message"][
"content"
] = completion_response["choices"][0]["message"]["content"]
except: except:
raise OobaboogaError(message=json.dumps(completion_response), status_code=response.status_code) raise OobaboogaError(
message=json.dumps(completion_response),
status_code=response.status_code,
)
model_response["created"] = int(time.time()) model_response["created"] = int(time.time())
model_response["model"] = model model_response["model"] = model
@ -103,25 +119,26 @@ def completion(
return model_response return model_response
def embedding( def embedding(
model: str, model: str,
input: list, input: list,
api_key: Optional[str] = None, api_key: Optional[str] = None,
api_base: str = None, api_base: Optional[str] = None,
logging_obj=None, logging_obj=None,
model_response=None, model_response=None,
optional_params=None, optional_params=None,
encoding=None, encoding=None,
): ):
# Create completion URL # Create completion URL
if "https" in model: if "https" in model:
embeddings_url = model embeddings_url = model
elif api_base: elif api_base:
embeddings_url = f"{api_base}/v1/embeddings" embeddings_url = f"{api_base}/v1/embeddings"
else: else:
raise OobaboogaError(status_code=404, message="API Base not set. Set one via completion(..,api_base='your-api-url')") raise OobaboogaError(
status_code=404,
message="API Base not set. Set one via completion(..,api_base='your-api-url')",
)
# Prepare request data # Prepare request data
data = {"input": input} data = {"input": input}
@ -130,7 +147,9 @@ def embedding(
# Logging before API call # Logging before API call
if logging_obj: if logging_obj:
logging_obj.pre_call(input=input, api_key=api_key, additional_args={"complete_input_dict": data}) logging_obj.pre_call(
input=input, api_key=api_key, additional_args={"complete_input_dict": data}
)
# Send POST request # Send POST request
headers = validate_environment(api_key) headers = validate_environment(api_key)
@ -141,16 +160,24 @@ def embedding(
# Check for errors in response # Check for errors in response
if "error" in completion_response: if "error" in completion_response:
raise OobaboogaError(message=completion_response["error"], status_code=completion_response.get('status_code', 500)) raise OobaboogaError(
message=completion_response["error"],
status_code=completion_response.get("status_code", 500),
)
# Process response data # Process response data
model_response["data"]=[{"embedding": completion_response["data"][0]["embedding"], "index": 0, "object": "embedding"}] model_response["data"] = [
{
"embedding": completion_response["data"][0]["embedding"],
"index": 0,
"object": "embedding",
}
]
num_tokens = len(completion_response["data"][0]["embedding"]) num_tokens = len(completion_response["data"][0]["embedding"])
#Adding metadata to response # Adding metadata to response
model_response.usage = Usage(prompt_tokens=num_tokens,total_tokens=num_tokens) model_response.usage = Usage(prompt_tokens=num_tokens, total_tokens=num_tokens)
model_response["object"]="list" model_response["object"] = "list"
model_response["model"]=model model_response["model"] = model
return model_response return model_response

2
litellm/main.py
View file

@ -2322,7 +2322,7 @@ def embedding(
api_base=api_base, api_base=api_base,
logging_obj=logging, logging_obj=logging,
optional_params=optional_params, optional_params=optional_params,
model_response= EmbeddingResponse() model_response=EmbeddingResponse(),
) )
elif custom_llm_provider == "ollama": elif custom_llm_provider == "ollama":
if aembedding == True: if aembedding == True: