use litellm_parent_otel_span as litellm_param

This commit is contained in:
Ishaan Jaff 2024-06-07 08:54:28 -07:00
parent c7545be275
commit 54ac848bfb
2 changed files with 7 additions and 2 deletions

View file

@ -600,6 +600,7 @@ def completion(
client = kwargs.get("client", None) client = kwargs.get("client", None)
### Admin Controls ### ### Admin Controls ###
no_log = kwargs.get("no-log", False) no_log = kwargs.get("no-log", False)
litellm_parent_otel_span = kwargs.get("litellm_parent_otel_span", None)
######## end of unpacking kwargs ########### ######## end of unpacking kwargs ###########
openai_params = [ openai_params = [
"functions", "functions",
@ -689,6 +690,7 @@ def completion(
"allowed_model_region", "allowed_model_region",
"model_config", "model_config",
"fastest_response", "fastest_response",
"litellm_parent_otel_span",
] ]
default_params = openai_params + litellm_params default_params = openai_params + litellm_params
@ -873,6 +875,7 @@ def completion(
input_cost_per_token=input_cost_per_token, input_cost_per_token=input_cost_per_token,
output_cost_per_second=output_cost_per_second, output_cost_per_second=output_cost_per_second,
output_cost_per_token=output_cost_per_token, output_cost_per_token=output_cost_per_token,
litellm_parent_otel_span=litellm_parent_otel_span,
) )
logging.update_environment_variables( logging.update_environment_variables(
model=model, model=model,

View file

@ -4918,6 +4918,7 @@ def get_litellm_params(
input_cost_per_token=None, input_cost_per_token=None,
output_cost_per_token=None, output_cost_per_token=None,
output_cost_per_second=None, output_cost_per_second=None,
litellm_parent_otel_span=None,
): ):
litellm_params = { litellm_params = {
"acompletion": acompletion, "acompletion": acompletion,
@ -4940,6 +4941,7 @@ def get_litellm_params(
"input_cost_per_second": input_cost_per_second, "input_cost_per_second": input_cost_per_second,
"output_cost_per_token": output_cost_per_token, "output_cost_per_token": output_cost_per_token,
"output_cost_per_second": output_cost_per_second, "output_cost_per_second": output_cost_per_second,
"litellm_parent_otel_span": litellm_parent_otel_span,
} }
return litellm_params return litellm_params
@ -7351,10 +7353,10 @@ def get_provider_fields(custom_llm_provider: str) -> List[ProviderField]:
if custom_llm_provider == "databricks": if custom_llm_provider == "databricks":
return litellm.DatabricksConfig().get_required_params() return litellm.DatabricksConfig().get_required_params()
elif custom_llm_provider == "ollama": elif custom_llm_provider == "ollama":
return litellm.OllamaConfig().get_required_params() return litellm.OllamaConfig().get_required_params()
else: else:
return [] return []