QingLong/AI/Vanna/VannaQwen.py

import os
import requests
from vanna.base import VannaBase
from typing import List, Dict, Any
from Config import *


class QwenChat(VannaBase):
    def __init__(self):
        super().__init__()
        self.api_key = MODEL_API_KEY
        self.model = "qwen-max"
        self.training_data = []
        self.chat_history = []  # 新增：用于存储对话历史

    # ---------- 必须实现的抽象方法 ----------
    def add_ddl(self, ddl: str, **kwargs) -> None:
        self.training_data.append({"type": "ddl", "content": ddl})

    def add_documentation(self, doc: str, **kwargs) -> None:
        self.training_data.append({"type": "documentation", "content": doc})

    def add_question_sql(self, question: str, sql: str, **kwargs) -> None:
        self.training_data.append({"type": "qa", "question": question, "sql": sql})

    def get_related_ddl(self, question: str, **kwargs) -> str:
        return "\n".join([item["content"] for item in self.training_data if item["type"] == "ddl"])

    def get_related_documentation(self, question: str, **kwargs) -> str:
        return "\n".join([item["content"] for item in self.training_data if item["type"] == "documentation"])

    def get_training_data(self, **kwargs) -> List[Dict[str, Any]]:
        return self.training_data

    # ---------- 新增的必需方法 ----------
    def system_message(self, message: str) -> None:
        """系统消息（用于初始化对话）"""
        self.chat_history.append({"role": "system", "content": message})

    def user_message(self, message: str) -> None:
        """用户消息"""
        self.chat_history.append({"role": "user", "content": message})

    def assistant_message(self, message: str) -> None:
        """助手消息"""
        self.chat_history.append({"role": "assistant", "content": message})

    def submit_prompt(self, prompt: str, **kwargs) -> str:
        """提交提示词（这里直接调用 Qwen API）"""
        return self.generate_sql(question=prompt)

    # ---------- 其他需要实现的方法 ----------
    def remove_training_data(self, id: str, **kwargs) -> bool:
        return True

    def generate_embedding(self, text: str, **kwargs) -> List[float]:
        return []

    def get_similar_question_sql(self, question: str, **kwargs) -> List[Dict[str, Any]]:
        return []

    # ---------- 原有生成 SQL 的方法 ----------
    def generate_sql(self, question: str, **kwargs) -> str:
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        prompt = f"基于以下 MySQL 表结构，生成 SQL 查询。只需返回 SQL 代码，不要解释。\n\n表结构：\n{self.get_related_ddl(question)}\n\n问题：{question}"

        data = {
            "model": self.model,
            "input": {"messages": [{"role": "user", "content": prompt}]},
            "parameters": {}
        }

        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            return response.json()["output"]["text"]
        else:
            raise Exception(f"Qwen API 错误: {response.text}")


# ----------------------------
# 初始化并使用实例
# ----------------------------
vn = QwenChat()

# 添加训练数据
vn.train(
    ddl="""
    CREATE TABLE employees (
        id INT PRIMARY KEY,
        name VARCHAR(100),
        department VARCHAR(50),
        salary DECIMAL(10, 2),
        hire_date DATE
    );
    """
)

vn.train(documentation="部门 'Sales' 的员工工资超过 50000")

# 生成 SQL
question = "显示 2020 年后入职的 Sales 部门员工姓名和工资，且工资大于 50000"
try:
    sql = vn.generate_sql(question=question)
    print("生成的 SQL:\n", sql)
except Exception as e:
    print("错误:", str(e))