import datetime import json import os import time from functools import wraps import tiktoken import requests from openai import OpenAI, Client from tenacity import before_sleep_log, retry, stop_after_attempt, wait_random_exponential, RetryError from consts.consts import GPT_DEFAULT, GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI, MODE_COMPLETION, MODE_CHAT, DebugActions, GenSQLErrorMsg, Claude3, Qwen from config import config from errors.errors import AlgorithmException, BadRequestException from svc.measuring import EDA_OPENAI_HISTOGRAM, measuring_openai from utils.local import get_local_metrics, set_local_metrics from models.user_preference import UserPreference from models.ailog import AILog from svc.claude import invoke_claude_3_with_text, extract_text from utils.logging import logging from json_repair import repair_json logger = logging.getLogger(__name__) # global client global_client = OpenAI( api_key=config.OPENAI_API_KEY, max_retries=0 ) # global assistant client global_assistant_client = OpenAI( api_key=config.OPENAI_API_KEY ) MAX_TOKEN = { GPT_4: 8000, GPT_4_TURBO: 128000, GPT_4O: 128000, GPT_4O_MINI: 128000, Qwen.TURBO: 8000, Qwen.PLUS: 128000, # qwen-plus 支持 128K 上下文 Qwen.MAX: 32000 } # Returns a maximum of 4,096 output tokens. Reserve at least 2,000 output tokens PROMPT_MAX_TOKEN = { GPT_4: 6000, GPT_4_TURBO: 126000, GPT_4O: 126000, GPT_4O_MINI: 126000, Qwen.TURBO: 6000, Qwen.PLUS: 126000, # qwen-plus 预留输出 token (128K - 2K buffer) Qwen.MAX: 30000 } PROMPT_MODE_LOCAL = "local" PROMPT_MODE = os.getenv("context_mode") or PROMPT_MODE_LOCAL SetupBlockPrompt = "setup_block" DataSetBlockingPrompt = "dataset_blocking" MainTaskBlockPrompt = "main_task_block" ChatLogsBlockPrompt = "chat_logs_block" DomainKnowledgeBlockPrompt = "domain_knowledge_block" SessionMainTaskBlockPrompt = "session_main_task_block" OutputJsonFormatBlockPrompt = "output_json_format_block" def format_prompt(template_name, variables=None): variables = variables or {} if PROMPT_MODE == PROMPT_MODE_LOCAL: return localPromptMapper[template_name] % variables url = "http://edaservice.nightly-ms.svc.cluster.local:11000/aipower/api/v1/prompt/format" payload = { "name": template_name, "vars": variables, } resp = requests.post(url, json=payload) resp.raise_for_status() resp_json = resp.json() return resp_json["result"]["full_prompt"] localPromptMapper = { SetupBlockPrompt: """ Act as a data analyst, your task is to perform data analysis using the provided information. """, DataSetBlockingPrompt: """ Information about data sources(DataBase Schema): \n %(db_schema_str)s """, MainTaskBlockPrompt: """ %(main_task)s """, ChatLogsBlockPrompt: """ --- The relevant chat logs: %(chat_logs_str)s --- """, DomainKnowledgeBlockPrompt: """ The relevant domain knowledge: %(domain_knowledge_embedding)s """, SessionMainTaskBlockPrompt: """ MainTask: %(main_task)s ChoosedAction: %(choosed_action)s SolutionPlan: %(solution_plan_str)s """, OutputJsonFormatBlockPrompt: """ Please provide the output in JSON format as described below: %(schema)s. Only output json, do not output anything else. Json: """, } """ Prompt 结构定义 * Stage 一个固定的阶段，不同阶段会有不同的 Prompt 主体结构 * Block 一个预定义结构的 Prompt 区块 * Prompt 由多个 Block 组合成的完整 Prompt Example AI Insight Stage: Prompt: --- ** Setup Block Act as a data analyst, your task is to perform a specific data analysis task based on the following information. --- ** Dataset Block DataBase Schema xxx --- ** Domain Knowledge Block xxxx --- ** Main Task Block xxxx -- ** Ouput Format Block Output Json Format follow xx schema -- """ SetupBlockPrompt = """ Act as a data analyst, your task is to perform various data analysis work based on the given information. """ class SetupBlock: @staticmethod def build_prompt(main_task, output_format=None, output_schema=None): # return format_prompt(SetupBlockPrompt) return SetupBlockPrompt class MainTaskBlock: @staticmethod def build_prompt(main_task, output_format=None, output_schema=None): return format_prompt(MainTaskBlockPrompt, {"main_task": main_task}) OutputJsonFormatBlockPrompt = """ Please provide the output in JSON format as described below: %s Only output json, do not output anything else. Json: """ class JsonFormatBlock: @staticmethod def build_prompt(main_task, output_format=None, output_schema=None): if not output_format or not output_schema: return "" if output_format == "json": return OutputJsonFormatBlockPrompt % output_schema raise AlgorithmException("not valid output format") UnderstandDBTemplate = [ MainTaskBlock, JsonFormatBlock ] Text2SQLTemplate = [ MainTaskBlock, JsonFormatBlock ] def build_prompt(stage, main_task, output_format=None, output_schema=None): if stage == "understand_db": templates = UnderstandDBTemplate elif stage == "text2sql": templates = Text2SQLTemplate elif stage == "data_analysis": templates = Text2SQLTemplate else: raise AlgorithmException("unknown stage") return "\n".join([block.build_prompt(main_task, output_format, output_schema).strip("\n") for block in templates]) def ai_completion(stage, main_task, output_format=None, output_schema=None, dry_run=False, ai_model=GPT_DEFAULT, extra_info={}): # Check param: ai_model model = ai_model or os.getenv("AI_MODEL", GPT_DEFAULT) prompt = build_prompt(stage, main_task, output_format, output_schema) # TODO: forbidden to turn to claude-3 if model in (Claude3.DEFAULT, Claude3.SONNET, Claude3.HAIKU): return invoke_claude_3_with_text(prompt, output_format, dry_run=dry_run, ai_model=ai_model, extra_info=extra_info) # Route to Qwen models if model in (Qwen.DEFAULT, Qwen.TURBO, Qwen.PLUS, Qwen.MAX): logging.info("session_id = %s, query_id = %s, module = %s, routing to Qwen model, model = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model) from svc.qwen import invoke_qwen_with_text return invoke_qwen_with_text(prompt, output_format, dry_run=dry_run, ai_model=ai_model, extra_info=extra_info) if model not in (GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI): logging.exception("session_id = %s, query_id = %s, module = %s, model not supported, model = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model) # No need to retry raise BadRequestException("Model not supported: {}".format(model)) encoding = tiktoken.encoding_for_model(model) num_tokens = len(encoding.encode(prompt)) if num_tokens == 0: raise BadRequestException("Empty prompt") if dry_run: return num_tokens <= PROMPT_MAX_TOKEN[model] if num_tokens > MAX_TOKEN[model] - 1500: # No need to retry logging.info("session_id = %s, query_id = %s, module = %s, Token limit exceeded for completion mode, model = %s, num_tokens = %s, prompt = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model, str(num_tokens), str(prompt)) raise BadRequestException(GenSQLErrorMsg.TOKEN_EXCEED_ERROR) metrics = get_local_metrics() app_id = metrics.get("app_id", 0) preference = UserPreference.load(app_id) result, response = _ai_request(model, prompt, num_tokens, output_format, mode=MODE_COMPLETION, extra_info=extra_info, openai_api_key=preference.llm_api_key) measuring_openai(stage, model, response) return result def ai_chat(messages, output_format=None, dry_run=False, ai_model=GPT_DEFAULT, extra_info={}, openai_api_key=None): # Check param: ai_model if len(messages) == 0: raise BadRequestException("Empty prompt") model = ai_model or os.getenv("AI_MODEL", GPT_DEFAULT) if model not in (GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI): raise BadRequestException("Model not supported: {}".format(model)) encoding = tiktoken.encoding_for_model(model) num_tokens = 0 # count for token for msg in messages: num_tokens += len(encoding.encode(msg.get("content"))) if dry_run: return num_tokens <= PROMPT_MAX_TOKEN[model] if num_tokens > PROMPT_MAX_TOKEN[model]: logging.info("session_id = %s, query_id = %s, module = %s, Token limit exceeded for chat mode, model = %s, num_tokens = %s, messages = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model, str(num_tokens), str(messages)) raise BadRequestException(GenSQLErrorMsg.TOKEN_EXCEED_ERROR) result, response = _ai_request(model, messages, num_tokens, output_format, mode=MODE_CHAT, extra_info=extra_info, openai_api_key=openai_api_key) measuring_openai("ai_chat", model, response) return result, response def get_job_id(): metrics = get_local_metrics() logging.info("get_job_id, metrics = %s", metrics) return metrics.get("job_id", "") def ai_request_wrapper(func): @wraps(func) def wrapper(model, prompt, num_tokens, output_format, mode=MODE_COMPLETION, extra_info=None, openai_api_key=None): # NOTE: This wrapper is only for _ai_request, please make sure the arguments are the same as _ai_request job_id = get_job_id() start_at = datetime.datetime.now() status = "success" result_text = "" attempt = 0 input_tokens = 0 output_tokens = 0 try: result, response = func(model, prompt, num_tokens, output_format, mode=mode, extra_info=extra_info, openai_api_key=openai_api_key) result_text = str(result) logging.info("Request OpenAI Success, model = %s, mode = %s, extra_info = %s, result = %s, prompt = %s, num_tokens = %s, openai_api_key = %s", model, mode, str(extra_info), result_text, str(prompt), str(num_tokens), openai_api_key) total_tokens = response.usage.total_tokens input_tokens = response.usage.prompt_tokens output_tokens = response.usage.completion_tokens m = get_local_metrics() m["total_tokens"] = total_tokens set_local_metrics(m) return result, response except Exception as e: if isinstance(e, RetryError): attempt = e.last_attempt.attempt_number logging.exception( "Request OpenAI Error, model = %s, mode = %s, extra_info = %s, prompt = %s, num_tokens = %s, openai_api_key = %s, exception = %s", model, mode, str(extra_info), str(prompt), str(num_tokens), openai_api_key, e, ) status = str(e) raise AlgorithmException(GenSQLErrorMsg.REQUEST_OPENAI_ERROR) finally: metrics = get_local_metrics() org_id = metrics.get("org_id", "") app_id = metrics.get("app_id", 0) endpoint = metrics.get("endpoint", "") app_key = metrics.get("app_key", "") user_email = metrics.get("user_email", "") time_cost = int(time.time()) - int(start_at.timestamp()) extra_info = extra_info or {} session_id = extra_info.get("session_id", "") query_id = extra_info.get("query_id", "") task_id = extra_info.get("task_id", "") module = extra_info.get("module", "") EDA_OPENAI_HISTOGRAM.labels(org_id, app_id, endpoint, app_key, status, attempt).observe(time_cost) AILog.upsert( start_at, job_id, status, prompt, result_text, model, openai_api_key or "", app_id, org_id, session_id, query_id, task_id, module, time_cost, input_tokens, output_tokens, user_email, ) return wrapper @ai_request_wrapper @retry(reraise=True, wait=wait_random_exponential(min=60, max=120), stop=stop_after_attempt(10), before_sleep=before_sleep_log(logger, logging.INFO)) def _ai_request(model, prompt, num_tokens, output_format, mode=MODE_COMPLETION, extra_info=None, openai_api_key=None): extra_info = extra_info or {} metrics = get_local_metrics() debug_action = metrics.get("debug_action", "") logging.info("session_id = %s, query_id = %s, module = %s, start to request OpenAI. model = %s, mode = %s, debug_action = %s, prompt = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model, mode, debug_action, str(prompt)) client = global_client if not openai_api_key else OpenAI(api_key=openai_api_key, max_retries=0) if mode == MODE_COMPLETION: response = client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], max_tokens=(4096 if MAX_TOKEN[model] - num_tokens > 4096 else MAX_TOKEN[model] - num_tokens), temperature=0, frequency_penalty=0, presence_penalty=0, top_p=1 ) result = response.choices[0].message.content else: response = client.chat.completions.create( model=model, messages=prompt, max_tokens=(4096 if MAX_TOKEN[model] - num_tokens > 4096 else MAX_TOKEN[model] - num_tokens), temperature=0, frequency_penalty=0, presence_penalty=0, top_p=1 ) result = response.choices[0].message.content logging.info("session_id = %s, query_id = %s, module = %s, end to request openAI. model = %s, mode = %s, input_tokens = %s, output_tokens = %s, raw_result = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model, mode, str(response.usage.prompt_tokens), str(response.usage.completion_tokens), str(result)) # extra result result = extract_text(result) if debug_action == DebugActions.OPENA_SLEEP_ONE_MINUTE: time.sleep(60) elif debug_action == DebugActions.OPENAI_SLEEP_SIX_MINUTES: time.sleep(60 * 6) elif debug_action == DebugActions.OPENAI_CRASH_IN_WORKER: raise Exception("Crash in worker") elif debug_action == DebugActions.OPENAI_INVALID_RESPONSE: result = "invalid response" logging.info("session_id = %s, query_id = %s, module = %s, extract result finished. model = %s, mode = %s, debug_action = %s, result = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), model, mode, debug_action, str(result)) if output_format == "json": return repair_json(result, True), response return result, response def ai_assistant_chat(assistant_name, thread_id, message, module, output_format, truncation_strategy={"type": "auto"}, instructions=None, ai_model=GPT_DEFAULT, llm_api_key=None, extra_info={}): # retrieval preference preference = UserPreference.load(get_local_metrics().get("app_id", 0)) if preference.llm_api_key: llm_api_key = preference.llm_api_key logging.info("session_id = %s, query_id = %s, module = %s, start to assistant chat, assistant_name = %s, thread_id = %s, message = %s, instructions = %s, module = %s, model = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(assistant_name), str(thread_id), str(message), instructions, module, ai_model) if ai_model not in (GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI): logging.info("session_id = %s, query_id = %s, module = %s, not supported for {}", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), ai_model) raise BadRequestException("Model not supported: {}".format(ai_model)) time_before = int(time.time()) input_message_id = "" output_message_id = "" try: # retrieval assistant client client = OpenAI(api_key=llm_api_key) if llm_api_key else global_assistant_client # retrieval assistant assistant = __retrieval_assistant(assistant_name, client) logging.info("session_id = %s, query_id = %s, module = %s, retrieval assistant successs, assistant_name = %s, assistant_id = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(assistant_name), str(assistant.id)) # append message to thread input_message = client.beta.threads.messages.create(thread_id, role="user", content=message) # set value input_message_id = input_message.id logging.info("session_id = %s, query_id = %s, module = %s, create input message success", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module")) # create Run run = client.beta.threads.runs.create( thread_id=thread_id, assistant_id=assistant.id, model=ai_model, temperature=0, top_p=1, instructions=instructions, truncation_strategy=truncation_strategy, tools=[] ) logging.info("session_id = %s, query_id = %s, module = %s, create run success, run = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(run), str(int(time.time()) - time_before)) # periodically retrieve the Run object if __retrieval_run_status(thread_id, run.id, client): logging.info("session_id = %s, query_id = %s, module = %s, loop run status finished, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(int(time.time()) - time_before)) # list run step to get the message_ids appended by assistant run_steps = client.beta.threads.runs.steps.list( thread_id=thread_id, run_id=run.id ) logging.info("session_id = %s, query_id = %s, module = %s, retrieval run steps, run_steps = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(run_steps), str(int(time.time()) - time_before)) total_tokens = 0 # get the message id of last step for step in run_steps: if step.status == "completed" and step.type == "message_creation": # log logging.info("session_id = %s, query_id = %s, module = %s, retrieval run step success, input_token = %s, output_token = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(step.usage.prompt_tokens), str(step.usage.completion_tokens)) total_tokens = step.usage.total_tokens # get the first message id output_message_id = step.step_details.message_creation.message_id break m = get_local_metrics() m["total_tokens"] = total_tokens set_local_metrics(m) # retrieval output message message = client.beta.threads.messages.retrieve(thread_id=thread_id, message_id=output_message_id) result = message.content[0].text.value logging.info("session_id = %s, query_id = %s, module = %s, end to assistant chat, raw_result = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(result), str(int(time.time()) - time_before)) result = extract_text(result) logging.info("session_id = %s, query_id = %s, module = %s, extract chat result finished, result = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(result)) if output_format == "json": return repair_json(result, True), input_message_id, output_message_id return result, input_message_id, output_message_id except Exception as e: logging.info("session_id = %s, query_id = %s, module = %s, assistant chat failed, error = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(e)) # delete message_id ai_assistant_delete_message(input_message_id, output_message_id, thread_id, llm_api_key, extra_info) raise e def ai_assistant_delete_message(input_message_id, output_message_id, thread_id, llm_api_key=None, extra_info={}): time_before = int(time.time()) preference = UserPreference.load(get_local_metrics().get("app_id", 0)) if preference.llm_api_key: llm_api_key = preference.llm_api_key try: client = OpenAI(api_key=llm_api_key) if llm_api_key else global_assistant_client if output_message_id: client.beta.threads.messages.delete(message_id=output_message_id, thread_id=thread_id) if input_message_id: client.beta.threads.messages.delete(message_id=input_message_id, thread_id=thread_id) logging.info("session_id = %s, query_id = %s, module = %s, delete assistant message, finished, input_message_id = %s, output_message_id = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(input_message_id), str(output_message_id), str(int(time.time()) - time_before)) except Exception as e: logging.info("session_id = %s, query_id = %s, module = %s, delete assistant message, error, input_message_id = %s, output_message_id = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(input_message_id), str(output_message_id), str(e)) def __retrieval_assistant(name, client): my_assistants = client.beta.assistants.list(order="desc", limit="20").data for assistant in my_assistants: if assistant.name == name and assistant.metadata.get("tag") == name: # exist return assistant # new create assistant = client.beta.assistants.create( name = name, model = GPT_DEFAULT, temperature=0, top_p=1, metadata={"tag": name} ) return assistant def __retrieval_run_status(thread_id, run_id, client, time_out=600): start_time = int(time.time()) while True: run_object = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run_id) if run_object.status == "completed": break elif int(time.time()) - start_time > time_out: # cancel the run, it's time out client.beta.threads.runs.cancel(thread_id=thread_id,run_id=run_id) raise Exception(GenSQLErrorMsg.ASSISTANT_TIME_OUT_ERROR) else: time.sleep(1) return True def create_thread(app_id: int, sqleditor_mode=False, qa_list=None): preference = UserPreference.load(app_id) api_key = config.OPENAI_API_KEY if sqleditor_mode: return "" if preference.llm_model not in (GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI): raise BadRequestException("Model not supported: {}".format(preference.llm_model)) if preference.llm_api_key: api_key = preference.llm_api_key messages = [] if qa_list: for qa in qa_list: messages.append({"role": "user", "content": qa.get("question")}) messages.append({"role": "assistant", "content": json.dumps(qa.get("answer"))}) logging.info("Before Create thread success, qa_list = %s, messages = %s", str(qa_list), str(messages)) client = OpenAI(api_key=api_key) thread = client.beta.threads.create(messages=messages) logging.info("Create thread success, thread_id = %s", thread.id) return thread.id def check_thread_status(app_id: int, thread_id: str): preference = UserPreference.load(app_id) api_key = config.OPENAI_API_KEY if app_id and preference.llm_api_key: api_key = preference.llm_api_key client = OpenAI(api_key=api_key) try: thread = client.beta.threads.retrieve(thread_id) logging.info("Check thread status success, thread_id = %s, thread = %s", thread_id, thread) return True except Exception as e: logging.exception("Check thread status failed, thread_id = %s, error = %s", thread_id, str(e)) return False def test_openai_api_key(api_key, ai_model): logging.info("Testing OpenAI API key %s", api_key) client = Client(api_key=api_key) result = client.models.list() logging.info(f"OpenAI API key is valid, response: {result}")