jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Feat/dify billing (#1679)

Browse source 
Co-authored-by: jyong <jyong@dify.ai>
Co-authored-by: takatost <takatost@users.noreply.github.com>
This commit is contained in:
Garfield Dai 2023-12-03 20:59:29 +08:00 • committed by GitHub
commit 053102f433
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 182 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

55
api/services/billing_service.py Normal file
View file

@ -0,0 +1,55 @@
import os
import requests
from services.dataset_service import DatasetService
class BillingService:
base_url = os.environ.get('BILLING_API_URL', 'BILLING_API_URL')
secret_key = os.environ.get('BILLING_API_SECRET_KEY', 'BILLING_API_SECRET_KEY')
@classmethod
def get_info(cls, tenant_id: str):
params = {'tenant_id': tenant_id}
billing_info = cls._send_request('GET', '/info', params=params)
vector_size = DatasetService.get_tenant_datasets_usage(tenant_id) / 1024
billing_info['vector_space']['size'] = int(vector_size)
return billing_info
@classmethod
def get_subscription(cls, plan: str, interval: str, prefilled_email: str = '', user_name: str = '', tenant_id: str = ''):
params = {
'plan': plan,
'interval': interval,
'prefilled_email': prefilled_email,
'user_name': user_name,
'tenant_id': tenant_id
}
return cls._send_request('GET', '/subscription', params=params)
@classmethod
def get_invoices(cls, prefilled_email: str = ''):
params = {'prefilled_email': prefilled_email}
return cls._send_request('GET', '/invoices', params=params)
@classmethod
def _send_request(cls, method, endpoint, json=None, params=None):
headers = {
"Content-Type": "application/json",
"Billing-Api-Secret-Key": cls.secret_key
}
url = f"{cls.base_url}{endpoint}"
response = requests.request(method, url, json=json, params=params, headers=headers)
return response.json()
@classmethod
def process_event(cls, event: dict):
json = {
"content": event,
}
return cls._send_request('POST', '/webhook/stripe', json=json)

33
api/services/dataset_service.py
View file

@ -227,6 +227,36 @@ class DatasetService:
return AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id) \
.order_by(db.desc(AppDatasetJoin.created_at)).all()
@staticmethod
def get_tenant_datasets_usage(tenant_id):
# get the high_quality datasets
dataset_ids = db.session.query(Dataset.id).filter(Dataset.indexing_technique == 'high_quality',
Dataset.tenant_id == tenant_id).all()
if not dataset_ids:
return 0
dataset_ids = [result[0] for result in dataset_ids]
document_ids = db.session.query(Document.id).filter(Document.dataset_id.in_(dataset_ids),
Document.tenant_id == tenant_id,
Document.completed_at.isnot(None),
Document.enabled == True,
Document.archived == False
).all()
if not document_ids:
return 0
document_ids = [result[0] for result in document_ids]
document_segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id.in_(document_ids),
DocumentSegment.tenant_id == tenant_id,
DocumentSegment.completed_at.isnot(None),
DocumentSegment.enabled == True,
).all()
if not document_segments:
return 0
total_words_size = sum(document_segment.word_count * 3 for document_segment in document_segments)
total_vector_size = 1536 * 4 * len(document_segments)
return total_words_size + total_vector_size
class DocumentService:
DEFAULT_RULES = {
@ -488,7 +518,8 @@ class DocumentService:
'score_threshold_enabled': False
}
dataset.retrieval_model = document_data.get('retrieval_model') if document_data.get('retrieval_model') else default_retrieval_model
dataset.retrieval_model = document_data.get('retrieval_model') if document_data.get(
'retrieval_model') else default_retrieval_model
documents = []
batch = time.strftime('%Y%m%d%H%M%S') + str(random.randint(100000, 999999))