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

improve replicate usage

Browse source
This commit is contained in:
ishaan-jaff 2023-09-06 12:29:32 -07:00
parent 1bf4bfa85f
commit dc0c084813
2 changed files with 30 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

16
docs/my-website/docs/providers/replicate.md
View file

@ -9,6 +9,22 @@ os.environ["REPLICATE_API_KEY"] = ""
``` ```
### Example Call
```python
from litellm import completion
## set ENV variables
os.environ["REPLICATE_API_KEY"] = "replicate key"
messages = [{ "content": "Hello, how are you?","role": "user"}]
# replicate llama-2 call
response = completion("replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf", messages)
```
### Replicate Models ### Replicate Models
liteLLM supports all replicate LLMs. For replicate models ensure to add a `replicate` prefix to the `model` arg. liteLLM detects it using this arg. liteLLM supports all replicate LLMs. For replicate models ensure to add a `replicate` prefix to the `model` arg. liteLLM detects it using this arg.
Below are examples on how to call replicate LLMs using liteLLM Below are examples on how to call replicate LLMs using liteLLM

22
litellm/main.py
View file

@ -116,7 +116,7 @@ def completion(
model model
] # update the model to the actual value if an alias has been passed in ] # update the model to the actual value if an alias has been passed in
model_response = ModelResponse() model_response = ModelResponse()
if deployment_id != None: if deployment_id != None: # azure llms
model=deployment_id model=deployment_id
custom_llm_provider="azure" custom_llm_provider="azure"
elif ( elif (
@ -124,10 +124,7 @@ def completion(
): # allow custom provider to be passed in via the model name "azure/chatgpt-test" ): # allow custom provider to be passed in via the model name "azure/chatgpt-test"
custom_llm_provider = model.split("/", 1)[0] custom_llm_provider = model.split("/", 1)[0]
model = model.split("/", 1)[1] model = model.split("/", 1)[1]
if (
"replicate" == custom_llm_provider and "/" not in model
): # handle the "replicate/llama2..." edge-case
model = custom_llm_provider + "/" + model
# check if user passed in any of the OpenAI optional params # check if user passed in any of the OpenAI optional params
optional_params = get_optional_params( optional_params = get_optional_params(
functions=functions, functions=functions,
@ -340,22 +337,19 @@ def completion(
model_response["model"] = model model_response["model"] = model
model_response["usage"] = response["usage"] model_response["usage"] = response["usage"]
response = model_response response = model_response
elif "replicate" in model or custom_llm_provider == "replicate": elif (
# import replicate/if it fails then pip install replicate "replicate" in model or
custom_llm_provider == "replicate" or
model in litellm.replicate_models
):
# Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN") # Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN")
replicate_key = os.environ.get("REPLICATE_API_TOKEN") replicate_key = None
if replicate_key == None:
# user did not set REPLICATE_API_TOKEN in .env
replicate_key = ( replicate_key = (
get_secret("REPLICATE_API_KEY") get_secret("REPLICATE_API_KEY")
or get_secret("REPLICATE_API_TOKEN") or get_secret("REPLICATE_API_TOKEN")
or api_key or api_key
or litellm.replicate_key or litellm.replicate_key
) )
# set replicate key
os.environ["REPLICATE_API_TOKEN"] = str(replicate_key)
model_response = replicate.completion( model_response = replicate.completion(
model=model, model=model,