forked from phoenix/litellm-mirror
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
GitHub
parent
3fc4ae0d65
commit
d57be47b0f
263 changed files with 1687 additions and 3320 deletions
|
|
@ -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={
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue