Add pyright to ci/cd + Fix remaining type-checking errors (BerriAI#6082)

* fix: fix type-checking errors

* fix: fix additional type-checking errors

* fix: additional type-checking error fixes

* fix: fix additional type-checking errors

* fix: additional type-check fixes

* fix: fix all type-checking errors + add pyright to ci/cd

* fix: fix incorrect import

* ci(config.yml): use mypy on ci/cd

* fix: fix type-checking errors in utils.py

* fix: fix all type-checking errors on main.py

* fix: fix mypy linting errors

* fix(anthropic/cost_calculator.py): fix linting errors

* fix: fix mypy linting errors

* fix: fix linting errors

Author: krrishdholakia

.circleci/config.yml

@@ -315,10 +315,10 @@ jobs:
             python -m pip install --upgrade pip
             pip install ruff
             pip install pylint  
+            pip install pyright
             pip install .
       - run: python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
       - run: ruff check ./litellm
-    
 
   build_and_test:
     machine:

enterprise/enterprise_callbacks/generic_api_callback.py

@@ -8,7 +8,7 @@
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.caching import DualCache
 
-from typing import Literal, Union
+from typing import Literal, Union, Optional
 
 import traceback
 
@@ -26,19 +26,25 @@
 
 class GenericAPILogger:
     # Class variables or attributes
-    def __init__(self, endpoint=None, headers=None):
+    def __init__(self, endpoint: Optional[str] = None, headers: Optional[dict] = None):
         try:
-            if endpoint == None:
+            if endpoint is None:
                 # check env for "GENERIC_LOGGER_ENDPOINT"
                 if os.getenv("GENERIC_LOGGER_ENDPOINT"):
                     # Do something with the endpoint
                     endpoint = os.getenv("GENERIC_LOGGER_ENDPOINT")
                 else:
                     # Handle the case when the endpoint is not found in the environment variables
                     raise ValueError(
-                        f"endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables"
+                        "endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables"
                     )
             headers = headers or litellm.generic_logger_headers
+
+            if endpoint is None:
+                raise ValueError("endpoint not set for GenericAPILogger")
+            if headers is None:
+                raise ValueError("headers not set for GenericAPILogger")
+
             self.endpoint = endpoint
             self.headers = headers
 

enterprise/enterprise_hooks/aporia_ai.py

@@ -48,8 +48,6 @@ def __init__(
         )
         self.aporia_api_key = api_key or os.environ["APORIO_API_KEY"]
         self.aporia_api_base = api_base or os.environ["APORIO_API_BASE"]
-        self.event_hook: GuardrailEventHooks
-
         super().__init__(**kwargs)
 
     #### CALL HOOKS - proxy only ####

enterprise/enterprise_hooks/blocked_user_list.py

@@ -84,7 +84,7 @@ async def async_pre_call_hook(
                     )
 
                 cache_key = f"litellm:end_user_id:{user}"
-                end_user_cache_obj: LiteLLM_EndUserTable = cache.get_cache(
+                end_user_cache_obj: Optional[LiteLLM_EndUserTable] = cache.get_cache(  # type: ignore
                     key=cache_key
                 )
                 if end_user_cache_obj is None and self.prisma_client is not None:

enterprise/enterprise_hooks/google_text_moderation.py

@@ -48,7 +48,7 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger):
     # Class variables or attributes
     def __init__(self):
         try:
-            from google.cloud import language_v1
+            from google.cloud import language_v1  # type: ignore
         except Exception:
             raise Exception(
                 "Missing google.cloud package. Run `pip install --upgrade google-cloud-language`"
@@ -57,8 +57,8 @@ def __init__(self):
         # Instantiates a client
         self.client = language_v1.LanguageServiceClient()
         self.moderate_text_request = language_v1.ModerateTextRequest
-        self.language_document = language_v1.types.Document
-        self.document_type = language_v1.types.Document.Type.PLAIN_TEXT
+        self.language_document = language_v1.types.Document  # type: ignore
+        self.document_type = language_v1.types.Document.Type.PLAIN_TEXT  # type: ignore
 
         default_confidence_threshold = (
             litellm.google_moderation_confidence_threshold or 0.8

enterprise/enterprise_hooks/llama_guard.py

@@ -8,6 +8,7 @@
 #  Thank you users! We ❤️ you! - Krrish & Ishaan
 
 import sys, os
+from collections.abc import Iterable
 
 sys.path.insert(
     0, os.path.abspath("../..")
@@ -19,11 +20,12 @@
 from litellm.integrations.custom_logger import CustomLogger
 from fastapi import HTTPException
 from litellm._logging import verbose_proxy_logger
-from litellm.utils import (
+from litellm.types.utils import (
     ModelResponse,
     EmbeddingResponse,
     ImageResponse,
     StreamingChoices,
+    Choices,
 )
 from datetime import datetime
 import aiohttp, asyncio
@@ -34,7 +36,10 @@
 class _ENTERPRISE_LlamaGuard(CustomLogger):
     # Class variables or attributes
     def __init__(self, model_name: Optional[str] = None):
-        self.model = model_name or litellm.llamaguard_model_name
+        _model = model_name or litellm.llamaguard_model_name
+        if _model is None:
+            raise ValueError("model_name not set for LlamaGuard")
+        self.model = _model
         file_path = litellm.llamaguard_unsafe_content_categories
         data = None
 
@@ -124,7 +129,13 @@ async def async_moderation_hook(
                 hf_model_name="meta-llama/LlamaGuard-7b",
             )
 
-            if "unsafe" in response.choices[0].message.content:
+            if (
+                isinstance(response, ModelResponse)
+                and isinstance(response.choices[0], Choices)
+                and response.choices[0].message.content is not None
+                and isinstance(response.choices[0].message.content, Iterable)
+                and "unsafe" in response.choices[0].message.content
+            ):
                 raise HTTPException(
                     status_code=400, detail={"error": "Violated content safety policy"}
                 )

enterprise/enterprise_hooks/llm_guard.py

@@ -8,7 +8,11 @@
 ## This provides an LLM Guard Integration for content moderation on the proxy
 
 from typing import Optional, Literal, Union
-import litellm, traceback, sys, uuid, os
+import litellm
+import traceback
+import sys
+import uuid
+import os
 from litellm.caching import DualCache
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.integrations.custom_logger import CustomLogger
@@ -21,8 +25,10 @@
     StreamingChoices,
 )
 from datetime import datetime
-import aiohttp, asyncio
+import aiohttp
+import asyncio
 from litellm.utils import get_formatted_prompt
+from litellm.secret_managers.main import get_secret_str
 
 litellm.set_verbose = True
 
@@ -38,7 +44,7 @@ def __init__(
         self.llm_guard_mode = litellm.llm_guard_mode
         if mock_testing == True:  # for testing purposes only
             return
-        self.llm_guard_api_base = litellm.get_secret("LLM_GUARD_API_BASE", None)
+        self.llm_guard_api_base = get_secret_str("LLM_GUARD_API_BASE", None)
         if self.llm_guard_api_base is None:
             raise Exception("Missing `LLM_GUARD_API_BASE` from environment")
         elif not self.llm_guard_api_base.endswith("/"):

enterprise/enterprise_hooks/openai_moderation.py

@@ -51,8 +51,8 @@ async def async_moderation_hook(  ### 👈 KEY CHANGE ###
             "audio_transcription",
         ],
     ):
+        text = ""
         if "messages" in data and isinstance(data["messages"], list):
-            text = ""
             for m in data["messages"]:  # assume messages is a list
                 if "content" in m and isinstance(m["content"], str):
                     text += m["content"]
@@ -67,7 +67,7 @@ async def async_moderation_hook(  ### 👈 KEY CHANGE ###
         )
 
         verbose_proxy_logger.debug("Moderation response: %s", moderation_response)
-        if moderation_response.results[0].flagged == True:
+        if moderation_response.results[0].flagged is True:
             raise HTTPException(
                 status_code=403, detail={"error": "Violated content safety policy"}
             )

enterprise/utils.py

@@ -6,7 +6,9 @@
 from datetime import datetime
 
 
-async def get_spend_by_tags(start_date=None, end_date=None, prisma_client=None):
+async def get_spend_by_tags(
+    prisma_client: PrismaClient, start_date=None, end_date=None
+):
     response = await prisma_client.db.query_raw(
         """
         SELECT

litellm/_redis.py

@@ -191,7 +191,7 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
         new_startup_nodes.append(ClusterNode(**item))
 
     redis_kwargs.pop("startup_nodes")
-    return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs)
+    return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs)  # type: ignore
 
 
 def _init_redis_sentinel(redis_kwargs) -> redis.Redis:

litellm/assistants/utils.py

@@ -1,5 +1,8 @@
-import litellm
 from typing import Optional, Union
+
+import litellm
+
+from ..exceptions import UnsupportedParamsError
 from ..types.llms.openai import *
 
 

litellm/caching.py

@@ -10,6 +10,7 @@
 import ast
 import asyncio
 import hashlib
+import inspect
 import io
 import json
 import logging
@@ -245,7 +246,8 @@ def __init__(
             self.redis_flush_size = redis_flush_size
         self.redis_version = "Unknown"
         try:
-            self.redis_version = self.redis_client.info()["redis_version"]
+            if not inspect.iscoroutinefunction(self.redis_client):
+                self.redis_version = self.redis_client.info()["redis_version"]  # type: ignore
         except Exception:
             pass
 
@@ -266,7 +268,8 @@ def __init__(
 
         ### SYNC HEALTH PING ###
         try:
-            self.redis_client.ping()
+            if hasattr(self.redis_client, "ping"):
+                self.redis_client.ping()  # type: ignore
         except Exception as e:
             verbose_logger.error(
                 "Error connecting to Sync Redis client", extra={"error": str(e)}
@@ -308,7 +311,7 @@ def increment_cache(
         _redis_client = self.redis_client
         start_time = time.time()
         try:
-            result = _redis_client.incr(name=key, amount=value)
+            result: int = _redis_client.incr(name=key, amount=value)  # type: ignore
 
             if ttl is not None:
                 # check if key already has ttl, if not -> set ttl
@@ -561,7 +564,7 @@ async def async_set_cache_sadd(
                 f"Set ASYNC Redis Cache: key: {key}\nValue {value}\nttl={ttl}"
             )
             try:
-                await redis_client.sadd(key, *value)
+                await redis_client.sadd(key, *value)  # type: ignore
                 if ttl is not None:
                     _td = timedelta(seconds=ttl)
                     await redis_client.expire(key, _td)
@@ -712,7 +715,7 @@ def batch_get_cache(self, key_list) -> dict:
             for cache_key in key_list:
                 cache_key = self.check_and_fix_namespace(key=cache_key)
                 _keys.append(cache_key)
-            results = self.redis_client.mget(keys=_keys)
+            results: List = self.redis_client.mget(keys=_keys)  # type: ignore
 
             # Associate the results back with their keys.
             # 'results' is a list of values corresponding to the order of keys in 'key_list'.
@@ -842,7 +845,7 @@ def sync_ping(self) -> bool:
         print_verbose("Pinging Sync Redis Cache")
         start_time = time.time()
         try:
-            response = self.redis_client.ping()
+            response: bool = self.redis_client.ping()  # type: ignore
             print_verbose(f"Redis Cache PING: {response}")
             ## LOGGING ##
             end_time = time.time()
@@ -911,8 +914,8 @@ async def delete_cache_keys(self, keys):
         async with _redis_client as redis_client:
             await redis_client.delete(*keys)
 
-    def client_list(self):
-        client_list = self.redis_client.client_list()
+    def client_list(self) -> List:
+        client_list: List = self.redis_client.client_list()  # type: ignore
         return client_list
 
     def info(self):

litellm/cost_calculator.py

@@ -39,8 +39,8 @@
 )
 from litellm.llms.OpenAI.cost_calculation import cost_per_token as openai_cost_per_token
 from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
-from litellm.rerank_api.types import RerankResponse
 from litellm.types.llms.openai import HttpxBinaryResponseContent
+from litellm.types.rerank import RerankResponse
 from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS
 from litellm.types.utils import PassthroughCallTypes, Usage
 from litellm.utils import (

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -39,11 +39,11 @@
     VirtualKeyEvent,
     WebhookEvent,
 )
+from litellm.types.integrations.slack_alerting import *
 from litellm.types.router import LiteLLM_Params
 
 from ..email_templates.templates import *
 from .batching_handler import send_to_webhook, squash_payloads
-from .types import *
 from .utils import _add_langfuse_trace_id_to_alert, process_slack_alerting_variables
 
 

litellm/integrations/custom_logger.py

@@ -172,14 +172,14 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
         ],
-    ):
+    ) -> Any:
         pass
 
     async def async_post_call_streaming_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
         response: str,
-    ):
+    ) -> Any:
         pass
 
     #### SINGLE-USE #### - https://docs.litellm.ai/docs/observability/custom_callback#using-your-custom-callback-function

litellm/integrations/email_alerting.py

@@ -42,8 +42,10 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list:
     )
     _team_member_user_ids: List[str] = []
     for member in _team_members:
-        if member and isinstance(member, dict) and member.get("user_id") is not None:
-            _team_member_user_ids.append(member.get("user_id"))
+        if member and isinstance(member, dict):
+            _user_id = member.get("user_id")
+            if _user_id and isinstance(_user_id, str):
+                _team_member_user_ids.append(_user_id)
 
     sql_query = """
         SELECT user_email

litellm/integrations/lunary.py

@@ -149,7 +149,7 @@ def log_event(
             else:
                 error_obj = None
 
-            self.lunary_client.track_event(
+            self.lunary_client.track_event(  # type: ignore
                 type,
                 "start",
                 run_id,
@@ -164,7 +164,7 @@ def log_event(
                 params=extra,
             )
 
-            self.lunary_client.track_event(
+            self.lunary_client.track_event(  # type: ignore
                 type,
                 event,
                 run_id,