Spaces:
Running
Running
File size: 2,595 Bytes
48ec86e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
# ruff: noqa: F401
import os
import sys
import rich
from loguru import logger
from smolagents import LiteLLMModel, OpenAIServerModel
from exit_gracefully import exit_gracefully
print = rich.get_console().print # noqa
def litellm_model(
model_id = None,
api_base = None,
api_key = None,
provider = "openai",
kwargs = None,
):
kwargs = kwargs or {}
# default siliconflow
api_base = api_base or "https://api.siliconflow.cn/v1"
api_key = api_key or os.getenv("SILICONFLOW_API_KEY")
model_id = model_id or "deepseek-ai/DeepSeek-V3"
if provider.lower() in ["gemini", "xai"]:
return LiteLLMModel(
f"{provider}/{model_id}",
# api_base=api_base, # must not set api_base to None or not set at all for gemini
api_key=api_key,
temperature=0.,
**kwargs,
)
else:
return LiteLLMModel(
f"{provider}/{model_id}",
api_base=api_base,
api_key=api_key,
temperature=0.,
**kwargs,
)
def main():
messages = [{'role': 'user', 'content': 'Say this is a test.'}]
logger.debug(sys.argv)
if not sys.argv[1:]:
model = litellm_model()
logger.debug(model(messages))
return
if len(sys.argv[1:]) < 4:
raise SystemExit("Provide at least four args (model_id, api_base, api_key, provider)")
model_id, api_base, api_key, provider, *_ = sys.argv[1:]
model = litellm_model(model_id, api_base, api_key, provider)
try:
response = model(messages)
logger.debug(response)
except Exception as e:
logger.error(e)
return
try:
print(response.content)
except Exception as e:
logger.error(e)
if __name__ == "__main__":
main()
# python litellm_model.py
# deepseek-ai/DeepSeek-V3 https://api.siliconflow.cn/v1 %SILICONFLOW_API_KEY%
# python litellm_model.py grok-3-beta https://api.x.ai/v1 %XAI_API_KEY%
# set proxy or use https://api-proxy.me/gemini/v1beta
# gemini-2.5-flash-preview-04-17 https://generativelanguage.googleapis.com/v1beta %GEMINI_API_KEY%
# python litellm_model.py gemini-2.5-flash-preview-04-17 https://api-proxy.me/gemini/v1beta %GEMINI_API_KEY%
# gemini-2.0-flash
# litellmmodel
# https://api.together.ai/models/deepseek-ai/DeepSeek-V3
# deepseek-ai/DeepSeek-V3 https://api.together.xyz/v1 %TOGETHER_API_KEY%
# deepseek-chat https://litellm.dattw.eu.org/v1 %LITELLM_API_KEY%
# LLM API proxy: https://linux.do/t/topic/290871
|