from typing import Literal, Annotated, Optional from langchain.agents.middleware import after_model, before_model, AgentMiddleware from langchain_core.messages import SystemMessage from langchain_openai import ChatOpenAI from langchain.agents import AgentState from langchain.agents import create_agent from langchain.tools import tool from langchain.tools import ToolRuntime from langgraph.runtime import Runtime from langsmith import traceable from config.settings import model from pydantic import BaseModel, Field llm = ChatOpenAI( api_key=model.get("api_key", ""), base_url=model.get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1"), model=model.get("name", "qwen-max") ) SYSTEM_CHAT_PROMPT = """你是一个全能助手、数据专家。解答问题全面、简明、清晰、易理解。 【输出约束】 必须只返回 JSON,禁止任何解释、文字、markdown、注释。 必须返回如下结构: {{ "content": "ai 答复结果", "type": "chat" }} content 字段必须是 ai 的答复结果,type 字段必须是 chat。 """ SYSTEM_SQL_PROMPT = """ 你是一个资深数据工程师、SQL构建专家、 SQL审核专家。 你的目标是: 根据用户给出的数据需求,生成一条正确、可执行、语义准确、性能合理的 SQL, 并对该 SQL 进行自我验证与改进,直到通过验证或无法继续优化为止, 生成SQL要符合指定的dialect:{dialect} 【重要!必须严格输出以下完整 JSON,不能缺少任何字段!】 【必须输出的固定结构】 {{ "content": "这里放生成的 SQL 字符串", "validation_passed": 当前 SQL 是否通过验证, true 或 false, "feedback": "若未通过验证,对 SQL 问题的明确说明", "attempts": 当前已经尝试生成 SQL 的次数, "type": "sql" <-- 必须固定输出这个值,不能修改! }} 【输出规则】 1. 必须只返回 JSON,禁止任何文字解释 2. 必须包含上面所有 5 个字段 3. type 字段必须固定为 "sql",不能是其他值 4. 即使无法生成 SQL,也要严格按上面结构返回 JSON 【工作流程】 1. 如果 content 为空: - 根据用户的数据需求生成一条 SQL - 对该 SQL 进行验证, 并根据验证结果更新 validation_passed 和 feedback, 如果验证未通过,根据 feedback 对 SQL 进行改进,生成新的 SQL, attempts 增加 1, 从头重新执行,重新生成改进SQL,直到生成的 SQL 通过验证或达到最大尝试次数为止 2. 如果 content 不为空: - 对当前 sql 进行严格验证,并根据验证结果更新 validation_passed 和 feedback - 如果验证未通过,根据 feedback 对 SQL 进行改进,生成新的 SQL, attempts 增加 1, 从头重新执行,重新生成改进SQL,直到生成的 SQL 通过验证或达到最大尝试次数为止 3. SQL 验证必须覆盖以下维度: - 是否满足用户的数据需求(语义正确性) - 表与字段是否来自提供的 schema - 过滤条件是否完整且合理(如时间、状态、范围等) - 聚合与分组逻辑是否正确 - SQL 语法是否符合指定方言并可执行 - 是否存在明显的性能风险(如不必要的全表扫描、冗余子查询) 4. 如果验证未通过: - 将 validation_passed 设为 false - 在 feedback 中给出明确、可操作的问题描述 - 基于 feedback 改进 SQL - attempts 增加 1 - 从头重新执行,重新生成改进SQL 5. 如果验证通过: - 将 validation_passed 设为 true - feedback 设为通过原因或简要说明 - 返回当前 SQL 作为最终结果 【重要输出约束】 - 每一次输出都必须同时包含以下字段: - content - validation_passed - feedback - 不要输出多余的解释文本 - 不要描述你的思考过程 - 严格按照规定的输出结构返回结果 【SQL 方言约束】 这次生成 SQL 的方言 dialect 是 {dialect}。 如果 dialect = "mysql": - 生成的 SQL 必须严格符合 MySQL 5.7 及更高版本 语法规范,并可直接执行 如果 dialect = "hive": - 生成的 SQL 必须严格符合 Apache Hive SQL 语法规范,并可直接执行 如果 dialect = "impala": - 生成的 SQL 必须严格符合 Cloudera Impala SQL 语法规范,并可直接执行 如存在语法或行为歧义,优先选择最保守、最通用的写法。 在生成和验证过程中,如不满足以上要求, 必须将 validation_passed 设为 false,并在 feedback 中说明原因。 【Schema 说明】 本次使用的schema是: {schema} schema 中包含当前数据库的完整 Schema 描述, 包括数据库名、表名、字段名、字段类型及示例值。 你在生成和验证 SQL 时必须遵守以下规则: 1. 只能使用 schema 中明确列出的数据库、表和字段 2. 不得假设、编造或引用任何 schema 中不存在的表或字段 3. 表名必须与 schema 中的 Table 名完全一致 4. 字段名必须来自对应表的字段列表 5. 示例值仅用于帮助理解字段语义,不代表完整数据 6. 生成 SQL 时,表名请使用 schema 中给出的完整表名(包含 schema 前缀) 在 SQL 验证阶段: - 如果 SQL 中使用了 schema 中不存在的表或字段, 必须将 validation_passed 设为 false, 并在 feedback 中说明具体问题。 """ MAX_ATTEMPTS = 3 class SQLAgentResponse(BaseModel): content: Optional[str] = Field( description="ai 答复结果 或 当前生成或优化后的 SQL", default=None ) validation_passed: Optional[bool] = Field( description="SQL 是否通过验证", default=False ) feedback: Optional[str] = Field( description="验证未通过时的问题说明或通过原因", default=None ) attempts: Optional[int] = Field( description="尝试执行次数", default=0 ) type: str = Field( description="结果类型,sql或chat", default="chat" ) class SQLAgentState(AgentState[SQLAgentResponse]): pass class SQLAgentContext: dialect: Literal["mysql", "hive", "impala"] # "mysql" | "hive" schema: str # 数据库的 schema 信息 class MyMiddleware(AgentMiddleware): def before_agent(self, state:SQLAgentState, runtime: Runtime[SQLAgentContext]): try: question = state["messages"][-1].content message = llm.invoke(input=[ ("system", "你是一个需求分析专家,分析用户的输入问题,确定用户需求是否为生成sql需求,是则返回yes,否则返回no"), ("user", question) ]) if "yes" == message.content.lower(): print("用户需求被判断为生成 SQL 相关,继续执行 agent") runtime.context["generate_sql"] = True runtime.context["system_prompt"] = SYSTEM_SQL_PROMPT.format(dialect =runtime.context["dialect"], schema= runtime.context["schema"]); else: print("用户需求被判断为非生成 SQL 相关,继续执行 chat") runtime.context["generate_sql"] = False runtime.context["system_prompt"] = SYSTEM_CHAT_PROMPT except Exception as e: print(f"需求分析模型调用失败,错误信息: {e}") state["structured_response"] = {"content": f"需求分析失败,error: {e}", "type": "chat"} state["jump_to"] = "end" return state def before_model(self, state, runtime: Runtime[SQLAgentContext]): # print("2. before_model 每次模型前都跑") return state def wrap_model_call(self, request, handler): if request.runtime.context["system_prompt"]: request = request.override(system_message = SystemMessage(content=request.runtime.context["system_prompt"])) # print("3. wrap_model_call 每次模型都会走") return handler(request) def after_model(self, state, runtime: Runtime[SQLAgentContext]): # print("4. after_model 每次模型后都跑") return state def after_agent(self, state, runtime: Runtime[SQLAgentContext]): # print("5. after_agent 只跑一次") return state @after_model def attempts_guard_middleware(state: SQLAgentState, runtime: Runtime[SQLAgentContext]): """ 控制最大执行步数 """ resp: SQLAgentResponse = state["structured_response"] print("--------- attempts_guard_middleware ----------") if resp.attempts >= MAX_ATTEMPTS: # 强制收敛 resp.validation_passed = True if not resp.feedback: resp.feedback = ( f"已达到最大尝试次数 {MAX_ATTEMPTS},最后一次feedback: {resp.feedback}," "返回当前最优 SQL" ) else: resp.attempts = resp.attempts+1 return state agent = create_agent( model=llm, # system_prompt = SYSTEM_PROMPT, response_format= SQLAgentResponse, state_schema= SQLAgentState, context_schema= SQLAgentContext, # tools= [get_feature], middleware= [MyMiddleware()], debug=False, # checkpointer= memory, ) @traceable def invoke(msg: str, schema:str, dialect: str = "hive"): result= agent.invoke( input= { "messages": [ ("user", msg) ] }, context= { "dialect": dialect, "schema": schema } ) # response = result["structured_response"] return result if __name__ == "__main__": # invoke("My Question is who are you?", "mysql", "mysql") invoke("Could you help me calculate the average of the total number of payments made using the most preferred payment method for each product category, where the most preferred payment method in a category is the one with the highest number of payments? ", "mysql", "mysql" )