diff --git a/dtable_events/app/config.py b/dtable_events/app/config.py index 0f36c87a..d0344b16 100644 --- a/dtable_events/app/config.py +++ b/dtable_events/app/config.py @@ -99,8 +99,12 @@ def get_llm_prices(models): # AI models and prices LLM_MODELS = configs.get('LLM_MODELS', []) -AI_PRICES = get_llm_prices(LLM_MODELS) -BAIDU_OCR_TOKENS = configs.get('BAIDU_OCR_TOKENS', default={}) + +## OCR +OCR_SERVICE_CONFIG = configs.get('OCR', {}) + +## AI price +AI_PRICES = get_llm_prices(LLM_MODELS + [OCR_SERVICE_CONFIG]) # IO server ARCHIVE_VIEW_EXPORT_ROW_LIMIT = configs.get('ARCHIVE_VIEW_EXPORT_ROW_LIMIT', default=250000) diff --git a/dtable_events/tasks/ai_stats_worker.py b/dtable_events/tasks/ai_stats_worker.py index b6561ba7..e77d04e4 100644 --- a/dtable_events/tasks/ai_stats_worker.py +++ b/dtable_events/tasks/ai_stats_worker.py @@ -10,10 +10,10 @@ from dateutil import relativedelta from sqlalchemy import text -from dtable_events.app.config import AI_PRICES, BAIDU_OCR_TOKENS, AI_STATS_ENABLED +from dtable_events.app.config import AI_PRICES, AI_STATS_ENABLED from dtable_events.app.event_redis import RedisClient, redis_cache from dtable_events.db import init_db_session_class -from dtable_events.utils import get_opt_from_conf_or_env, parse_bool, uuid_str_to_36_chars, uuid_str_to_32_chars +from dtable_events.utils import uuid_str_to_36_chars, uuid_str_to_32_chars logger = logging.getLogger(__name__) @@ -46,10 +46,17 @@ def save_to_memory(self, usage_info, session): model = usage_info['model'] usage = usage_info.get('usage') or {} + org_id = usage_info.get('org_id') if model not in AI_PRICES: logger.warning('model %s price not defined', model) return + + try: + org_id = int(org_id) + except (TypeError, ValueError): + org_id = -1 + if 'prompt_tokens' in usage: usage['input_tokens'] = usage['prompt_tokens'] @@ -61,9 +68,6 @@ def save_to_memory(self, usage_info, session): if not isinstance(usage.get('output_tokens'), int): usage['output_tokens'] = 0 - if model in BAIDU_OCR_TOKENS: - usage['output_tokens'] = BAIDU_OCR_TOKENS[model] - if usage_info.get('assistant_uuid'): # for assistant call, set stat object to the assistant owner, i.e., table admin, group admin ... @@ -73,9 +77,9 @@ def save_to_memory(self, usage_info, session): if not owner_info: logger.warning('assistant %s has no owner', assistant_uuid) return - if owner_info['org_id'] != -1: - self.org_stats[owner_info['org_id']][model]['input_tokens'] += usage.get('input_tokens') or 0 - self.org_stats[owner_info['org_id']][model]['output_tokens'] += usage.get('output_tokens') or 0 + if org_id != -1: + self.org_stats[org_id][model]['input_tokens'] += usage.get('input_tokens') or 0 + self.org_stats[org_id][model]['output_tokens'] += usage.get('output_tokens') or 0 else: self.owner_stats[owner_info['owner_id']][model]['input_tokens'] += usage.get('input_tokens') or 0 self.owner_stats[owner_info['owner_id']][model]['output_tokens'] += usage.get('output_tokens') or 0