113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
from datetime import datetime
|
|
import json
|
|
import logging
|
|
import re
|
|
from pydantic import BaseModel, Field
|
|
|
|
from configs.model_config import LLM_MODELS
|
|
from server.agent.tools.search_tool import search_tool
|
|
from server.chat import utils
|
|
from server.chat.policy_fun_iast import get_llm_model_response
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def math_count(query: str):
|
|
try:
|
|
matches = re.findall(r'\{.*?\}', query)
|
|
if len(matches)>=2:
|
|
query = matches[0]
|
|
parsed_query = json.loads(query)["query"]
|
|
# 继续使用解析后的查询进行后续操作
|
|
time_based_uuid = json.loads(matches[1])["uuid"]
|
|
# temp = utils.get_shared_variable(time_based_uuid)
|
|
# temp["END"] ="ok"
|
|
# utils.set_shared_variable(time_based_uuid,temp)
|
|
|
|
tip = {}
|
|
# tip["END"]="ok"
|
|
tip["source_docs"]=[]
|
|
tip["num"]=0
|
|
tip["title"]=[]
|
|
|
|
utils.set_shared_variable(time_based_uuid+"q",tip)
|
|
first_json = {
|
|
"query": parsed_query,
|
|
"knowledge_name": [],
|
|
"keywords": []
|
|
}
|
|
second_json = {
|
|
"uuid": time_based_uuid+"q"
|
|
}
|
|
math_doc = search_tool(json.dumps(first_json) + json.dumps(second_json))
|
|
utils.remove_shared_variable(time_based_uuid+"q")
|
|
res = get_llm_model_response(
|
|
strategy_name="default_math",
|
|
llm_model_name=LLM_MODELS[3],
|
|
template_prompt_name="default_math",
|
|
prompt_param_dict={"input": parsed_query, "math_doc": f"{math_doc}", "time": datetime.now().strftime("%Y%m%d")},
|
|
temperature=0.01,
|
|
max_tokens=512
|
|
)
|
|
return f"{res}"
|
|
except Exception as e:
|
|
logging.error(f"Error occurred while processing math query: {e}")
|
|
return "<system>不要再调用该工具了</system>"
|
|
|
|
except Exception as e:
|
|
return "<system>不要再调用该工具了</system>"
|
|
|
|
def code_count(query: str):
|
|
try:
|
|
matches = re.findall(r'\{.*?\}', query)
|
|
if len(matches)>=2:
|
|
query = matches[0]
|
|
parsed_query = json.loads(query)["query"]
|
|
# 继续使用解析后的查询进行后续操作
|
|
time_based_uuid = json.loads(matches[1])["uuid"]
|
|
# temp = utils.get_shared_variable(time_based_uuid)
|
|
# temp["END"] ="ok"
|
|
# utils.set_shared_variable(time_based_uuid,temp)
|
|
|
|
tip = {}
|
|
# tip["END"]="ok"
|
|
tip["source_docs"]=[]
|
|
tip["num"]=0
|
|
tip["title"]=[]
|
|
|
|
utils.set_shared_variable(time_based_uuid+"q",tip)
|
|
first_json = {
|
|
"query": parsed_query,
|
|
"knowledge_name": [],
|
|
"keywords": []
|
|
}
|
|
second_json = {
|
|
"uuid": time_based_uuid+"q"
|
|
}
|
|
code_doc = search_tool(json.dumps(first_json) + json.dumps(second_json))
|
|
utils.remove_shared_variable(time_based_uuid+"q")
|
|
res = get_llm_model_response(
|
|
strategy_name="default_code",
|
|
llm_model_name=LLM_MODELS[2],
|
|
template_prompt_name="default_code",
|
|
prompt_param_dict={"input": parsed_query, "code_doc": f"{code_doc}", "time": datetime.now().strftime("%Y%m%d")},
|
|
temperature=0.01,
|
|
max_tokens=512
|
|
)
|
|
res = res.replace("<think>","")
|
|
return f"{res}"
|
|
except Exception as e:
|
|
logging.error(f"Error occurred while processing math query: {e}")
|
|
return "<system>不要再调用该工具了</system>"
|
|
|
|
except Exception as e:
|
|
return "<system>不要再调用该工具了</system>"
|
|
|
|
|
|
|
|
|
|
class RagSearchInput(BaseModel):
|
|
query: str = Field(...,description="查询对象") |