Litellm ruff linting enforcement (#5992)

* ci(config.yml): add a 'check_code_quality' step

Addresses https://github.com/BerriAI/litellm/issues/5991

* ci(config.yml): check why circle ci doesn't pick up this test

* ci(config.yml): fix to run 'check_code_quality' tests

* fix(__init__.py): fix unprotected import

* fix(__init__.py): don't remove unused imports

* build(ruff.toml): update ruff.toml to ignore unused imports

* fix: fix: ruff + pyright - fix linting + type-checking errors

* fix: fix linting errors

* fix(lago.py): fix module init error

* fix: fix linting errors

* ci(config.yml): cd into correct dir for checks

* fix(proxy_server.py): fix linting error

* fix(utils.py): fix bare except

causes ruff linting errors

* fix: ruff - fix remaining linting errors

* fix(clickhouse.py): use standard logging object

* fix(__init__.py): fix unprotected import

* fix: ruff - fix linting errors

* fix: fix linting errors

* ci(config.yml): cleanup code qa step (formatting handled in local_testing)

* fix(_health_endpoints.py): fix ruff linting errors

* ci(config.yml): just use ruff in check_code_quality pipeline for now

* build(custom_guardrail.py): include missing file

* style(embedding_handler.py): fix ruff check
This commit is contained in:
Krish Dholakia 2024-10-01 16:44:20 -07:00 committed by GitHub
parent 4fa8991a90
commit 94a05ca5d0
263 changed files with 1687 additions and 3320 deletions

View file

@ -26,11 +26,16 @@ from typing import (
overload,
)
import backoff
try:
import backoff
except ImportError:
raise ImportError(
"backoff is not installed. Please install it via 'pip install backoff'"
)
import httpx
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
from typing_extensions import overload
import litellm
import litellm.litellm_core_utils
@ -1049,7 +1054,7 @@ class PrismaClient:
)
try:
from prisma import Prisma # type: ignore
except Exception as e:
except Exception:
os.environ["DATABASE_URL"] = database_url
# Save the current working directory
original_dir = os.getcwd()
@ -1105,7 +1110,7 @@ class PrismaClient:
if isinstance(v, dict):
try:
db_data[k] = json.dumps(v)
except:
except Exception:
# This avoids Prisma retrying this 5 times, and making 5 clients
db_data[k] = "failed-to-serialize-json"
return db_data
@ -1204,7 +1209,7 @@ class PrismaClient:
)
)
except Exception as e:
except Exception:
raise
# try:
@ -1990,7 +1995,7 @@ class PrismaClient:
if response is not None:
try:
_data = response.model_dump() # type: ignore
except Exception as e:
except Exception:
_data = response.dict()
return {"token": token, "data": _data}
elif (
@ -2083,7 +2088,7 @@ class PrismaClient:
data_json = self.jsonify_object(
data=t.model_dump(exclude_none=True)
)
except:
except Exception:
data_json = self.jsonify_object(data=t.dict(exclude_none=True))
batcher.litellm_verificationtoken.update(
where={"token": t.token}, # type: ignore
@ -2109,7 +2114,7 @@ class PrismaClient:
data_json = self.jsonify_object(
data=user.model_dump(exclude_none=True)
)
except:
except Exception:
data_json = self.jsonify_object(data=user.dict())
batcher.litellm_usertable.upsert(
where={"user_id": user.user_id}, # type: ignore
@ -2138,7 +2143,7 @@ class PrismaClient:
data_json = self.jsonify_object(
data=team.model_dump(exclude_none=True)
)
except:
except Exception:
data_json = self.jsonify_object(
data=team.dict(exclude_none=True)
)
@ -2268,7 +2273,7 @@ class PrismaClient:
verbose_proxy_logger.debug(
"PrismaClient: connect() called Attempting to Connect to DB"
)
if self.db.is_connected() == False:
if self.db.is_connected() is False:
verbose_proxy_logger.debug(
"PrismaClient: DB not connected, Attempting to Connect to DB"
)
@ -2384,69 +2389,6 @@ class PrismaClient:
return
class DBClient:
"""
Routes requests for CustomAuth
[TODO] route b/w customauth and prisma
"""
def __init__(
self, custom_db_type: Literal["dynamo_db"], custom_db_args: dict
) -> None:
if custom_db_type == "dynamo_db":
from litellm.proxy.db.dynamo_db import DynamoDBWrapper
self.db = DynamoDBWrapper(database_arguments=DynamoDBArgs(**custom_db_args))
async def get_data(self, key: str, table_name: Literal["user", "key", "config"]):
"""
Check if key valid
"""
return await self.db.get_data(key=key, table_name=table_name)
async def insert_data(
self, value: Any, table_name: Literal["user", "key", "config"]
):
"""
For new key / user logic
"""
return await self.db.insert_data(value=value, table_name=table_name)
async def update_data(
self, key: str, value: Any, table_name: Literal["user", "key", "config"]
):
"""
For cost tracking logic
key - hash_key value \n
value - dict with updated values
"""
return await self.db.update_data(key=key, value=value, table_name=table_name)
async def delete_data(
self, keys: List[str], table_name: Literal["user", "key", "config"]
):
"""
For /key/delete endpoints
"""
return await self.db.delete_data(keys=keys, table_name=table_name)
async def connect(self):
"""
For connecting to db and creating / updating any tables
"""
return await self.db.connect()
async def disconnect(self):
"""
For closing connection on server shutdown
"""
if self.db is not None:
return await self.db.disconnect() # type: ignore
return asyncio.sleep(0) # Return a dummy coroutine if self.db is None
### CUSTOM FILE ###
def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
module_name = value
@ -2494,9 +2436,7 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
### HELPER FUNCTIONS ###
async def _cache_user_row(
user_id: str, cache: DualCache, db: Union[PrismaClient, DBClient]
):
async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient):
"""
Check if a user_id exists in cache,
if not retrieve it.
@ -2504,10 +2444,7 @@ async def _cache_user_row(
cache_key = f"{user_id}_user_api_key_user_id"
response = cache.get_cache(key=cache_key)
if response is None: # Cache miss
if isinstance(db, PrismaClient):
user_row = await db.get_data(user_id=user_id)
elif isinstance(db, DBClient):
user_row = await db.get_data(key=user_id, table_name="user")
user_row = await db.get_data(user_id=user_id)
if user_row is not None:
print_verbose(f"User Row: {user_row}, type = {type(user_row)}")
if hasattr(user_row, "model_dump_json") and callable(
@ -2815,9 +2752,8 @@ async def update_spend(
end_user_id,
response_cost,
) in prisma_client.end_user_list_transactons.items():
max_end_user_budget = None
if litellm.max_end_user_budget is not None:
max_end_user_budget = litellm.max_end_user_budget
pass
batcher.litellm_endusertable.upsert(
where={"user_id": end_user_id},
data={