forked from open-webui/open-webui
-
Notifications
You must be signed in to change notification settings - Fork 0
/
credits.py
146 lines (119 loc) · 4.5 KB
/
credits.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import os
from fastapi import Depends, HTTPException, status, Request
from typing import Optional
import requests
import json
from fastapi import APIRouter
from pydantic import BaseModel
from apps.web.models.users import Users
from utils.utils import (
bearer_scheme,
)
from constants import ERROR_MESSAGES
import stripe
# This is your test secret API key.
stripe.api_key = os.environ.get('STRIPE_API_KEY')
price_id = os.environ.get('STRIPE_PRICE_ID')
webhook_secret = os.environ.get('STRIPE_WEBHOOK_SECRET')
app_url = os.environ.get('WEBUI_API_BASE_URL', 'https://chat.meteron.ai')
class UserCredits(BaseModel):
email: str
balance: float
router = APIRouter()
# Get user's credits from Meteron's API
@router.get("/", response_model=Optional[UserCredits])
async def get_user_credits(cred=Depends(bearer_scheme)):
token = cred.credentials
user = Users.get_user_by_token(token)
if user:
url = "https://app.meteron.ai/api/credits?user=" + user.email
headers = {
"Content-Type": "application/json",
"Authorization": "Bearer " + os.environ["METERON_API_KEY"]
}
response = requests.get(url, headers=headers)
return {
"email": user.email,
"balance": response.json()['balance']
}
else:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
)
class UserCreditsTopUp(BaseModel):
url: str
# Checkout session is created in Stripe and the URL is returned to the client. UI
# then redirects to the given URL to being the payment
@router.post("/checkout-session", response_model=Optional[UserCreditsTopUp])
async def start_user_checkout(cred=Depends(bearer_scheme)):
token = cred.credentials
user = Users.get_user_by_token(token)
if user:
try:
checkout_session = stripe.checkout.Session.create(
customer_email=user.email,
line_items=[
{
'price': price_id,
'quantity': 5,
},
],
mode='payment',
success_url=app_url + '/',
cancel_url=app_url + '/',
)
except Exception as e:
print(e)
return str(e)
return {
"url": checkout_session.url
}
else:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
)
# One-off payments in Stripe produce 'checkout.session.completed' events. We need to
# read the event, extract the email and call Meteron's API to top up the user's credits.
@router.post("/webhook")
async def stripe_webhook(request: Request):
payload = await request.body()
sig_header = request.headers.get('Stripe-Signature')
try:
event = stripe.Webhook.construct_event(
payload, sig_header, webhook_secret
)
# Handle the event
if event['type'] == 'checkout.session.completed':
invoice = event['data']['object']
customer_email = invoice['customer_email']
amount = invoice["amount_total"]/100 # Stripe amounts are in cents
checkout_session_id = invoice["id"]
# Print the customer's email address
print(f"Customer email: {customer_email}, amount: {amount}, checkout session id: {checkout_session_id}")
# Call Meteron's API to top up the user's credits
headers = {
"Content-Type": "application/json",
"Authorization": "Bearer " + os.environ["METERON_API_KEY"]
}
try:
requests.put("https://app.meteron.ai/api/credits", headers=headers, data=json.dumps({
"user": customer_email,
"amount": amount
}))
except Exception as e:
print(e)
# Failed to topup
return {"status": "error", "message": str(e)}
return
return {"status": "success"}
except ValueError as e:
# Invalid payload
return {"status": "error", "message": str(e)}
except stripe.error.SignatureVerificationError as e:
# Invalid signature
return {"status": "error", "message": str(e)}
except Exception as e:
# Other error
return {"status": "error", "message": str(e)}