Begin the Stripe Integration #56
Conversation
- Introduced a new PUT endpoint `/private-ai-keys/{key_id}/extend-token-life` to allow users to extend the life of their private AI keys.
- Added `TokenDurationUpdate` schema for specifying the duration to extend.
- Implemented logic to verify user access and update the key duration in LiteLLM.
- Enhanced `LiteLLMService` with a method to update key duration and handle potential errors.
- Updated tests to validate the new functionality and ensure correct duration handling based on environment variables.
- Updated `requests-mock` to version 1.12.1 in `requirements-test.txt`. - Upgraded several packages in `requirements.txt`: - `uvicorn` to 0.34.2 - `sqlalchemy` to 2.0.41 - `pydantic` to 2.11.4 - `pydantic-settings` to 2.9.1 - `asyncpg` to 0.30.0 - `psycopg2-binary` to 2.9.10 - `requests` to 2.32.3 - `python-dotenv` to 1.1.0 - `boto3` to 1.38.20 - `markdown` to 3.8.0 - `prometheus-client` to 0.21.1 - `prometheus-fastapi-instrumentator` to 7.0.0
- Introduced a new GET endpoint `/private-ai-keys/{key_id}` to fetch detailed information about a specific private AI key.
- Implemented access control to ensure only system administrators can access this endpoint.
- Enhanced the `PrivateAIKeyDetail` schema to include additional fields such as spend, key name, key alias, and LiteLLM-specific data.
- Updated logging to capture key retrieval processes and potential errors.
- Add basic test cases for the new API
- Introduced a new billing module with endpoints for creating checkout sessions and handling Stripe webhook events. - Added a new `billing.py` API file to manage team subscriptions and payment processing. - Implemented `create_checkout_session` and `handle_stripe_event` functions in the Stripe service for managing subscriptions. - Updated the main application to include the billing router and integrated it with existing middleware. - Added a new `DBSystemSecret` model to store sensitive Stripe webhook secrets. - Enhanced the configuration settings to include AWS and SES parameters. - Created a migration script to add the `system_secrets` table to the database. - Updated the initialization script to set up the Stripe webhook in specific environments. - Added a new `CheckoutSessionCreate` schema for validating checkout session requests.
- Updated the `APIRouter` instances in multiple API files to use consistent lowercase tags. - Changed tags for `auth`, `private-ai-keys`, `regions`, `teams`, and `users` to improve clarity and maintainability. - Removed redundant router definitions in `private_ai_keys.py` to streamline the codebase.
- Introduced a new `products` API module with endpoints for creating, listing, retrieving, updating, and deleting products. - Implemented access control to restrict product management actions to system and team administrators. - Added a new `DBProduct` model to represent products in the database, including fields for name, stripe lookup key, and active status. - Updated the main application to include the products router. - Enhanced configuration settings to include a Stripe key for product management. - Created migration scripts to add the `products` table and update the `teams` table with a new `stripe_customer_id` field. - Developed comprehensive tests to validate product management functionality and access control.
- Refactored the billing API to handle Stripe webhook events asynchronously using background tasks. - Introduced a new function to retrieve Stripe events and updated the webhook setup to use a consistent key. - Added comprehensive tests for handling various Stripe events, including successful checkouts and subscription deletions. - Created a new diagram to illustrate the Stripe flow for better understanding of the billing process.
- Added a new `DBTeamProduct` model to establish a many-to-many relationship between teams and products. - Enhanced the `products` API to support new fields for product management, including user limits and budget settings. - Introduced a `check_team_user_limit` function to enforce user limits based on active products. - Updated billing API to handle product application for teams, including extending key durations and setting budgets. - Created comprehensive tests for product management, user limits, and billing integration. - Refactored existing code to improve maintainability and clarity.
- Introduced a new `check_key_limits` function to validate LLM token creation against team and user limits. - Updated `check_team_user_limit` to remove default limit checks for teams without products. - Enhanced tests to cover scenarios for creating LLM tokens, including exceeding total, user, and service key limits. - Refactored existing tests to utilize the new key limit checks and ensure comprehensive coverage for user and key limits.
- Introduced a new `check_vector_db_limits` function to validate vector DB creation against team limits. - Updated resource limits tests to include scenarios for creating vector DBs, including exceeding limits and handling user-owned keys. - Refactored existing tests for clarity and to ensure comprehensive coverage of resource limits functionality.
- Updated the `get_portal` endpoint to create a Stripe customer if one does not exist for the team. - Refactored the `create_portal_session` function to accept a Stripe customer ID and return URL. - Improved error handling for team not found scenarios in the billing API. - Added comprehensive tests for the `get_portal` functionality, covering both existing and new customer cases.
- Refactored the billing API to improve handling of various Stripe webhook events, including success and failure scenarios for checkout sessions and subscriptions. - Updated the event processing logic to utilize a unified list of known events, enhancing clarity and maintainability. - Introduced a new `remove_product_from_team` function to manage product removals based on event types. - Renamed functions for consistency, such as `get_product_id_from_sub` to `get_product_id_from_subscription`. - Added comprehensive tests for new event handling scenarios, ensuring robust coverage for product application and removal processes.
…ing API
- Introduced a new endpoint `/teams/{team_id}/pricing-table-session` to create a Stripe Customer Session for team subscription management.
- Enhanced error handling for scenarios where the team is not found and when the Stripe webhook secret is not configured.
- Updated the subscription success events to include `customer.subscription.resumed`.
- Added comprehensive tests for the new pricing table session functionality, covering both existing and new customer cases, as well as error scenarios.
- Added functionality to enforce limits on team users and LLM tokens based on configuration settings. - Introduced `check_key_limits` and `check_team_user_limit` functions to validate resource usage against defined limits. - Updated the `create_user` and `create_llm_token` endpoints to incorporate these checks. - Enhanced the `check_vector_db_limits` function to validate vector DB creation against team limits. - Refactored existing tests to ensure comprehensive coverage for new limit checks and updated error handling for limit exceedances.
- Added a new page for product management in the frontend, allowing users to create, edit, and delete products with detailed forms and validation. - Updated the backend to include checks preventing the deletion of products associated with teams. - Introduced a new function to retrieve the Stripe customer ID from payment intents. - Enhanced the README with instructions for testing Stripe integration. - Added tests to ensure product creation and deletion functionality, including checks for inconsistent key counts and team associations.
- Updated the `create_llm_token` function to check key limits based on team and user configurations. - Modified the `Settings` class to allow environment variable overrides for limit settings. - Adjusted default user count limit to improve resource management. - Enhanced tests to validate behavior when limits are enabled, ensuring proper error handling for exceeding key creation limits. - Removed outdated test for inconsistent key counts in product creation.
- Introduced a new `get_token_restrictions` function to retrieve token limits based on team products and payment history. - Updated the `create_llm_token` function to apply token restrictions dynamically based on the retrieved values. - Refactored the `LiteLLMService` to include a new method for setting key restrictions, consolidating budget and duration updates. - Enhanced tests to validate the behavior of token restrictions under various scenarios, ensuring proper handling of default and product-specific limits. - Updated existing tests to reflect changes in key management and resource limits functionality.
Also update event management and improve billing APIs
This listing allows admins to see which products a team is subscribed to.
There was a problem hiding this comment.
As always i left mostly drive-by commentary and didn't review front end or tests closely.
I think my biggest pieces of feedback are:
- Maybe try to reduce DB queries (especially in the resource limit checks)
- Maybe make the Stripe webhook use cases you expect to deal with the various business scenarios (new sub, renewal, upgrade, downgrade, cancellation) a bit more explicit.
In terms of your questions from Slack:
Should the background processor for webhook events stay in the API file, or move to the worker file?
I'd vote move them.
I'd prefer to have the Webhook endpoint configured dynamically so that the signing secret is never exposed (outside of a local testing environment), does that fit with normal practice, or should I just accept the manual step?
I think either is fine. I've seen apps put it in an env var or in the DB. The Stripe private API key is a much more important secret than the webhook signing secret.
app/api/billing.py
Outdated
| """ | ||
| Background task to handle Stripe webhook events. | ||
| This runs in a separate thread to avoid blocking the webhook response. | ||
| """ |
There was a problem hiding this comment.
I'm not sure what the fastAPI conventions are, but if we stick with this setup I wonder if it makes sense to pull this out into a different file? just to sort of signal that it's not a route / http handler.
There was a problem hiding this comment.
Oh, ha I re-read your slack message after writing this. Yeah, if it's gonna stay like this I'd suggest moving it I think.
app/api/billing.py
Outdated
|
|
||
| try: | ||
| # Create Stripe customer if one doesn't exist | ||
| if not team.stripe_customer_id: |
There was a problem hiding this comment.
when would a team go to the billing portal without an existing customer (or subscription)? i had the sense it was only useful after you've already gone through a subscription workflow first.
app/api/billing.py
Outdated
| db.commit() | ||
|
|
||
| # Get the frontend URL from environment | ||
| frontend_url = os.getenv("FRONTEND_URL", "http://localhost:3000") |
There was a problem hiding this comment.
nit: could extract this bit to a helper function, just in case it ever gets more complicated or changes later.
| logger.info(f"Creating Stripe customer for team {team.id}") | ||
| team.stripe_customer_id = await create_stripe_customer(team) | ||
| db.add(team) | ||
| db.commit() |
There was a problem hiding this comment.
another nit: sems like these same few lines appear in several places. could extract to _create_customer_if_necessary() or similar.
app/api/private_ai_keys.py
Outdated
| "created_at": info.get("created_at"), | ||
| "updated_at": info.get("updated_at"), | ||
| "metadata": info.get("metadata") | ||
| }) |
There was a problem hiding this comment.
assume there was a reason to not just do key_data.update(info)? If the missing values needed to be present but empty, could that be handled in the serializer?
app/core/worker.py
Outdated
| return | ||
|
|
||
| # Update the last payment date | ||
| team.last_payment = datetime.now(UTC) |
There was a problem hiding this comment.
this feels a bit weird. should it not be passed in / recorded from Stripe objects?
app/core/worker.py
Outdated
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| async def apply_product_for_team(db: Session, customer_id: str, product_id: str): |
There was a problem hiding this comment.
nit: feels like given how much this is doing it could maybe benefit from a more descriptive name
There was a problem hiding this comment.
the logic is the same for new subs and renewals, is that right?
There was a problem hiding this comment.
the logic is the same for new subs and renewals, is that right?
Yes. I suspect we may add complexity in the future, but for now a new subscription and a renewal will have the same effects.
app/api/billing.py
Outdated
| await apply_product_for_team(db, customer_id, product_id) | ||
| elif event_type in session_success_events: | ||
| product_id = await get_product_id_from_session(event_object.id) | ||
| await apply_product_for_team(db, customer_id, product_id) |
There was a problem hiding this comment.
is it clear in your mind which of these conditionals catch which steps in the workflow? normally when i have a subscriptoin stripe integration there's basically a single event i listen to for new subs and then a single one for renewals. so was surprised to see three different events that all land on the same workflow
There was a problem hiding this comment.
a single event i listen to for new subs and then a single one for renewals
I would prefer that! I may have read too much of the documentation... I think I have managed to whittle it down. Now it comes down to testing.
app/api/billing.py
Outdated
| frontend_url = os.getenv("FRONTEND_URL", "http://localhost:3000") | ||
|
|
||
| # Create checkout session using the service | ||
| checkout_url = await create_checkout_session(team, request_data.price_lookup_token, frontend_url) |
There was a problem hiding this comment.
where does the price_lookup_token on the request come from? i couldn't figure that out.
There was a problem hiding this comment.
This entire route is all wrong, so I'm removing it. We pivoted away from it once we discovered easier methods.
app/services/stripe.py
Outdated
| url=webhook_url, | ||
| enabled_events=[ | ||
| "checkout.session.completed", | ||
| "customer.subscription.deleted" |
There was a problem hiding this comment.
It seems like the webhook handling logic tries to handle many more updates than just these. What's the reason for the discrepancy?
There was a problem hiding this comment.
Me forgetting what I was doing because the change was too big. Good catch. I need to go through and remove things which aren't actually going to be used in the first iteration, as well as fix up some other aspects.
- Removed the `checkout` endpoint and replaced it with a `get_return_url` function to streamline return URL generation for team dashboards. - Consolidated Stripe event handling logic into the `handle_stripe_event_background` function, improving clarity and maintainability. - Updated the Stripe webhook event processing to include success and failure scenarios for subscriptions and checkout sessions. - Enhanced error handling for cases where teams are not registered with Stripe. - Removed outdated tests related to the checkout session and added new tests for the updated event handling logic.
- Refactored the `check_team_user_limit`, `check_key_limits`, and `check_vector_db_limits` functions to consolidate database queries, improving performance and reducing redundancy. - Updated the logic to retrieve current counts and maximum limits in a single query, enhancing efficiency. - Enhanced error handling to provide clearer messages when limits are exceeded. - Added new tests to validate the behavior of resource limits under various scenarios, ensuring comprehensive coverage for the updated functionality.
- Modified the `apply_product_for_team` function to accept a `start_date` parameter, allowing for accurate tracking of the last payment date. - Updated the `handle_stripe_event_background` function to pass the correct start date when applying products for new subscriptions and renewals. - Adjusted related tests to ensure compatibility with the new function signature and validate the correct application of products with the specified start date.
|
|
||
| # Combine database key info with LiteLLM info | ||
| key_data = private_ai_key.to_dict() | ||
| key_data.update(info) |
There was a problem hiding this comment.
I was looking at this commit and noticed that it previously changed "expires" to "expires_at". just flagging in case that's a bug/problem.
| db.query(DBUser.id).filter(DBUser.team_id == DBTeam.id) | ||
| ) | ||
| ) | ||
| ).first() |
There was a problem hiding this comment.
you really ran with the sql optimization! 😂 😵
app/services/stripe.py
Outdated
|
|
||
| success_events = invoice_success_events + subscription_success_events | ||
| failure_events = session_failure_events + subscription_failure_events + invoice_failure_events | ||
| known_events = success_events + failure_events |
There was a problem hiding this comment.
nit: it might be slightly more pythonic to name these in CONSTANT_CASE
- Renamed event handling constants in `stripe.py` for improved clarity and consistency. - Updated the `handle_stripe_event_background` function in `worker.py` to use the new constant names. - Changed the `expires_at` attribute to `expires` in the `PrivateAIKeyDetail` model for better semantic alignment.
This is a fairly large change to begin the integration with Stripe for billing purposes. Once this is in we still need to refine things, but that can be done in smaller chunks.
Key changes in this PR: