# 基礎依賴：llama_index 核心庫、工作流工具、SQLAlchemy（數據庫連接）
pip install llama-index-core llama-index-utils-workflow sqlalchemy
# 向量存儲依賴：Milvus（存儲表結構信息）
pip install llama-index-vector-stores-milvus pymilvus
# 數據庫連接依賴：MySQL 驅動、加密庫（解決 MySQL 8.0+ 認證問題）
pip install pymysql cryptography
# LLM 依賴：根據實際使用的模型安裝（如 deepseek 等）
pip install deepseek-ai

from llama_index.core import SQLDatabase, StorageContext
from llama_index.core.objects import SQLTableSchema, SQLTableNodeMapping, ObjectIndex
from sqlalchemy import create_engine
from llama_index.vector_stores.milvus import MilvusVectorStore
import re
from .config import RagConfig  # 配置類：存儲數據庫連接字符串、Milvus 地址等

class DatabaseRAG:
    def __init__(self, **kwargs):
        # 1. 連接數據庫：從配置中獲取連接字符串
        self.engine = create_engine(RagConfig.db_connection_string)
        self.sql_database = SQLDatabase(self.engine)
    
    async def load_data(self):
        # 2. 加載數據庫表結構：獲取所有表名，提取表描述與字段信息
        tables = self.sql_database.get_usable_table_names()
        table_schema_objs = []
        for table in tables:
            # 獲取單表信息（包含字段、註釋）
            table_info = self.sql_database.get_single_table_info(table)
            # 提取表描述（若不存在，用表名作為描述）
            match = re.search(r"with comment: \((.*?)\)", table_info)
            table_desc = match.group(1) if match else f"{table} table"
            # 創建表結構對象（用於後續向量存儲）
            table_schema_objs.append(
                SQLTableSchema(table_name=table, context_str=table_desc)
            )
        return table_schema_objs
    
    async def create_index(self, collection_name="database"):
        # 3. 創建向量索引：將表結構信息存儲到 Milvus
        # 加載表結構數據
        table_data = await self.load_data()
        # 建立表與節點的映射關係
        table_node_mapping = SQLTableNodeMapping(sql_database=self.sql_database)
        # 初始化 Milvus 向量存儲
        vector_store = MilvusVectorStore(
            uri=RagConfig.milvus_uri,
            collection_name=collection_name,
            dim=RagConfig.embedding_model_dim,  # 嵌入模型維度（如 768）
            overwrite=True
        )
        # 創建存儲上下文
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        # 生成向量索引（用於後續檢索表信息）
        index = ObjectIndex.from_objects(
            objects=table_data,
            object_mapping=table_node_mapping,
            storage_context=storage_context
        )
        return index

from llama_index.core.indices.struct_store.sql_retriever import SQLRetriever
from llama_index.core.workflow import Workflow, Context, StartEvent, Event, step, StopEvent
from .database_rag import DatabaseRAG
from .prompts import TEXT_TO_SQL_PROMPT, RESPONSE_SYNTHESIS_PROMPT  # 預定義提示模板
from .llms import deepseek_llm  # 自定義 LLM 加載函數
from .utils import parse_response_to_sql  # 自定義 SQL 解析函數

# 自定義事件：表檢索完成事件（攜帶查詢與表信息）
class TableRetrieveEvent(Event):
    query: str
    table_content_str: str

# 自定義事件：SQL 生成完成事件（攜帶查詢與 SQL）
class TextToSQLEvent(Event):
    query: str
    sql: str

class TextToSQLWorkflow(Workflow, DatabaseRAG):
    def __init__(self, **kwargs):
        # 初始化工作流與數據庫 RAG
        Workflow.__init__(self, **kwargs)
        DatabaseRAG.__init__(self, **kwargs)
    
    @step
    async def retrieve_tables(self, ctx: Context, ev: StartEvent) -> TableRetrieveEvent:
        """步驟 1：檢索與用户查詢相關的表信息"""
        # 加載 Milvus 中的表結構索引
        index = await self.create_index(collection_name="database")
        # 檢索前 20 個最相關的表結構（避免遺漏關鍵信息）
        retriever = index.as_retriever(similarity_top_k=20)
        table_schema_objs = retriever.retrieve(ev.query)
        # 整理表信息為字符串（便於 LLM 理解）
        table_content_str = "\n\n".join([node.text for node in table_schema_objs])
        return TableRetrieveEvent(query=ev.query, table_content_str=table_content_str)
    
    @step
    async def generate_sql(self, ctx: Context, ev: TableRetrieveEvent) -> TextToSQLEvent:
        """步驟 2：根據查詢與表信息生成 SQL"""
        # 格式化提示模板（注入用户查詢與表信息）
        prompt = TEXT_TO_SQL_PROMPT.format_messages(
            query=ev.query,
            table_context_str=ev.table_content_str
        )
        # 調用 LLM 生成 SQL
        llm = deepseek_llm()
        chat_response = llm.chat(prompt)
        # 解析 LLM 響應，提取純 SQL 語句（去除多餘文本）
        sql = parse_response_to_sql(chat_response)
        return TextToSQLEvent(query=ev.query, sql=sql)
    
    @step
    async def generate_response(self, ctx: Context, ev: TextToSQLEvent) -> StopEvent:
        """步驟 3：執行 SQL 並生成自然語言響應"""
        # 執行 SQL 查詢（獲取結構化結果）
        sql_retriever = SQLRetriever(sql_database=self.sql_database)
        query_result = sql_retriever.retrieve(ev.sql)
        # 格式化響應模板（注入查詢、SQL、結果）
        response_prompt = RESPONSE_SYNTHESIS_PROMPT.format_messages(
            query=ev.query,
            sql=ev.sql,
            context_str=query_result
        )
        # 返回結果（包含響應提示與原始 SQL，便於調試）
        return StopEvent(result={"message": response_prompt, "sql": ev.sql})

from llama_index.core import PromptTemplate
from llama_index.core.prompts import PromptType

# Text2SQL 提示模板：限定 LLM 只使用指定表、生成正確語法的 MySQL 語句
TEXT_TO_SQL_PROMPT_STRING = """
給定一個輸入問題，首先創建語法正確的 MySQL 查詢來運行，然後查看查詢結果並返回答案。
注意事項：
1. 只查詢問題中需要的列，不要查詢表中所有列；
2. 僅使用下面列出的表結構信息，不要使用其他表或列；
3. 若涉及多表查詢，需通過外鍵（如 student_id）關聯，確保表關係正確；
4. 生成的 SQL 語句需包含完整的表名、列名，避免語法錯誤。

你必須使用以下格式：
Question: {query}
SQLQuery: 要運行的 SQL 查詢
SQLResult: （無需填寫，僅為格式要求）
Answer: （無需填寫，僅為格式要求）

可用表結構信息：
{table_context_str}

SQLQuery:
"""
TEXT_TO_SQL_PROMPT = PromptTemplate(
    TEXT_TO_SQL_PROMPT_STRING,
    prompt_type=PromptType.TEXT_TO_SQL
)

# 響應合成提示模板：將 SQL 結果整理為自然語言
RESPONSE_SYNTHESIS_PROMPT_STRING = """
根據用户問題、執行的 SQL 語句和查詢結果，合成簡潔易懂的自然語言回答。
要求：
1. 直接回答用户問題，不要包含 SQL 語句；
2. 若結果為表格數據，用清晰的列表或段落描述，避免展示原始數據；
3. 語言口語化，讓非技術用户能快速理解。

用户問題：{query}
執行的 SQL：{sql}
查詢結果：{context_str}

Response:
"""
RESPONSE_SYNTHESIS_PROMPT = PromptTemplate(RESPONSE_SYNTHESIS_PROMPT_STRING)

from pydantic import Field
from pydantic_settings import BaseSettings
import os

class RagConfig(BaseSettings):
    # 數據庫連接字符串（從環境變量讀取，避免硬編碼）
    db_connection_string: str = Field(
        default=os.getenv("DB_CONNECTION_STRING"),
        description="MySQL 數據庫連接字符串"
    )
    # Milvus 向量存儲配置
    milvus_uri: str = Field(default="http://localhost:19530", description="Milvus 服務地址")
    embedding_model_dim: int = Field(default=768, description="嵌入模型維度（如 BGE 模型為 768）")

# 實例化配置（全局複用）
rag_config = RagConfig()

import chainlit as cl
from llama_index.core import Settings
from .text_to_sql_workflow import TextToSQLWorkflow
from .llms import deepseek_llm

# 初始化 LLM（用於響應流式輸出）
Settings.llm = deepseek_llm()

@cl.on_chat_start
async def start():
    # 向用户發送歡迎消息
    await cl.Message(content="歡迎使用數據庫智能查詢系統！請輸入您的查詢（如“獲取所有學生的基本信息”）").send()

@cl.on_message
async def main(message: cl.Message):
    # 1. 創建 Text2SQL 工作流實例
    workflow = TextToSQLWorkflow()
    # 2. 運行工作流（傳入用户查詢）
    result = await workflow.run(query=message.content)
    # 3. 提取結果（SQL 語句與響應提示）
    sql = result.get("sql")
    response_prompt = result.get("message")
    
    # 4. 向用户展示生成的 SQL（便於調試）
    await cl.Message(content=f"生成的 SQL 語句：\n```sql\n{sql}\n```", author="系統").send()
    
    # 5. 流式輸出自然語言響應
    msg = cl.Message(content="", author="系統")
    response_gen = Settings.llm.stream_chat(response_prompt)
    for token in response_gen:
        await msg.stream_token(token.delta)
    await msg.send()

# Windows
set DB_CONNECTION_STRING=mysql+pymysql://root:root123@localhost:3306/student_db
# Linux/Mac
export DB_CONNECTION_STRING=mysql+pymysql://root:root123@localhost:3306/student_db

chainlit run rag/chainlit_ui.py -w