Merge pull request #135 from amazeeio/async-litellm

Async litellm

Author: PhiRho

app/core/resource_limits.py

@@ -28,7 +28,7 @@ def check_team_user_limit(db: Session, team_id: int) -> None:
     """
     # Get current user count and max allowed users in a single query
     result = db.query(
-        func.count(DBUser.id).label('current_user_count'),
+        func.count(func.distinct(DBUser.id)).label('current_user_count'),
         func.coalesce(func.max(DBProduct.user_count), DEFAULT_USER_COUNT).label('max_users')
     ).select_from(DBUser).filter(
         DBUser.team_id == team_id
@@ -64,14 +64,14 @@ def check_key_limits(db: Session, team_id: int, owner_id: Optional[int] = None)
         func.coalesce(func.max(DBProduct.total_key_count), DEFAULT_TOTAL_KEYS).label('max_total_keys'),
         func.coalesce(func.max(DBProduct.keys_per_user), DEFAULT_KEYS_PER_USER).label('max_keys_per_user'),
         func.coalesce(func.max(DBProduct.service_key_count), DEFAULT_SERVICE_KEYS).label('max_service_keys'),
-        func.count(DBPrivateAIKey.id).filter(
+        func.count(func.distinct(DBPrivateAIKey.id)).filter(
             DBPrivateAIKey.litellm_token.isnot(None)
         ).label('current_team_keys'),
-        func.count(DBPrivateAIKey.id).filter(
+        func.count(func.distinct(DBPrivateAIKey.id)).filter(
             DBPrivateAIKey.owner_id == owner_id,
             DBPrivateAIKey.litellm_token.isnot(None)
         ).label('current_user_keys') if owner_id else None,
-        func.count(DBPrivateAIKey.id).filter(
+        func.count(func.distinct(DBPrivateAIKey.id)).filter(
             DBPrivateAIKey.owner_id.is_(None),
             DBPrivateAIKey.litellm_token.isnot(None)
         ).label('current_service_keys')
@@ -126,7 +126,7 @@ def check_vector_db_limits(db: Session, team_id: int) -> None:
     # Get vector DB limits and current count in a single query
     result = db.query(
         func.coalesce(func.max(DBProduct.vector_db_count), DEFAULT_VECTOR_DB_COUNT).label('max_vector_db_count'),
-        func.count(DBPrivateAIKey.id).filter(
+        func.count(func.distinct(DBPrivateAIKey.id)).filter(
             DBPrivateAIKey.database_name.isnot(None)
         ).label('current_vector_db_count')
     ).select_from(DBTeam).filter(

app/services/litellm.py

@@ -1,4 +1,4 @@
-import requests
+import httpx
 from fastapi import HTTPException, status
 import logging
 from app.core.resource_limits import DEFAULT_KEY_DURATION, DEFAULT_MAX_SPEND, DEFAULT_RPM_PER_KEY
@@ -57,20 +57,21 @@ async def create_key(self, email: str, name: str, user_id: int, team_id: str, du
             if user_id is not None:
                 request_data["user_id"] = str(user_id)
 
-            logger.info(f"Making request to LiteLLM API to generate key with data: {request_data}")
-            response = requests.post(
-                f"{self.api_url}/key/generate",
-                json=request_data,
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                }
-            )
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/generate",
+                    json=request_data,
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    }
+                )
 
-            response.raise_for_status()
-            key = response.json()["key"]
-            logger.info("Successfully generated new LiteLLM API key")
-            return key
-        except requests.exceptions.RequestException as e:
+                response.raise_for_status()
+                response_data = response.json()
+                key = response_data["key"]
+                logger.info("Successfully generated new LiteLLM API key")
+                return key
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:
@@ -87,21 +88,22 @@ async def create_key(self, email: str, name: str, user_id: int, team_id: str, du
     async def delete_key(self, key: str) -> bool:
         """Delete a LiteLLM API key"""
         try:
-            response = requests.post(
-                f"{self.api_url}/key/delete",
-                json={"keys": [key]},  # API expects an array of keys
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                }
-            )
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/delete",
+                    json={"keys": [key]},  # API expects an array of keys
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    }
+                )
 
-            # Treat 404 (key not found) as success
-            if response.status_code == 404:
-                return True
+                # Treat 404 (key not found) as success
+                if response.status_code == 404:
+                    return True
 
-            response.raise_for_status()
-            return True
-        except requests.exceptions.RequestException as e:
+                response.raise_for_status()
+                return True
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:
@@ -118,19 +120,21 @@ async def delete_key(self, key: str) -> bool:
     async def get_key_info(self, litellm_token: str) -> dict:
         """Get information about a LiteLLM API key"""
         try:
-            response = requests.get(
-                f"{self.api_url}/key/info",
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                },
-                params={
-                    "key": litellm_token
-                }
-            )
-            response.raise_for_status()
-            logger.info(f"LiteLLM key information: {response.json()}")
-            return response.json()
-        except requests.exceptions.RequestException as e:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    f"{self.api_url}/key/info",
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    },
+                    params={
+                        "key": litellm_token
+                    }
+                )
+                response.raise_for_status()
+                response_data = response.json()
+                logger.info(f"LiteLLM key information: {response_data}")
+                return response_data
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             logger.error(f"Error getting LiteLLM key information: {error_msg}")
             if hasattr(e, 'response') and e.response is not None:
@@ -156,15 +160,16 @@ async def update_budget(self, litellm_token: str, budget_duration: str, budget_a
             if budget_amount:
                 request_data["max_budget"] = budget_amount
 
-            response = requests.post(
-                f"{self.api_url}/key/update",
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                },
-                json=request_data
-            )
-            response.raise_for_status()
-        except requests.exceptions.RequestException as e:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/update",
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    },
+                    json=request_data
+                )
+                response.raise_for_status()
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:
@@ -180,18 +185,19 @@ async def update_budget(self, litellm_token: str, budget_duration: str, budget_a
     async def update_key_duration(self, litellm_token: str, duration: str):
         """Update the duration for a LiteLLM API key"""
         try:
-            response = requests.post(
-                f"{self.api_url}/key/update",
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                },
-                json={
-                    "key": litellm_token,
-                    "duration": duration
-                }
-            )
-            response.raise_for_status()
-        except requests.exceptions.RequestException as e:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/update",
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    },
+                    json={
+                        "key": litellm_token,
+                        "duration": duration
+                    }
+                )
+                response.raise_for_status()
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:
@@ -207,21 +213,22 @@ async def update_key_duration(self, litellm_token: str, duration: str):
     async def set_key_restrictions(self, litellm_token: str, duration: str, budget_amount: float, rpm_limit: int, budget_duration: Optional[str] = None):
         """Set the restrictions for a LiteLLM API key"""
         try:
-            response = requests.post(
-                f"{self.api_url}/key/update",
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                },
-                json={
-                    "key": litellm_token,
-                    "duration": duration,
-                    "budget_duration": budget_duration,
-                    "max_budget": budget_amount,
-                    "rpm_limit": rpm_limit
-                }
-            )
-            response.raise_for_status()
-        except requests.exceptions.RequestException as e:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/update",
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    },
+                    json={
+                        "key": litellm_token,
+                        "duration": duration,
+                        "budget_duration": budget_duration,
+                        "max_budget": budget_amount,
+                        "rpm_limit": rpm_limit
+                    }
+                )
+                response.raise_for_status()
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:
@@ -237,18 +244,19 @@ async def set_key_restrictions(self, litellm_token: str, duration: str, budget_a
     async def update_key_team_association(self, litellm_token: str, new_team_id: str):
         """Update the team association for a LiteLLM API key"""
         try:
-            response = requests.post(
-                f"{self.api_url}/key/update",
-                headers={
-                    "Authorization": f"Bearer {self.master_key}"
-                },
-                json={
-                    "key": litellm_token,
-                    "team_id": new_team_id
-                }
-            )
-            response.raise_for_status()
-        except requests.exceptions.RequestException as e:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.api_url}/key/update",
+                    headers={
+                        "Authorization": f"Bearer {self.master_key}"
+                    },
+                    json={
+                        "key": litellm_token,
+                        "team_id": new_team_id
+                    }
+                )
+                response.raise_for_status()
+        except httpx.HTTPStatusError as e:
             error_msg = str(e)
             if hasattr(e, 'response') and e.response is not None:
                 try:

requirements.txt

@@ -18,4 +18,5 @@ prometheus-client==0.21.1
 prometheus-fastapi-instrumentator==7.0.0
 stripe==12.1.0
 six==1.17.0
-apscheduler==3.10.4
+apscheduler==3.10.4
+httpx==0.28.1

scripts/add_test_data.py

@@ -4,14 +4,16 @@
 import sys
 from datetime import datetime, timedelta, UTC
 import random
+import asyncio
 
 # Add the parent directory to the Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from sqlalchemy.orm import sessionmaker, Session
 from app.db.database import engine
-from app.db.models import DBTeam, DBUser, DBProduct, DBTeamProduct
+from app.db.models import DBTeam, DBUser, DBProduct, DBTeamProduct, DBRegion, DBPrivateAIKey
 from app.core.security import get_password_hash
+from app.services.litellm import LiteLLMService
 
 def create_test_data():
     """Create test data for teams, users, and products"""
@@ -264,10 +266,49 @@ def create_test_data():
     finally:
         db.close()
 
+async def create_test_keys(count: int):
+    # Create database session
+    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+    db = SessionLocal()
+
+    try:
+        team = db.query(DBTeam).first()
+        region = db.query(DBRegion).filter(DBRegion.is_active == True).first()
+        if not region:
+            print(f"No active regions, not creating test keys")
+            return
+        litellm = LiteLLMService(region.litellm_api_url, region.litellm_api_key)
+        team_id = team.id
+        for i in range(0, count):
+            key_name = f"auto_test_{i}"
+            litellm_token = await litellm.create_key(
+                email=team.admin_email,
+                name=key_name,
+                user_id=team_id,
+                team_id=LiteLLMService.format_team_id(region.name, team_id),
+            )
+
+            # Create response object
+            db_token = DBPrivateAIKey(
+                litellm_token=litellm_token,
+                litellm_api_url=region.litellm_api_url,
+                owner_id=None,
+                team_id=None if team_id is None else team_id,
+                name=key_name,
+                region_id = region.id
+            )
+            db.add(db_token)
+            db.commit()
+            print(f"Created LLM token {key_name} int team {team.name}")
+
+    except Exception as e:
+        print(f"failed to create test keys {str(e)}")
+
 def main():
     """Main function to run the script"""
     try:
         create_test_data()
+        asyncio.run(create_test_keys(50))
     except Exception as e:
         print(f"Script failed: {str(e)}")
         sys.exit(1)