fix notion estimate (#2090)

Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Jyong 2024-01-19 13:27:12 +08:00 committed by GitHub
parent a3c7c07ecc
commit 869690c485
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -274,6 +274,8 @@ class IndexingRunner:
tokens = 0 tokens = 0
preview_texts = [] preview_texts = []
total_segments = 0 total_segments = 0
total_price = 0
currency = 'USD'
for file_detail in file_details: for file_detail in file_details:
processing_rule = DatasetProcessRule( processing_rule = DatasetProcessRule(
@ -344,11 +346,13 @@ class IndexingRunner:
price_type=PriceType.INPUT, price_type=PriceType.INPUT,
tokens=tokens tokens=tokens
) )
total_price = '{:f}'.format(embedding_price_info.total_amount)
currency = embedding_price_info.currency
return { return {
"total_segments": total_segments, "total_segments": total_segments,
"tokens": tokens, "tokens": tokens,
"total_price": '{:f}'.format(embedding_price_info.total_amount) if embedding_model_instance else 0, "total_price": total_price,
"currency": embedding_price_info.currency if embedding_model_instance else 'USD', "currency": currency,
"preview": preview_texts "preview": preview_texts
} }
@ -388,6 +392,8 @@ class IndexingRunner:
tokens = 0 tokens = 0
preview_texts = [] preview_texts = []
total_segments = 0 total_segments = 0
total_price = 0
currency = 'USD'
for notion_info in notion_info_list: for notion_info in notion_info_list:
workspace_id = notion_info['workspace_id'] workspace_id = notion_info['workspace_id']
data_source_binding = DataSourceBinding.query.filter( data_source_binding = DataSourceBinding.query.filter(
@ -470,7 +476,7 @@ class IndexingRunner:
"qa_preview": document_qa_list, "qa_preview": document_qa_list,
"preview": preview_texts "preview": preview_texts
} }
if embedding_model_instance:
embedding_model_type_instance = embedding_model_instance.model_type_instance embedding_model_type_instance = embedding_model_instance.model_type_instance
embedding_model_type_instance = cast(TextEmbeddingModel, embedding_model_type_instance) embedding_model_type_instance = cast(TextEmbeddingModel, embedding_model_type_instance)
embedding_price_info = embedding_model_type_instance.get_price( embedding_price_info = embedding_model_type_instance.get_price(
@ -479,11 +485,13 @@ class IndexingRunner:
price_type=PriceType.INPUT, price_type=PriceType.INPUT,
tokens=tokens tokens=tokens
) )
total_price = '{:f}'.format(embedding_price_info.total_amount)
currency = embedding_price_info.currency
return { return {
"total_segments": total_segments, "total_segments": total_segments,
"tokens": tokens, "tokens": tokens,
"total_price": '{:f}'.format(embedding_price_info.total_amount) if embedding_model_instance else 0, "total_price": total_price,
"currency": embedding_price_info.currency if embedding_model_instance else 'USD', "currency": currency,
"preview": preview_texts "preview": preview_texts
} }