Spaces:
Running
Running
| import psycopg2 | |
| import datetime | |
| from bin_public.config.presets import * | |
| from dateutil import tz | |
| import os | |
| def current_time(type): | |
| if type == 'ymd': | |
| return datetime.datetime.now().strftime("%Y-%m-%d") | |
| if type == 'ymdhms': | |
| return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| # hologres 基础函数：查询 | |
| def holo_query_func(run_sql, is_query=0): | |
| conn = psycopg2.connect(host=os.environ['HOST'], | |
| port=os.environ['PORT'], | |
| dbname=os.environ['DBNAME'], | |
| user=os.environ['AK'], | |
| password=os.environ['SK']) | |
| cur = conn.cursor() | |
| cur.execute(run_sql) | |
| if is_query: | |
| data = cur.fetchall() | |
| cur.close() | |
| conn.close() | |
| if is_query: | |
| return data | |
| def holo_query_account_mapping(invite_code): | |
| run_sql = f""" | |
| select end_date, status, mapping_ak | |
| from s_account_invite_code | |
| where invite_code = '{invite_code}' | |
| order by gmt_modify desc | |
| limit 1 | |
| """ | |
| data = holo_query_func(run_sql, is_query=1) | |
| # 数据库中查不到，则返回no_invite_code_msg | |
| if len(data) == 0: | |
| status_text = standard_error_msg + no_invite_code_msg | |
| return status_text, None | |
| # 数据库中查到，判断是否可用 | |
| if len(data) == 1: | |
| end_date = data[0][0] | |
| status = data[0][1] | |
| mapping_ak = data[0][2] | |
| if end_date < datetime.datetime.now().strftime("%Y%m%d") or status != '1': | |
| status_text = standard_error_msg + no_useful_invite_code_msg | |
| return status_text, None | |
| return 'Success status: ready', mapping_ak | |
| def key_preprocessing(keyTxt): | |
| invite_code = keyTxt | |
| # 这里先用这个逻辑，到时候等实际的邀请码来了就改一下这个函数就行 | |
| if keyTxt.startswith("dteam_"): | |
| status_display, keyTxt = holo_query_account_mapping(keyTxt) | |
| yield status_display, keyTxt, invite_code | |
| return | |
| else: | |
| if len(keyTxt) != 51: | |
| status_display = standard_error_msg + no_apikey_msg | |
| yield status_display, keyTxt, invite_code | |
| return | |
| yield 'Success status: ready', keyTxt, invite_code | |
| return | |
| def holo_query_insert_chat_message(invite_code, prompt, response, all_token_cnt, history): | |
| run_sql = f""" | |
| insert into s_account_chat_message( | |
| gmt_create | |
| ,invite_code | |
| ,prompt | |
| ,response | |
| ,all_token_cnt | |
| ,history | |
| ,chat_seq | |
| ,log_timestamp | |
| ) | |
| select | |
| '{datetime.datetime.now().replace(tzinfo=tz.gettz('Asina/Shanghai')).strftime("%Y-%m-%d %H:%M:%S")}' as gmt_create | |
| ,'{str(invite_code).replace("'", '"')}' as invite_code | |
| ,'{str(prompt).replace("'", '"')}' as prompt | |
| ,'{str(response).replace("'", '"')}' as response | |
| ,'{str(all_token_cnt).replace("'", '"')}' as all_token_cnt | |
| ,'{str(history).replace("'", '"')}' as history | |
| ,'{len(history)}' as chat_seq | |
| ,localtimestamp as log_timestamp | |
| """ | |
| holo_query_func(run_sql, is_query=0) |