import os
import dotenv
import logging
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory


logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

dotenv.load_dotenv()

# 1. 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("SILICONFLOW_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("SILICONFLOW_BASE_URL")

# 2. 初始化模型
llm = ChatOpenAI(model="deepseek-ai/DeepSeek-V3.1")

# 3. 定义 Prompt (现代版无需手动处理 question 变量)
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个万能的人工智能AI"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

# 4. 【核心改動】使用 LCEL 組合鏈
# 這裡不需要 LLMChain，直接用管道符
chain = prompt | llm

# 5. 管理記憶體 (現代版做法：使用字典存儲不同 Session 的歷史)
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

# 包裝成帶有記憶功能的鏈
with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="question",
    history_messages_key="history",
)

# 6. 執行調用
config = {"configurable": {"session_id": "xiaoming_test"}}

res1 = with_message_history.invoke({"question": "我是小明"}, config=config)
print(f"回答1: {res1.content}")

res2 = with_message_history.invoke({"question": "我是谁？"}, config=config)
print(f"回答2: {res2.content}")