# 安装LangChain和相关依赖
pip install langchain langchain-community mysql-connector-python ollama

# 启动Ollama服务(如果未设置为自动启动)
ollama serve

CREATE DATABASE langchain_demo;
USE langchain_demo;

CREATE TABLE products (
    id INT AUTO_INCREMENT PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    category VARCHAR(50) NOT NULL,
    price DECIMAL(10,2) NOT NULL,
    stock INT NOT NULL
);

INSERT INTO products (name, category, price, stock) VALUES 
('Laptop', 'Electronics', 999.99, 50),
('Smartphone', 'Electronics', 699.99, 100),
('Desk Chair', 'Furniture', 149.99, 30),
('Coffee Maker', 'Appliances', 79.99, 25);

CREATE USER 'langchain_user'@'localhost' IDENTIFIED BY 'securepassword';
GRANT ALL PRIVILEGES ON langchain_demo.* TO 'langchain_user'@'localhost';
FLUSH PRIVILEGES;

from mysql.connector import connect, Error

def get_db_connection():
    try:
        connection = connect(
            host="localhost",
            user="langchain_user",
            password="securepassword",
            database="langchain_demo"
        )
        return connection
    except Error as e:
        print(f"Error connecting to MySQL: {e}")
        return None

from langchain_community.llms import Ollama
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from db_connection import get_db_connection

# 初始化Ollama LLM (使用mistral模型)
llm = Ollama(model="mistral")

# MySQL连接测试函数
def test_mysql_connection():
    conn = get_db_connection()
    if conn:
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM products LIMIT 1")
        result = cursor.fetchone()
        print("MySQL连接测试成功，示例数据:", result)
        conn.close()

# SQL查询执行函数
def execute_sql_query(query):
    conn = get_db_connection()
    if not conn:
        return "无法连接到数据库"

    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(query)

        if query.strip().lower().startswith("select"):
            results = cursor.fetchall()
            return results if results else "查询无结果"
        else:
            conn.commit()
            return f"操作成功，影响行数: {cursor.rowcount}"

    except Exception as e:
        return f"SQL执行错误: {str(e)}"
    finally:
        if conn:
            conn.close()

# NL转SQL的提示模板
prompt_template = """
你是一个专业的SQL编写助手。根据用户的问题和下面的表结构，编写合适的MySQL查询语句。

表结构:
products(id, name, category, price, stock)

问题: {question}

请只返回SQL查询语句，不要包含其他解释或说明。
"""

prompt = PromptTemplate.from_template(prompt_template)
sql_chain = LLMChain(llm=llm, prompt=prompt)

def natural_language_to_sql(question):
    # Step1: NL转SQL
    generated_sql = sql_chain.run(question=question)

    # Step2: SQL执行并获取结果
    query_result = execute_sql_query(generated_sql)

    # Step3: LLM解释结果(可选)
    explanation_prompt = f"""
    原始问题: {question}
    生成的SQL: {generated_sql}
    查询结果: {query_result}

    请用自然语言解释这些结果。
    """

    explanation = llm(explanation_prompt)

    return {
        "generated_sql": generated_sql,
        "query_result": query_result,
        "explanation": explanation.strip()
    }

if __name__ == "__main__":
    test_mysql_connection()

    while True:
        question = input("\n请输入您的问题(输入q退出): ")
        if question.lower() == 'q':
            break

        response = natural_language_to_sql(question)

        print("\n=== SQL生成 ===")
        print(response["generated_sql"])

        print("\n=== 查询结果 ===")
        print(response["query_result"])

        print("\n=== AI解释 ===")
        print(response["explanation"])

python ollama_mysql.py