andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-05-03 16:22:49 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

ruff updates

Browse Source
This commit is contained in:
Rip&Tear
2024-10-23 17:20:05 +08:00
parent 098a4312ab
commit 263544524d
1 changed files with 44 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/crewai/cli/provider.py
View File

@@ -7,6 +7,7 @@ from pathlib import Path
from crewai.cli.constants import PROVIDERS, MODELS, JSON_URL
def select_choice(prompt_message, choices):
"""
Presents a list of choices to the user and prompts them to select one.
@@ -18,17 +19,19 @@ def select_choice(prompt_message, choices):
Returns:
- str: The selected choice from the list, or None if the user chooses to quit.
"""
provider_models = get_provider_data()
if not provider_models:
return
click.secho(prompt_message, fg="cyan")
for idx, choice in enumerate(choices, start=1):
click.secho(f"{idx}. {choice}", fg="cyan")
click.secho("q. Quit", fg="cyan")
while True:
choice = click.prompt(
"Enter the number of your choice or 'q' to quit", type=str
)
choice = click.prompt("Enter the number of your choice or 'q' to quit", type=str)
if choice.lower() == "q":
if choice.lower() == 'q':
return None
try:
@@ -38,11 +41,7 @@ def select_choice(prompt_message, choices):
except ValueError:
pass
click.secho(
"Invalid selection. Please select a number between 1 and 6 or 'q' to quit.",
fg="red",
)
click.secho("Invalid selection. Please select a number between 1 and 6 or 'q' to quit.", fg="red")
def select_provider(provider_models):
"""
@@ -58,20 +57,17 @@ def select_provider(provider_models):
predefined_providers = [p.lower() for p in PROVIDERS]
all_providers = sorted(set(predefined_providers + list(provider_models.keys())))
provider = select_choice(
"Select a provider to set up:", predefined_providers + ["other"]
)
provider = select_choice("Select a provider to set up:", predefined_providers + ['other'])
if provider is None: # User typed 'q'
return None
if provider == "other":
if provider == 'other':
provider = select_choice("Select a provider from the full list:", all_providers)
if provider is None: # User typed 'q'
return None
return provider.lower() if provider else False
def select_model(provider, provider_models):
"""
Presents a list of models for a given provider to the user and prompts them to select one.
@@ -94,12 +90,9 @@ def select_model(provider, provider_models):
click.secho(f"No models available for provider '{provider}'.", fg="red")
return None
selected_model = select_choice(
f"Select a model to use for {provider.capitalize()}:", available_models
)
selected_model = select_choice(f"Select a model to use for {provider.capitalize()}:", available_models)
return selected_model
def load_provider_data(cache_file, cache_expiry):
"""
Loads provider data from a cache file if it exists and is not expired. If the cache is expired or corrupted, it fetches the data from the web.
@@ -112,24 +105,15 @@ def load_provider_data(cache_file, cache_expiry):
- dict or None: The loaded provider data or None if the operation fails.
"""
current_time = time.time()
if (
cache_file.exists()
and (current_time - cache_file.stat().st_mtime) < cache_expiry
):
if cache_file.exists() and (current_time - cache_file.stat().st_mtime) < cache_expiry:
data = read_cache_file(cache_file)
if data:
return data
click.secho(
"Cache is corrupted. Fetching provider data from the web...", fg="yellow"
)
click.secho("Cache is corrupted. Fetching provider data from the web...", fg="yellow")
else:
click.secho(
"Cache expired or not found. Fetching provider data from the web...",
fg="cyan",
)
click.secho("Cache expired or not found. Fetching provider data from the web...", fg="cyan")
return fetch_provider_data(cache_file)
def read_cache_file(cache_file):
"""
Reads and returns the JSON content from a cache file. Returns None if the file contains invalid JSON.
@@ -146,7 +130,6 @@ def read_cache_file(cache_file):
except json.JSONDecodeError:
return None
def fetch_provider_data(cache_file):
"""
Fetches provider data from a specified URL and caches it to a file.
@@ -170,7 +153,6 @@ def fetch_provider_data(cache_file):
click.secho("Error parsing provider data. Invalid JSON format.", fg="red")
return None
def download_data(response):
"""
Downloads data from a given HTTP response and returns the JSON content.
@@ -181,19 +163,16 @@ def download_data(response):
Returns:
- dict: The JSON content of the response.
"""
total_size = int(response.headers.get("content-length", 0))
total_size = int(response.headers.get('content-length', 0))
block_size = 8192
data_chunks = []
with click.progressbar(
length=total_size, label="Downloading", show_pos=True
) as progress_bar:
with click.progressbar(length=total_size, label='Downloading', show_pos=True) as progress_bar:
for chunk in response.iter_content(block_size):
if chunk:
data_chunks.append(chunk)
progress_bar.update(len(chunk))
data_content = b"".join(data_chunks)
return json.loads(data_content.decode("utf-8"))
data_content = b''.join(data_chunks)
return json.loads(data_content.decode('utf-8'))
def get_provider_data():
"""
@@ -202,9 +181,9 @@ def get_provider_data():
Returns:
- dict or None: A dictionary of providers mapped to their models or None if the operation fails.
"""
cache_dir = Path.home() / ".crewai"
cache_dir = Path.home() / '.crewai'
cache_dir.mkdir(exist_ok=True)
cache_file = cache_dir / "provider_cache.json"
cache_file = cache_dir / 'provider_cache.json'
cache_expiry = 24 * 3600
data = load_provider_data(cache_file, cache_expiry)
@@ -214,7 +193,7 @@ def get_provider_data():
provider_models = defaultdict(list)
for model_name, properties in data.items():
provider = properties.get("litellm_provider", "").strip().lower()
if "http" in provider or provider == "other":
if 'http' in provider or provider == 'other':
continue
if provider:
provider_models[provider].append(model_name)