agent.py

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.deepseek import DeepSeekProvider
from tools import get_db_structure, execute_sql_tool, execute_sql_tool_with_screenshot,log_print
from dotenv import load_dotenv
import os  # 添加os模块用于加载环境变量
from divide import extract_tasks
import json
load_dotenv()

# from be_dict import my_json
api_key = os.getenv("DEEPSEEK_API_KEY")
# 配置Deepseek模型
model = OpenAIModel(
    "deepseek-chat",
    provider=DeepSeekProvider(api_key=api_key),
)

def get_sysprompt():

    # 这里存储所有的数据库名称
    db_names = ['studentcourse']
    result = []
    # 遍历 db_names，获取每个数据库的结构
    for i in db_names:
        result.append(get_db_structure(i))
    # 获取结果数量
    count = len(result)
    # 连接所有数据库的结构和视图信息
    all_structures = ""
    for db_info in result:  # db_info 是列表中的每个字典
        db_name = db_info["db_name"]
        structure = db_info["structure"]
        views = db_info.get("views", "无视图")  # 如果没有视图，默认显示“无视图”

        all_structures += f"数据库名称：{db_name}\n数据库的结构：\n{structure}\n数据库视图：\n{views}\n\n"
    return  all_structures

# print(all_structures)
sysprompt=f"""
严禁自行编造表名或字段名或者视图，必须完全使用上述结构中提供的表名和字段。
    禁止向用户提问或者让用户选择，自行生成对应sql server 语句
    
    
【任务要求】
当用户给出自然语言描述时：
1. 你需要准确理解用户的问题逻辑；
2. 根据题意生成 一条或多条合法的 SQL Server 语句；
3. 所有调用都必须用 JSON 对象形式输出，不要有解释或文字说明。

严禁生成 Markdown、注释、说明或任何除 JSON 外的输出；
不得简化或模糊用户的题目要求

调用结束返回的结果必须是有效的 JSON 格式，不得包含任何文字说明或额外内容。
返回的 JSON 应该包含三个字段：`before_img`，`sql`，和 `after_img`。
`before_img` 是 执行前截图 SQL的路径，可为空 
`after_img` 执行后截图 SQL的路径，可为
，`sql` 是实际执行的 SQL 语句。
示例格式：
{{
  "before_img": [],  
  "sql": "INSERT INTO table_name VALUES (...);",  
  "after_img": []   // 
}}


生成的 JSON 必须严格遵守上面的格式，不得包含额外的描述文字或说明。
【输出要求】
- 所有的sql server 语句 前面不能包含 use db_name 的语句
- 不允许分多次输出字段，必须一次性输出完整 JSON。
- 不得添加解释、说明、Markdown。
-before_img 最多三个  after_img 最多三个
- 主 SQL（sql 字段）不能为空，必须符合用户任务逻辑。
- 所有 SQL 均使用 SQL Server 语法（例如 TOP、GETDATE() 等）。
- 若操作会影响表数据（INSERT/UPDATE/DELETE），必须同时提供 before_img 和 after_img
【注意事项】
- 若用户未指定数据库名，则默认使用当前 db_name；
- 若用户请求查看或比较执行前后变化，必须同时生成 before_img 和 after_img；
- 所有操作均使用已存在的数据库与表结构。
"""
structure=get_sysprompt()
# 创建agent对象
agent = Agent(
    model,
    system_prompt=f"""
你是一个 SQL Server 数据库实验报告助手。请根据用户的自然语言描述，生成并执行 SQL Server 语句。
"""+sysprompt,
    tools=[
        execute_sql_tool,
    ],
)

def main():

    # user_promot
    user_prompt = '''
    任务0   0 是 num 的值
    以StudentCourse数据库为基础，用SQL语句实现以下操作：
    
    
    1、创建一个行列子集视图Stu _View，可以看到计算机科学与技术专业的学生情况。
    （1）建立视图的SQL语句（文本）：
    （2）查看原基本表、视图内容：
    '''
    # 运行 Agent
    # agent.run_sync() 就是 pydantic_ai 提供的一个同步（非异步）执行版本
    response = agent.run_sync(user_prompt)
    print("模型回复：")
    print(response)

def is_valid_json(json_string):
    try:
        json.loads(json_string)
        return True  # 如果解析成功，返回 True
    except json.JSONDecodeError:
        return False  # 如果解析失败，返回 False
ans=[]
def run_tasks():

    file_path = os.getenv("FILE_PATH", "")  # 读取文件路径
    # 获取任务列表
    tasks = extract_tasks(file_path)
    pre_tasks=''
    for index, task in enumerate(tasks):
        o_task=task
        task=f"调用execute_sql_tool函数时，num 的值是{index}"+task
        log_print(task)
        structure = '共有两个数据库。它们的结构分别是\n'+get_sysprompt()
        user_prompt=structure+task+pre_tasks
        response = agent.run_sync(user_prompt)
        log_print('agent.py中LLM的返回值')
        log_print(response.output)

        if index == 0:
            pre_tasks = "之前执行过的任务描述为" + o_task
        else:
            pre_tasks = pre_tasks + o_task

        my_dict=dict()
        if is_valid_json(response.output):
            my_dict = json.loads(response.output)
            ans.append(my_dict)
        else:
            log_print(f"返回值不是有效的 JSON: {response.output}")
            ans.append(my_dict)




if __name__=='__main__':
    # main()
    run_tasks()
    log_print('插入信息所用的ans字典')
    log_print(ans)
    with open("./ans.json", 'w', encoding='utf-8') as f:
        json.dump(ans, f, ensure_ascii=False, indent=4)