# %%
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv

import os
import logging 
# %%
load_dotenv()
vllm_api_key = os.getenv("VLLM_API_KEY")
vllm_base_url = os.getenv("VLLM_BASE_URL")
vllm_model_name = os.getenv("VLLM_MODEL_NAME")
logger = logging.getLogger("vLLM Use Log")
logger_console = logging.StreamHandler()
logger.setLevel(logging.INFO)
logger_console.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s:%(message)s")
logger_console.setFormatter(formatter)
logger.addHandler(logger_console)
# %%
model = ChatOpenAI(base_url = vllm_base_url, model = vllm_model_name, api_key = vllm_api_key)
# %%
prompt = ChatPromptTemplate.from_messages(
    [
        ('system', '你是一个超级无敌厉害的物理学家！你可以回答一切关于物理的问题！'),
        ('human', '{input}')
    ]
)

chain_qwen = prompt | model | StrOutputParser()
# %%
chain_qwen.invoke({"input":"牛顿第一定律是什么"})
# %%
while True:
    input_text = input()
    if input_text == 'exit':
        print("退出模型交互！")
        break
    response = chain_qwen.invoke({"input": input_text})
    logger.info(f"Question:{input_text}")
    logger.info(f"Answer:{response}")
'''
    2025-08-20 00:06:21,317:Question:牛顿第一定律是什么
    2025-08-20 00:06:21,319:Answer:牛顿第一定律，也称为惯性定律，是经典力学中的一个基本原理。它表明：如果一个物体不受外力作用，那么这个物体将保持静止状态或者以恒定速度沿直线运动。换句话说，在没有外力作用的情况下，物体将继续在原来的状态中移动。
    这个定律由艾萨克·牛顿（Isaac Newton）在1687年提出，并且是所有现代物理学的基础之一。它是理解物体运动、惯性和能量守恒定律的关键概念。
    2025-08-20 00:06:45,389:Question:我刚刚问的第一个问题是什么
    2025-08-20 00:06:45,390:Answer:好的，我会尽力帮助您解答任何物理问题。请随时提出您的第一个问题。
    退出模型交互！
'''
# %%
# 创建消息历史
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
# %%
store = {}
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()

    return store[session_id]
# %%
with_message_history = RunnableWithMessageHistory(chain_qwen, get_session_history)
# input_messages_key="input"
config = {
    "configurable": {"session_id": "11146"}
}
# %%
while True:
    input_text = input()
    if input_text == 'exit':
        print("退出模型交互！")
        break
    response = with_message_history.invoke({"input": input_text}, config=config)
    logger.info(f"Question:{input_text}")
    logger.info(f"Answer:{response}")
'''
    2025-08-20 01:22:14,285:Question:牛顿第一定律是什么
    2025-08-20 01:22:14,286:Answer:牛顿第一定律，也称为惯性定律，是物理学中的一个基本原理。它表明：任何物体都会保持其静止状态或匀速直线运动的状态，除非外力迫使它改变这种状态。具体来说，如果没有外力作用于物体上，该物体将保持其当前的运动状态。换句话说，物体会继续保持它的速度不变或者沿原方向移动，直到有外力迫使它改变这个状态.简单来说，牛顿第一定律描述了物体在没有外部力作用时的行为规律，即它们会在保持原来的速度和方向上进行匀速直线运动。
    2025-08-20 01:22:49,850:Question:结合历史消息，我刚刚问的第一个问题题目是什么？
    2025-08-20 01:22:49,852:Answer:结合历史消息，您刚刚问的第一个问题题目是：“牛顿第一定律，也称为惯性定律，是物理学中的一个基本原理。”
    退出模型交互！
'''
# %%