Defer litellm import until first chat message; only import streamlit if gui is activated

This commit is contained in:
Paul Gauthier 2024-07-03 13:04:13 -03:00
parent ee203deef0
commit 2fc358a02f
2 changed files with 25 additions and 22 deletions

View file

@ -8,7 +8,6 @@ import git
import httpx import httpx
from dotenv import load_dotenv from dotenv import load_dotenv
from prompt_toolkit.enums import EditingMode from prompt_toolkit.enums import EditingMode
from streamlit.web import cli
from aider import __version__, models, utils from aider import __version__, models, utils
from aider.args import get_parser from aider.args import get_parser
@ -150,6 +149,8 @@ def scrub_sensitive_info(args, text):
def launch_gui(args): def launch_gui(args):
from streamlit.web import cli
from aider import gui from aider import gui
print() print()

View file

@ -1,9 +1,11 @@
import difflib import difflib
import importlib
import json import json
import math import math
import os import os
import sys import sys
from dataclasses import dataclass, fields from dataclasses import dataclass, fields
from pathlib import Path
from typing import Optional from typing import Optional
import yaml import yaml
@ -40,7 +42,7 @@ gpt-3.5-turbo-16k
gpt-3.5-turbo-16k-0613 gpt-3.5-turbo-16k-0613
""" """
OPENAI_MODELS = [ln.strip for ln in OPENAI_MODELS.splitlines() if ln.strip()] OPENAI_MODELS = [ln.strip() for ln in OPENAI_MODELS.splitlines() if ln.strip()]
ANTHROPIC_MODELS = """ ANTHROPIC_MODELS = """
claude-2 claude-2
@ -51,7 +53,7 @@ claude-3-sonnet-20240229
claude-3-5-sonnet-20240620 claude-3-5-sonnet-20240620
""" """
ANTHROPIC_MODELS = [ln.strip for ln in ANTHROPIC_MODELS.splitlines() if ln.strip()] ANTHROPIC_MODELS = [ln.strip() for ln in ANTHROPIC_MODELS.splitlines() if ln.strip()]
@dataclass @dataclass
@ -365,25 +367,7 @@ class Model:
def __init__(self, model, weak_model=None): def __init__(self, model, weak_model=None):
self.name = model self.name = model
# Do we have the model_info? self.info = self.get_model_info(model)
try:
self.info = litellm.get_model_info(model)
except Exception:
self.info = dict()
if not self.info and "gpt-4o" in self.name:
self.info = {
"max_tokens": 4096,
"max_input_tokens": 128000,
"max_output_tokens": 4096,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 1.5e-5,
"litellm_provider": "openai",
"mode": "chat",
"supports_function_calling": True,
"supports_parallel_function_calling": True,
"supports_vision": True,
}
# Are all needed keys/params available? # Are all needed keys/params available?
res = self.validate_environment() res = self.validate_environment()
@ -404,6 +388,24 @@ class Model:
else: else:
self.get_weak_model(weak_model) self.get_weak_model(weak_model)
def get_model_info(self, model):
# Try and do this quickly, without triggering the litellm import
spec = importlib.util.find_spec("litellm")
if spec:
origin = Path(spec.origin)
fname = origin.parent / "model_prices_and_context_window_backup.json"
if fname.exists():
data = json.loads(fname.read_text())
info = data.get(model)
if info:
return info
# Do it the slow way...
try:
self.info = litellm.get_model_info(model)
except Exception:
self.info = dict()
def configure_model_settings(self, model): def configure_model_settings(self, model):
for ms in MODEL_SETTINGS: for ms in MODEL_SETTINGS:
# direct match, or match "provider/<model>" # direct match, or match "provider/<model>"