aiagent/backend/app/services/conversation_branch_service.py

"""
对话分支服务 — 创建、列举、恢复、删除对话分支

参考 Claude Code src/commands/branch/branch.ts
"""
from __future__ import annotations

import logging
import re
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy.orm import Session
from sqlalchemy import desc

from app.core.database import SessionLocal
from app.models.conversation_branch import ConversationBranch

logger = logging.getLogger(__name__)


def derive_title_from_messages(messages: List[Dict[str, Any]]) -> str:
    """从第一条用户消息中提取标题（截断100字）。

    参考 Claude Code branch.ts L38-54
    """
    for msg in messages:
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if content:
                # 压缩空白，截断100字符
                title = re.sub(r"\s+", " ", str(content)).strip()
                return title[:100] if len(title) > 100 else title
    return "已分支的对话"


def _get_unique_branch_title(db: Session, base_name: str, user_id: str) -> str:
    """生成唯一分支标题，避免冲突。

    参考 Claude Code branch.ts L179-220
    格式: "Base (分支)", "Base (分支 2)", "Base (分支 3)", ...
    """
    candidate = f"{base_name} (分支)"

    # 检查精确匹配
    existing = db.query(ConversationBranch).filter(
        ConversationBranch.user_id == user_id,
        ConversationBranch.title == candidate,
    ).first()

    if not existing:
        return candidate

    # 查找已有编号
    pattern = re.compile(rf"^{re.escape(base_name)} \(分支(?: (\d+))?\)$")
    used_numbers: set[int] = {1}
    all_forks = db.query(ConversationBranch).filter(
        ConversationBranch.user_id == user_id,
        ConversationBranch.title.like(f"{base_name} (分支%"),
    ).all()

    for fork in all_forks:
        match = pattern.match(fork.title or "")
        if match:
            num_str = match.group(1)
            if num_str:
                used_numbers.add(int(num_str))
            else:
                used_numbers.add(1)

    next_num = 2
    while next_num in used_numbers:
        next_num += 1

    return f"{base_name} (分支 {next_num})"


def create_branch(
    db: Session,
    user_id: str,
    parent_session_id: str,
    messages: List[Dict[str, Any]],
    agent_id: Optional[str] = None,
    agent_name: Optional[str] = None,
    custom_title: Optional[str] = None,
) -> ConversationBranch:
    """从当前会话创建新分支。

    Args:
        db: 数据库会话
        user_id: 用户ID
        parent_session_id: 原会话ID
        messages: 完整消息列表（含 system prompt）
        agent_id: 关联 Agent ID
        agent_name: Agent 名称
        custom_title: 自定义标题（可选）

    Returns:
        创建的 ConversationBranch 实例
    """
    # 生成标题
    base_name = custom_title or derive_title_from_messages(messages)
    title = _get_unique_branch_title(db, base_name, user_id)

    # 提取第一条用户消息用于列表展示
    first_user_msg = None
    for msg in messages:
        if msg.get("role") == "user":
            content = msg.get("content", "")
            first_user_msg = content[:200] if len(content) > 200 else content
            break

    branch = ConversationBranch(
        id=str(uuid.uuid4()),
        user_id=user_id,
        agent_id=agent_id,
        agent_name=agent_name,
        parent_session_id=parent_session_id,
        branch_session_id=str(uuid.uuid4()),
        title=title,
        message_count=len(messages),
        first_user_message=first_user_msg,
        messages=messages,
        is_active=True,
        created_at=datetime.now(),
    )

    db.add(branch)
    db.commit()
    db.refresh(branch)

    logger.info(
        "对话分支已创建: id=%s title=%s messages=%d parent=%s",
        branch.id, title, len(messages), parent_session_id,
    )
    return branch


def list_branches(
    db: Session,
    user_id: str,
    agent_id: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
) -> List[ConversationBranch]:
    """列出用户的所有分支。

    Args:
        db: 数据库会话
        user_id: 用户ID
        agent_id: 可选，按 Agent 过滤
        limit: 每页数量
        offset: 偏移量

    Returns:
        分支列表（按创建时间倒序）
    """
    query = db.query(ConversationBranch).filter(
        ConversationBranch.user_id == user_id,
        ConversationBranch.is_active == True,
    )
    if agent_id:
        query = query.filter(ConversationBranch.agent_id == agent_id)

    return query.order_by(desc(ConversationBranch.created_at)).offset(offset).limit(limit).all()


def get_branch(db: Session, branch_id: str, user_id: str) -> Optional[ConversationBranch]:
    """获取单个分支详情。"""
    return db.query(ConversationBranch).filter(
        ConversationBranch.id == branch_id,
        ConversationBranch.user_id == user_id,
    ).first()


def delete_branch(db: Session, branch_id: str, user_id: str) -> bool:
    """软删除分支（标记为不活跃）。

    Args:
        db: 数据库会话
        branch_id: 分支ID
        user_id: 用户ID (验证所有权)

    Returns:
        是否成功删除
    """
    branch = db.query(ConversationBranch).filter(
        ConversationBranch.id == branch_id,
        ConversationBranch.user_id == user_id,
    ).first()
    if not branch:
        return False

    branch.is_active = False
    db.commit()
    logger.info("对话分支已删除: id=%s title=%s", branch.id, branch.title)
    return True


def get_messages_from_session(db: Session, session_id: str) -> Optional[List[Dict[str, Any]]]:
    """从 execution_logs 表中恢复会话的完整消息链。

    通过 session_id 查询 agent_execution_logs，按时间顺序组装消息。
    """
    from app.models.agent_execution_log import AgentExecutionLog
    from app.models.agent_llm_log import AgentLLMLog

    logs = db.query(AgentExecutionLog).filter(
        AgentExecutionLog.session_id == session_id,
    ).order_by(AgentExecutionLog.created_at).all()

    if not logs:
        return None

    # 重建消息链 — 每一条 execution log 代表一次用户-助手交互
    messages: List[Dict[str, Any]] = []
    for log in logs:
        if log.input_text:
            messages.append({"role": "user", "content": log.input_text})
        if log.output_text:
            messages.append({"role": "assistant", "content": log.output_text})
        # 如果有步骤详情，可以重建工具调用
        if log.steps:
            for step in log.steps:
                if step.get("type") == "tool_result" and step.get("tool_name"):
                    messages.append({
                        "role": "tool",
                        "tool_call_id": step.get("tool_call_id", "unknown"),
                        "content": step.get("tool_result", ""),
                        "name": step.get("tool_name", "unknown"),
                    })

    return messages if messages else None