from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import PromptTemplate

from src.server.deep_agent.init_llm import title_llm


async def conversation_title(agent, config):
    state = agent.get_state(config)
    messages = state.values.get("messages", [])
    if len(messages) < 2:
        return None

    user_msg = None
    ai_msg = None

    for m in messages:
        if isinstance(m, HumanMessage) and user_msg is None:
            user_msg = m.content

        if isinstance(m, AIMessage) and ai_msg is None:
            ai_msg = m.content

        if user_msg and ai_msg:
            break

    prompt = f"""
    请根据以下首次对话内容，生成一个简洁、精准的标题（2-15个字）：

    用户：{user_msg}
    助手：{ai_msg}

    要求：
    1. 标题需概括对话核心内容
    2. 语言简洁、符合中文表达习惯
    3. 仅返回标题，不要额外解释
    """
    response = await title_llm.ainvoke(prompt)
    title = response.content.strip()

    # 去掉可能的符号
    title = title.replace("《", "").replace("》", "")
    return title