feat - watsonx refractoring, removed dependency, and added support for embedding calls

This commit is contained in:
Simon Sanchez Viloria 2024-04-23 11:53:38 +02:00
parent a77537ddd4
commit 74d2ba0a23
4 changed files with 477 additions and 366 deletions

View file

@ -1862,7 +1862,7 @@ def completion(
response = response
elif custom_llm_provider == "watsonx":
custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict
response = watsonx.completion(
response = watsonx.IBMWatsonXAI().completion(
model=model,
messages=messages,
custom_prompt_dict=custom_prompt_dict,
@ -2976,6 +2976,15 @@ def embedding(
client=client,
aembedding=aembedding,
)
elif custom_llm_provider == "watsonx":
response = watsonx.IBMWatsonXAI().embedding(
model=model,
input=input,
encoding=encoding,
logging_obj=logging,
optional_params=optional_params,
model_response=EmbeddingResponse(),
)
else:
args = locals()
raise ValueError(f"No valid embedding model args passed in - {args}")