JARVIS/backend/app/services/agent_service.py

"""
Jarvis Agent 服务层
负责 LangGraph Agent 的调用、流式输出、对话历史管理
"""

import json
import uuid
import logging
from datetime import UTC, datetime
from typing import Any, AsyncGenerator
import asyncio
from openai import BadRequestError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from langchain_core.messages import HumanMessage, AIMessage

from app.database import async_session
from app.logging_utils import summarize_llm_config

from app.models.conversation import Conversation, Message
from app.models.user import User
from app.agents.graph import get_agent_graph
from app.agents.context import set_current_user, clear_current_user
from app.services import memory_service
from app.services.brain_service import BrainService
from app.services.llm_service import create_llm_from_config, resolve_provider_capabilities
from app.agents.tools.time_reasoning import extract_reference_datetime
from app.agents.state import initial_state

logger = logging.getLogger(__name__)


def _is_streaming_rejection_error(error: Exception, user_llm_config: dict | None) -> bool:
    capabilities = resolve_provider_capabilities(user_llm_config)
    error_text = str(error).lower()
    markers = [
        "invalid chat setting",
        "invalid params",
        "stream",
        "streaming",
        "unsupported",
        "bad_request_error",
        "http 400",
        "error code: 400",
    ]

    if isinstance(error, BadRequestError):
        return (
            getattr(capabilities, "provider", None) not in {"openai", "claude"}
            and any(marker in error_text for marker in markers)
        )

    return any(marker in error_text for marker in markers)


def _coerce_event_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict):
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
        return "".join(parts)
    return str(content) if content else ""


_CONTINUITY_STATE_VERSION = 1
_CONTINUITY_SNAPSHOT_FIELDS = (
    "turn_context",
    "routing_decision",
    "continuity_state",
    "pending_action",
    "last_completed_action",
    "clarification_context",
    "tool_outcomes",
    "pending_tasks",
    "completed_tasks",
    "created_entities",
    "current_agent",
    "next_step",
    "agent_trace",
)


def _build_continuity_snapshot(state: dict[str, Any]) -> dict[str, Any] | None:
    snapshot = {
        field: state.get(field)
        for field in _CONTINUITY_SNAPSHOT_FIELDS
        if state.get(field) is not None
    }
    if not snapshot:
        return None
    return {
        "version": _CONTINUITY_STATE_VERSION,
        "state": snapshot,
    }


def _extract_continuity_snapshot(payload: Any) -> dict[str, Any] | None:
    if isinstance(payload, list):
        for item in payload:
            snapshot = _extract_continuity_snapshot(item)
            if snapshot:
                return snapshot
        return None
    if not isinstance(payload, dict):
        return None
    if payload.get("kind") != "agent_continuity_state":
        return None
    if payload.get("version") != _CONTINUITY_STATE_VERSION:
        return None
    state = payload.get("state")
    if isinstance(state, dict):
        return state
    return None


class AgentService:
    """对话 Agent 服务"""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def _try_auto_summarize_background(self, user_id: str, conversation_id: str) -> None:
        async with async_session() as session:
            await memory_service.try_auto_summarize(session, user_id, conversation_id)

    def _build_progress_event(
        self,
        stage: str,
        label: str,
        *,
        agent: str | None = None,
        tool_name: str | None = None,
        step: str | None = None,
        steps: list[str] | None = None,
    ) -> dict[str, Any]:
        return {
            "type": "progress",
            "stage": stage,
            "label": label,
            "agent": agent,
            "tool_name": tool_name,
            "step": step,
            "steps": steps or [],
        }

    def _build_current_datetime_context(self) -> tuple[str, dict[str, str]]:
        now_utc = datetime.now(UTC)
        reference = {
            "current_time_iso": now_utc.isoformat(),
            "current_date_iso": now_utc.date().isoformat(),
        }
        context = (
            "【当前时间】\n"
            f"- current_time_utc: {reference['current_time_iso']}\n"
            f"- current_date_utc: {reference['current_date_iso']}\n"
            "说明：解析‘今天/明天/后天/本周/下周’等相对时间时，请以 current_time_utc 为准。"
        )
        return context, reference

    async def _get_user_llm_config(self, user_id: str, model_name: str | None = None) -> dict | None:
        """获取用户的 LLM 模型配置"""
        user = await self.db.get(User, user_id)
        if not user or not user.llm_config:
            return None

        llm_config = user.llm_config

        if model_name:
            models = llm_config.get("chat", [])
            for m in models:
                if m.get("name") == model_name:
                    return m
            return None

        chat_models = llm_config.get("chat", [])
        for m in chat_models:
            if m.get("enabled"):
                return m

        return None

    async def _load_continuity_snapshot(self, conversation: Conversation) -> dict[str, Any] | None:
        snapshot = _extract_continuity_snapshot(conversation.agent_state)
        if snapshot:
            return snapshot

        result = await self.db.execute(
            select(Message)
            .where(Message.conversation_id == conversation.id, Message.role == "assistant")
            .order_by(Message.created_at.desc())
        )
        for message in result.scalars():
            snapshot = _extract_continuity_snapshot(message.attachments)
            if snapshot:
                return snapshot
        return None

    async def _build_agent_state(
        self,
        *,
        user_id: str,
        conversation: Conversation,
        full_message: str,
        memory_context: str | None,
        current_datetime_context: str,
        current_datetime_reference: dict[str, str],
        user_llm_config: dict | None,
    ) -> dict[str, Any]:
        state = initial_state(user_id, conversation.id)
        state.update({
            "messages": [HumanMessage(content=full_message)],
            "memory_context": memory_context,
            "current_datetime_context": current_datetime_context,
            "current_datetime_reference": current_datetime_reference,
            "user_llm_config": user_llm_config,
        })
        previous_snapshot = await self._load_continuity_snapshot(conversation)
        if previous_snapshot:
            state.update(previous_snapshot)
            state["messages"] = [HumanMessage(content=full_message)]
        return state

    async def chat(
        self,
        user_id: str,
        message: str,
        conversation_id: str | None = None,
        file_ids: list[str] | None = None,
        model_name: str | None = None,
    ) -> tuple[str, str, AsyncGenerator[dict[str, Any], None]]:
        """
        处理对话请求（流式）
        """
        user_llm_config = await self._get_user_llm_config(user_id, model_name)
        model_name_used = model_name
        if model_name and not user_llm_config:
            raise ValueError("所选模型不可用于聊天，请切换到聊天模型")
        if user_llm_config:
            model_name_used = user_llm_config.get("name", model_name)

        logger.info(
            "agent_chat_started",
            extra={
                "details": {
                    "mode": "stream",
                    "requested_model_name": model_name,
                    "resolved_model_name": model_name_used,
                    "message_length": len(message or ""),
                }
            },
        )

        if conversation_id:
            result = await self.db.execute(
                select(Conversation).where(
                    Conversation.id == conversation_id,
                    Conversation.user_id == user_id,
                )
            )
            conv = result.scalar_one_or_none()
            if conv is None:
                raise ValueError("会话不存在或无权访问")
        else:
            conv = None

        if not conv:
            conv = Conversation(user_id=user_id, title=message[:50])
            self.db.add(conv)
            await self.db.commit()
            await self.db.refresh(conv)
            conversation_id = conv.id
        else:
            conversation_id = conv.id

        file_context = ""
        if file_ids:
            from app.services.document_service import DocumentService
            doc_svc = DocumentService(self.db)
            for file_id in file_ids:
                content = await doc_svc.get_document_content(user_id, file_id)
                if content:
                    file_context += f"\n\n[用户上传文件内容]\n{content}\n[/文件内容]"

        full_message = f"{message}\n{file_context}" if file_context else message

        user_msg = Message(
            conversation_id=conversation_id,
            role="user",
            content=message,
            attachments=[{"file_ids": file_ids}] if file_ids else None,
        )
        self.db.add(user_msg)
        await self.db.commit()
        await self.db.refresh(user_msg)

        brain_service = BrainService(self.db)
        await brain_service.create_event(
            user_id,
            source_type="conversation",
            source_id=conversation_id,
            event_type="message_created",
            title="User message",
            content_summary=message[:500],
            raw_excerpt=message[:2000],
            metadata_={"role": "user"},
            importance_signal=1.0,
        )
        await self.db.commit()

        memory_ctx = await memory_service.build_memory_context(
            self.db, user_id, conversation_id, message
        )

        assistant_msg = Message(
            conversation_id=conversation_id,
            role="assistant",
            content="",
            model=model_name_used or "jarvis",
            attachments=None,
        )
        self.db.add(assistant_msg)
        await self.db.commit()
        await self.db.refresh(assistant_msg)

        def _build_assistant_event_payload(content: str) -> dict[str, Any]:
            return {
                "source_type": "conversation",
                "source_id": conversation_id,
                "event_type": "message_created",
                "title": "Assistant message",
                "content_summary": content[:500],
                "raw_excerpt": content[:2000],
                "metadata_": {"role": "assistant"},
                "importance_signal": 0.8,
            }

        async def run_agent():
            collected = ""
            state: dict[str, Any] | None = None
            set_current_user(user_id)
            try:
                graph = get_agent_graph()
                current_datetime_context, current_datetime_reference = self._build_current_datetime_context()

                state = await self._build_agent_state(
                    user_id=user_id,
                    conversation=conv,
                    full_message=full_message,
                    memory_context=memory_ctx,
                    current_datetime_context=current_datetime_context,
                    current_datetime_reference=current_datetime_reference,
                    user_llm_config=user_llm_config,
                )

                yield self._build_progress_event("thinking", "Jarvis 正在分析请求", agent="master", step="理解你的问题")

                try:
                    async for event in graph.astream_events(state, version="v2"):
                        kind = event.get("event")
                        event_name = event.get("name", "")
                        metadata = event.get("metadata", {})
                        data = event.get("data", {})

                        if kind == "on_chain_start" and event_name in {"master", "schedule_planner", "executor", "librarian", "analyst"}:
                            stage_map = {
                                "master": ("thinking", "Jarvis 正在理解请求"),
                                "schedule_planner": ("planning", "Jarvis 正在编排日程"),
                                "executor": ("tool", "Jarvis 正在执行操作"),
                                "librarian": ("tool", "Jarvis 正在检索知识"),
                                "analyst": ("thinking", "Jarvis 正在分析信息"),
                            }
                            stage, label = stage_map.get(event_name, ("thinking", "Jarvis 正在思考"))
                            yield self._build_progress_event(stage, label, agent=event_name, step=label)
                        
                        elif kind == "on_tool_start":
                            yield self._build_progress_event(
                                "tool",
                                f"Jarvis 正在调用工具 {event_name}",
                                agent="executor",
                                tool_name=event_name,
                                step=f"正在执行 {event_name}",
                            )
                        
                        elif kind == "on_tool_end":
                            tool_result = data.get("output")
                            step = f"已完成 {event_name}"
                            if isinstance(tool_result, str) and len(tool_result) > 0:
                                step = tool_result[:100]
                            yield self._build_progress_event(
                                "tool",
                                f"工具 {event_name} 已完成",
                                agent="executor",
                                tool_name=event_name,
                                step=step,
                            )
                        
                        elif kind == "on_chat_model_stream":
                            chunk = data.get("chunk")
                            content = _coerce_event_text(getattr(chunk, "content", "") if chunk else "")
                            if content:
                                collected += content
                                yield {"type": "chunk", "content": content}
                        
                        elif kind == "on_chain_end":
                            output = data.get("output")
                            final_resp = None
                            if isinstance(output, dict):
                                state.update(output)
                                final_resp = output.get("final_response")
                            if final_resp:
                                final_text = str(final_resp)
                                if final_text != collected:
                                    collected = final_text
                                    yield {"type": "chunk", "content": final_text}

                        elif kind == "on_chat_model_end":
                            output = data.get("output")
                            final_content = _coerce_event_text(getattr(output, "content", "") if output else "")
                            if final_content:
                                final_text = final_content
                                if final_text != collected:
                                    collected = final_text
                                    yield {"type": "chunk", "content": final_text}

                except Exception as e:
                    if _is_streaming_rejection_error(e, user_llm_config) and not collected:
                        yield self._build_progress_event("responding", "Jarvis 正在生成回复", agent="master", step="fallback")
                        try:
                            result_state = await graph.ainvoke(state)
                            if isinstance(result_state, dict):
                                state.update(result_state)
                            fallback_content = result_state.get("final_response") or str(result_state.get("messages", [AIMessage(content="")])[-1].content)
                            collected = str(fallback_content)
                            yield {"type": "chunk", "content": collected}
                        except Exception:
                            logger.exception("llm_sync_fallback_failed")
                            safe_error = "模型服务暂不可用，请稍后再试。"
                            yield {"type": "error", "error": safe_error}
                            collected = f"抱歉，发生错误: {safe_error}"
                            yield {"type": "chunk", "content": collected}
                    else:
                        logger.exception("agent_streaming_failed")
                        if not collected:
                            safe_error = "模型服务暂不可用，请稍后再试。"
                            yield {"type": "error", "error": safe_error}
                            collected = f"抱歉，发生错误: {safe_error}"
                            yield {"type": "chunk", "content": collected}
                        else:
                            yield {"type": "error", "error": str(e)}
            finally:
                clear_current_user()
                try:
                    if collected:
                        assistant_msg.content = collected
                        continuity_snapshot = _build_continuity_snapshot(state or {})
                        assistant_msg.attachments = ([{
                            "kind": "agent_continuity_state",
                            **continuity_snapshot,
                        }] if continuity_snapshot else None)
                        conv.agent_state = continuity_snapshot
                        await BrainService(self.db).create_event(
                            user_id,
                            **_build_assistant_event_payload(collected),
                        )
                        await self.db.commit()
                        await self.db.refresh(assistant_msg)
                except Exception:
                    logger.exception("save_assistant_message_failed")
                asyncio.create_task(self._try_auto_summarize_background(user_id, conversation_id))

        return conversation_id, assistant_msg.id, run_agent()

    async def chat_simple(
        self,
        user_id: str,
        message: str,
        conversation_id: str | None = None,
        file_ids: list[str] | None = None,
        model_name: str | None = None,
    ) -> tuple[str, str, str, str | None]:
        """
        简单同步版对话
        """
        user_llm_config = await self._get_user_llm_config(user_id, model_name)
        model_name_used = model_name
        if model_name and not user_llm_config:
            raise ValueError("所选模型不可用于聊天，请切换到聊天模型")
        if user_llm_config:
            model_name_used = user_llm_config.get("name", model_name)

        if conversation_id:
            result = await self.db.execute(
                select(Conversation).where(
                    Conversation.id == conversation_id,
                    Conversation.user_id == user_id,
                )
            )
            conv = result.scalar_one_or_none()
            if conv is None:
                raise ValueError("会话不存在或无权访问")
        else:
            conv = None

        if not conv:
            conv = Conversation(user_id=user_id, title=message[:50])
            self.db.add(conv)
            await self.db.commit()
            await self.db.refresh(conv)
            conversation_id = conv.id
        else:
            conversation_id = conv.id

        user_msg = Message(conversation_id=conversation_id, role="user", content=message)
        self.db.add(user_msg)

        assistant_msg = Message(
            conversation_id=conversation_id,
            role="assistant",
            content="",
            model=model_name_used or "jarvis",
            attachments=None,
        )
        self.db.add(assistant_msg)

        brain_service = BrainService(self.db)
        await brain_service.create_event(
            user_id,
            source_type="conversation",
            source_id=conversation_id,
            event_type="message_created",
            title="User message",
            content_summary=message[:500],
            raw_excerpt=message[:2000],
            metadata_={"role": "user"},
            importance_signal=1.0,
        )

        memory_ctx = await memory_service.build_memory_context(self.db, user_id, conversation_id, message)

        set_current_user(user_id)
        try:
            graph = get_agent_graph()
            current_datetime_context, current_datetime_reference = self._build_current_datetime_context()
            state = await self._build_agent_state(
                user_id=user_id,
                conversation=conv,
                full_message=message,
                memory_context=memory_ctx,
                current_datetime_context=current_datetime_context,
                current_datetime_reference=current_datetime_reference,
                user_llm_config=user_llm_config,
            )

            result_state = await graph.ainvoke(state)
            response_content = result_state.get("final_response") or str(result_state.get("messages", [AIMessage(content="")])[-1].content)
        except Exception as e:
            logger.exception("agent_chat_simple_failed")
            response_content = "抱歉，发生错误。"
        finally:
            clear_current_user()

        brain_service = BrainService(self.db)
        await brain_service.create_event(
            user_id,
            source_type="conversation",
            source_id=conversation_id,
            event_type="message_created",
            title="Assistant message",
            content_summary=response_content[:500],
            raw_excerpt=response_content[:2000],
            metadata_={"role": "assistant"},
            importance_signal=0.8,
        )

        assistant_msg.content = response_content
        continuity_snapshot = _build_continuity_snapshot(result_state) if 'result_state' in locals() else None
        assistant_msg.attachments = ([{
            "kind": "agent_continuity_state",
            **continuity_snapshot,
        }] if continuity_snapshot else None)
        conv.agent_state = continuity_snapshot
        await self.db.commit()
        await self.db.refresh(assistant_msg)

        return conversation_id, assistant_msg.id, response_content, model_name_used