File size: 1,488 Bytes
dfd6145 |
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 |
import os
from smolagents import CodeAgent, LiteLLMRouterModel, WebSearchTool
# Make sure to setup the necessary environment variables!
llm_loadbalancer_model_list = [
{
"model_name": "model-group-1",
"litellm_params": {
"model": "gpt-4o-mini",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "model-group-1",
"litellm_params": {
"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
"aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
"aws_region_name": os.getenv("AWS_REGION"),
},
},
# {
# "model_name": "model-group-2",
# "litellm_params": {
# "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
# "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
# "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
# "aws_region_name": os.getenv("AWS_REGION"),
# },
# },
]
model = LiteLLMRouterModel(
model_id="model-group-1",
model_list=llm_loadbalancer_model_list,
client_kwargs={"routing_strategy": "simple-shuffle"},
)
agent = CodeAgent(tools=[WebSearchTool()], model=model, stream_outputs=True, return_full_result=True)
full_result = agent.run("How many seconds would it take for a leopard at full speed to run through Pont des Arts?")
print(full_result)
|