|
|
import time, gc
|
|
|
from config import *
|
|
|
from agent.question.question_utils import preprocessing_question, process_question, process_dict, add_question_mark, question_rew
|
|
|
from agent.llm_client import create_chat_completion, create_chat_completion_change_zhushi
|
|
|
from agent.schema.schema_utils import get_table_schema, find_table_with_emb, get_table_schema_with_emb, deep_find_tables
|
|
|
from agent.sql.exec import exec_sql_s
|
|
|
from agent.sql.process_sql import to_select, extract_table_names, all_tables_in_prompt
|
|
|
from agent.sql.knowledge import way_string_2
|
|
|
from agent.optimization import optimize_answer, find_kl_answer
|
|
|
from agent.utils.async_util import async_llm_chain_call
|
|
|
from agent.utils.io_utils import load_external_answers
|
|
|
from agent.utils.json_utils import find_json
|
|
|
|
|
|
def run_conversation_xietong(question, org_question):
|
|
|
global QUESTION_TYPE
|
|
|
global QUESTION_TYPE_LIST
|
|
|
global LL
|
|
|
global table_maps
|
|
|
global table_maps_LL
|
|
|
global all_tables_name_list
|
|
|
str_process_question = process_question(question)
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"-------------已查询获得事实: {str_process_question}")
|
|
|
|
|
|
if "ConstantDB.HK_SecuMain" in str(str_process_question) and "港股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("港股")
|
|
|
if "ConstantDB.US_SecuMain" in str(str_process_question) and "美股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("美股")
|
|
|
if "ConstantDB.SecuMain" in str(str_process_question) and "A股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("A股")
|
|
|
|
|
|
content_p_1 = content_ALL
|
|
|
if len(QUESTION_TYPE_LIST) == 0:
|
|
|
content_p_1 = content_ALL
|
|
|
elif len(QUESTION_TYPE_LIST) >= 2:
|
|
|
content_p_1 = content_ALL
|
|
|
elif "港股" in QUESTION_TYPE_LIST:
|
|
|
content_p_1 = content_HK
|
|
|
elif "美股" in QUESTION_TYPE_LIST:
|
|
|
content_p_1 = content_US
|
|
|
elif "A股" in QUESTION_TYPE_LIST:
|
|
|
content_p_1 = content_CN
|
|
|
else:
|
|
|
content_p_1 = content_ALL
|
|
|
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"---------QUESTION_TYPE_LIST:{QUESTION_TYPE_LIST}")
|
|
|
content_p = content_p_1.replace('<<question>>', str(question)).replace('<<fact_1>>',
|
|
|
str(str_process_question))
|
|
|
main_sql_prompts,highest_prompt_similar = way_string_2(org_question)
|
|
|
content_p = content_p + main_sql_prompts
|
|
|
# 执行函数部分
|
|
|
messages = []
|
|
|
messages.append({"role": "user", "content": "您好阿"})
|
|
|
messages.append({"role": "user", "content": content_p})
|
|
|
response = create_chat_completion(messages)
|
|
|
|
|
|
if "全部完成" in response.choices[0].message.content:
|
|
|
if "XX" in response.choices[0].message.content:
|
|
|
response.choices[0].message.content += "ConstantDB.SecuMain,ConstantDB.HK_SecuMain,ConstantDB.US_SecuMain"
|
|
|
else:
|
|
|
if DEBUG_VER == 3:
|
|
|
print("全部完成：response：",response.choices[0].message.content)
|
|
|
if "需要查询的数据库表" in response.choices[0].message.content or "需要查询数据库表" in response.choices[0].message.content:
|
|
|
pass
|
|
|
else:
|
|
|
return response.choices[0].message.content
|
|
|
|
|
|
messages.append({"role": "assistant", "content": response.choices[0].message.content})
|
|
|
table_maps = get_table_schema(question)
|
|
|
LL = [i for i in table_maps if i.get('数据表名') in response.choices[0].message.content]
|
|
|
tables_name_list = [i.get('数据表名') for i in table_maps if i.get('数据表名') in response.choices[0].message.content]
|
|
|
try:
|
|
|
sql_prompt_db_tables = extract_table_names(main_sql_prompts)
|
|
|
except Exception as e:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"------->e:{e}")
|
|
|
if sql_prompt_db_tables:
|
|
|
all_tables_name_list += sql_prompt_db_tables
|
|
|
all_tables_name_list += tables_name_list
|
|
|
all_tables_name_list = list(set(all_tables_name_list))
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"-----------------all_tables_name_list:{all_tables_name_list}")
|
|
|
all_tables_name_list += find_table_with_emb(all_tables_name_list, question)
|
|
|
all_tables_name_list = list(set(all_tables_name_list))
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"-----------------find_table_with_emb all_tables_name_list:{all_tables_name_list}")
|
|
|
if all_tables_name_list == []:
|
|
|
highest_similarity_score_list = [0]
|
|
|
else:
|
|
|
table_maps_LL, highest_similarity_score_list = get_table_schema_with_emb(all_tables_name_list, question)
|
|
|
# print(f"--------------table_maps_LL:{table_maps_LL}")
|
|
|
if all_tables_in_prompt(tables_name_list, main_sql_prompts) and highest_prompt_similar > START_DEEP_THRESHOLD:
|
|
|
content_p_2 = """获取的表结构如下<list>,表结构中列名可以引用使用,表结构中数据示例只是参考不能引用。
|
|
|
我们现在开始查询当前问题，可参考示例sql写出查询sql语句，如果示例sql是一句的（一步），你也尽可能用一句sql得到结果。
|
|
|
如果需要多步，我把查询结果告诉你，你再告诉我下一步，注意如果我返回的结果为空或者错误影响下一步调用，请重新告诉我sql语句。
|
|
|
等你全部回答完成，不需要进行下一步调用时，记得提示我：<全部完成，答案如下>,将答案总结以json格式给我，只需要总结当前问题。
|
|
|
查询技巧: 涉及计算的使用sql进行加减乘除数学计算。查询尽量使用InnerCode。查年用year()函数、月用month()、某日则用date()，如：date(TradingDay) = 'YYYY-MM-DD'。sql查询语句不需要注释，不然会报错。尽量利用表格中已有的字段。\n"""
|
|
|
|
|
|
else:
|
|
|
if DEEP_COLUMNS_RESEARCH:
|
|
|
if max(highest_similarity_score_list) < START_DEEP_THRESHOLD: # 没有大于0.5的说明这个表可能不对，新增表
|
|
|
print(f"-----------Start Deeping table search")
|
|
|
deep_table_maps_LL = deep_find_tables(all_tables_name_list, org_question) # 直接用原始的扫
|
|
|
print(f"--------------len(deep_table_maps_LL):{len(deep_table_maps_LL)}")
|
|
|
deep_tables_name_list = [i.get('数据表名') for i in deep_table_maps_LL]
|
|
|
all_tables_name_list += deep_tables_name_list
|
|
|
all_tables_name_list = list(set(all_tables_name_list))
|
|
|
|
|
|
if max(highest_similarity_score_list) < START_DEEP_THRESHOLD:
|
|
|
table_maps_LL, highest_similarity_score_list = get_table_schema_with_emb(all_tables_name_list, question)
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"--------------End Deeping table_maps_LL:{table_maps_LL}")
|
|
|
|
|
|
content_p_2 = """获取的表结构如下<list>,表结构中列名可以引用使用,表结构中数据示例只是参考不能引用。
|
|
|
我们现在开始查询当前问题，请你分步写出查询sql语句，我把查询结果告诉你，你再告诉我下一步，
|
|
|
注意如果我返回的结果为空或者错误影响下一步调用，请重新告诉我sql语句。
|
|
|
等你全部回答完成，不需要进行下一步调用时，记得提示我：<全部完成，答案如下>,将答案总结以json格式给我，只需要总结当前问题。
|
|
|
查询技巧: 涉及计算的使用sql进行加减乘除数学计算。查询尽量使用InnerCode。查年用year()、查月用month()、某日用date()，如：date(TradingDay) = 'YYYY-MM-DD'。sql查询语句不需要注释，不然会报错。尽量利用表格中已有的字段。\n"""
|
|
|
if len(str_process_question[1]) == 0:
|
|
|
if DEBUG_VER == 3:
|
|
|
print("--------------查询获得事实为空，跳过处理")
|
|
|
else:
|
|
|
if FRASH_DATA_DESCRIBE:
|
|
|
table_maps_LL = create_chat_completion_change_zhushi(str_process_question, table_maps_LL)
|
|
|
|
|
|
content_p_2 = main_sql_prompts + content_p_2.replace('<list>', str(table_maps_LL))
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"-------------content_p_2: {content_p_2}")
|
|
|
messages.append({"role": "user", "content": content_p_2}) ###开始对话
|
|
|
last_answer = run_conversation_until_complete(messages, max_rounds=MAX_ROUNDS)
|
|
|
return last_answer
|
|
|
|
|
|
def run_conversation_until_complete(messages, max_rounds=9):
|
|
|
"""
|
|
|
Test function to run a conversation loop until the assistant indicates completion.
|
|
|
"""
|
|
|
last_response = None # 用于存储最后一次对话的响应
|
|
|
round_count = 0 # 对话轮数计数器
|
|
|
# response = create_chat_completion(messages)
|
|
|
if ASYNC_LLM_TIME >= 2:
|
|
|
response = async_llm_chain_call(messages, MODEL)
|
|
|
else:
|
|
|
response = create_chat_completion(messages)
|
|
|
|
|
|
while True:
|
|
|
if round_count >= max_rounds:
|
|
|
break # 如果对话轮数超过最大值，则退出循环
|
|
|
|
|
|
question = response.choices[0].message.content
|
|
|
select_result = to_select(question)
|
|
|
messages.append({"role": "assistant", "content": question})
|
|
|
messages.append({"role": "user", "content": str(select_result)})
|
|
|
|
|
|
# response = create_chat_completion(messages)
|
|
|
if ASYNC_LLM_TIME >= 2:
|
|
|
response = async_llm_chain_call(messages, MODEL)
|
|
|
else:
|
|
|
response = create_chat_completion(messages)
|
|
|
|
|
|
last_response = response.choices[0].message.content # 存储最后一次响应
|
|
|
if "全部完成" in response.choices[0].message.content:
|
|
|
break # 如果检测到“回答完成”，则停止循环
|
|
|
round_count += 1 # 增加对话轮数计数
|
|
|
return last_response # 返回最后一次对话的内容
|
|
|
|
|
|
def run_conversation(question, org_question):
|
|
|
"""
|
|
|
Run a conversation flow given a question by:
|
|
|
1. Using run_conversation_xietong(question) to get an answer.
|
|
|
2. Attempting to find and parse JSON from the answer using find_json.
|
|
|
3. Converting the parsed result (or original text if parsing fails) into a descriptive sentence using process_dict.
|
|
|
|
|
|
Parameters:
|
|
|
question (str): The question to ask.
|
|
|
|
|
|
Returns:
|
|
|
str: The final processed answer as a descriptive string.
|
|
|
"""
|
|
|
org_answer = run_conversation_xietong(question, org_question)
|
|
|
parsed_data = find_json(org_answer)
|
|
|
last_answer = process_dict(parsed_data)
|
|
|
return last_answer,org_answer
|
|
|
|
|
|
def get_answer(question, org_question):
|
|
|
"""
|
|
|
Attempt to answer the given question by interacting with the
|
|
|
conversation model. If an error occurs, return a default error message.
|
|
|
|
|
|
Parameters:
|
|
|
question (str): The question that needs an answer.
|
|
|
|
|
|
Returns:
|
|
|
str: The answer string or an error message if an exception occurs.
|
|
|
"""
|
|
|
try:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"Attempting to answer the question: {question}")
|
|
|
last_answer,org_answer = run_conversation(question, org_question)
|
|
|
return last_answer,org_answer
|
|
|
except Exception as e:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"Error occurred while executing get_answer: {e}")
|
|
|
return "An error occurred while retrieving the answer."
|
|
|
|
|
|
def main_answer(q_json_list, out_path, idx_ranges):
|
|
|
"""
|
|
|
处理整个 JSON 数据列表：
|
|
|
- 如果当前索引在 idx_ranges 中任一范围内，则按照原有逻辑提取问题、重写问题、调用 get_answer 得到答案，
|
|
|
并更新每个问题对应的 answer 字段；
|
|
|
- 否则将每个 team 成员的 answer 字段更新为空。
|
|
|
参数：
|
|
|
q_json_list: JSON 数据列表
|
|
|
out_path: 输出文件路径
|
|
|
idx_ranges: 包含多个 (start, end) 范围的列表，例如 [(21, 22), (74, 75), (43, 44)]
|
|
|
"""
|
|
|
global QUESTION_TYPE, QUESTION_TYPE_LIST, all_tables_name_list, prev_tables_name_list, org_question
|
|
|
|
|
|
# 如果 idx_ranges 为空，则默认处理全部数据
|
|
|
if not idx_ranges:
|
|
|
idx_ranges = [(0, len(q_json_list))]
|
|
|
|
|
|
if REFRASH and Knowledge_file:
|
|
|
external_answers = load_external_answers(Knowledge_file)
|
|
|
total = len(q_json_list)
|
|
|
# 以流式写入的方式生成 JSON 数组
|
|
|
with open(out_path, 'w', encoding='utf-8') as out_f:
|
|
|
out_f.write('[\n')
|
|
|
for i in range(total):
|
|
|
try:
|
|
|
print("--> cleaned mem:", gc.collect())
|
|
|
print("--> id:", i)
|
|
|
item = q_json_list[i]
|
|
|
start_time = time.time()
|
|
|
|
|
|
# 判断当前索引是否在任一给定范围内
|
|
|
process_this_item = any(start <= i < end for start, end in idx_ranges)
|
|
|
|
|
|
if process_this_item:
|
|
|
# 处理需要调用 get_answer 生成答案的数据
|
|
|
questions_list = [(member["id"], member["question"]) for member in item["team"]]
|
|
|
answers_dict = {}
|
|
|
history_dict = {}
|
|
|
all_previous = ''
|
|
|
# 初始化相关全局变量
|
|
|
QUESTION_TYPE = ""
|
|
|
QUESTION_TYPE_LIST = []
|
|
|
question_num = 0
|
|
|
all_tables_name_list = []
|
|
|
prev_tables_name_list = []
|
|
|
|
|
|
if REFRASH:
|
|
|
get_ext_answer_num = 1
|
|
|
# 遍历当前 item 内所有问题
|
|
|
for question_id, question_text in questions_list:
|
|
|
try:
|
|
|
answer = ""
|
|
|
# 更新历史表变量（超过 MAX_TABLES_LIST 后清空）
|
|
|
all_tables_name_list = prev_tables_name_list
|
|
|
if len(all_tables_name_list) > MAX_TABLES_LIST:
|
|
|
if DEBUG_VER == 3:
|
|
|
print(f"----------------历史表过多，进行清空")
|
|
|
all_tables_name_list = []
|
|
|
|
|
|
question_num += 1
|
|
|
org_question = question_text
|
|
|
# 根据关键词更新 QUESTION_TYPE_LIST
|
|
|
if "港股" in question_text and "港股份" not in question_text and "港股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("港股")
|
|
|
if "美股" in question_text and "美股份" not in question_text and "美股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("美股")
|
|
|
if "A股" in question_text and "A股" not in QUESTION_TYPE_LIST:
|
|
|
QUESTION_TYPE_LIST.append("A股")
|
|
|
|
|
|
# 问题预处理
|
|
|
# if DEBUG_VER == 3:
|
|
|
print(f'---------->Org question_text:{question_text}')
|
|
|
question_text = preprocessing_question(question_text)
|
|
|
question_text = add_question_mark(question_text)
|
|
|
|
|
|
# 如果没有历史问答，则直接使用原问题；否则重写问题
|
|
|
if all_previous == '':
|
|
|
rewritten_question = question_text
|
|
|
else:
|
|
|
rewritten_question = question_rew(all_previous, question_text)
|
|
|
|
|
|
# 调用 get_answer 获取答案，并进行优化
|
|
|
if REFRASH:
|
|
|
matched_answer = find_kl_answer(question_num, history_dict, org_question, external_answers)
|
|
|
if matched_answer:
|
|
|
if get_ext_answer_num == question_num:
|
|
|
answer = matched_answer
|
|
|
history_dict[org_question] = answer
|
|
|
# 累加问答历史
|
|
|
all_previous += f"第{question_num}个问题：" + question_text + f" 第{question_num}个的回答：" + answer + " \n"
|
|
|
get_ext_answer_num += 1
|
|
|
if not answer or answer == "":
|
|
|
answer,org_answer = get_answer(rewritten_question, org_question)
|
|
|
# 累加问答历史
|
|
|
all_previous += f"第{question_num}个问题：" + question_text + f" 第{question_num}个的回答：" + answer + " \n"
|
|
|
answer = optimize_answer(rewritten_question, answer, org_answer)
|
|
|
# if DEBUG_VER == 3:
|
|
|
print(f"---------->answer:{answer}")
|
|
|
answer = answer.replace(";", "")
|
|
|
answers_dict[question_id] = answer
|
|
|
del answer
|
|
|
|
|
|
except Exception as e_inner:
|
|
|
print(f"Error processing question id {question_id}: {e_inner}")
|
|
|
answers_dict[question_id] = "异常"
|
|
|
|
|
|
# 更新 item 中各 team 成员的 answer 字段
|
|
|
for member in item["team"]:
|
|
|
member["answer"] = answers_dict.get(member["id"], "无答案")
|
|
|
else:
|
|
|
# 对于不在处理范围内的数据，直接置空 answer 字段
|
|
|
for member in item["team"]:
|
|
|
member["answer"] = ""
|
|
|
|
|
|
updated_data = {"tid": item["tid"], "team": item["team"]}
|
|
|
|
|
|
elapsed_time = time.time() - start_time
|
|
|
print(f"Completed processing JSON index {i} in {elapsed_time:.2f} seconds")
|
|
|
|
|
|
# 写入当前处理完的数据项。若不是最后一个，则尾部添加逗号。
|
|
|
json_str = json.dumps(updated_data, ensure_ascii=False, indent=4)
|
|
|
if i < total - 1:
|
|
|
out_f.write(json_str + ',\n')
|
|
|
else:
|
|
|
out_f.write(json_str + '\n')
|
|
|
out_f.flush() # 每处理完一道题就 flush 到磁盘，降低内存占用
|
|
|
|
|
|
except Exception as e_outer:
|
|
|
print(f"Error processing JSON item index {i}: {e_outer}")
|
|
|
out_f.write(']\n')
|
