JARVIS/backend/app/services/agent_service.py

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

import json
import uuid
from datetime import datetime
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from langchain_core.messages import HumanMessage, AIMessage

from app.models.conversation import Conversation, Message
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


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

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

    async def chat(
        self,
        user_id: str,
        message: str,
        conversation_id: str | None = None,
    ) -> tuple[str, str, AsyncGenerator[str, None]]:
        """
        处理对话请求（流式）

        Returns:
            (conversation_id, message_id, response_stream)
        """
        # 获取或创建对话
        if conversation_id:
            result = await self.db.execute(
                select(Conversation).where(Conversation.id == conversation_id)
            )
            conv = result.scalar_one_or_none()
        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)
        await self.db.commit()
        await self.db.refresh(user_msg)

        # 预创建助手消息（后续更新内容）
        assistant_msg = Message(
            conversation_id=conversation_id,
            role="assistant",
            content="",
            model="jarvis",
        )
        self.db.add(assistant_msg)
        await self.db.commit()
        await self.db.refresh(assistant_msg)

        # 加载记忆上下文
        memory_ctx = await memory_service.build_memory_context(
            self.db, user_id, conversation_id, message
        )

        # 调用 LangGraph Agent
        async def run_agent():
            set_current_user(user_id)
            try:
                graph = get_agent_graph()
                langgraph_state = {
                    "messages": [HumanMessage(content=message)],  # type: ignore[arg-type]
                    "user_id": user_id,
                    "conversation_id": conversation_id,
                    "current_agent": "master",
                    "active_agents": ["master"],
                    "pending_tasks": [],
                    "completed_tasks": [],
                    "tool_calls": [],
                    "last_tool_result": None,
                    "knowledge_context": None,
                    "graph_context": None,
                    "plan": None,
                    "plan_steps": [],
                    "analysis_report": None,
                    "final_response": None,
                    "should_respond": True,
                    "memory_context": memory_ctx,
                }

                collected = ""
                async for event in graph.astream_events(langgraph_state, version="v2"):
                    kind = event.get("event")
                    if kind == "on_chat_model_end":
                        content = event.get("data", {}).get("output", {})
                        if isinstance(content, dict):
                            content = content.get("content", "")
                        if content:
                            delta = content[len(collected):]
                            if delta:
                                collected += delta
                                yield delta
                    elif kind == "on_tool_end":
                        name = event.get("name", "")
                        yield f"\n[工具执行: {name}]\n"
            except Exception as e:
                yield f"\n执行出错: {str(e)}"
            finally:
                clear_current_user()
                # 异步触发自动摘要和记忆提取（不阻塞响应）
                import asyncio
                try:
                    loop = asyncio.get_running_loop()
                    loop.create_task(
                        memory_service.try_auto_summarize(self.db, user_id, conversation_id)
                    )
                except Exception:
                    pass

            # 最终更新数据库中的消息内容
            if collected:
                try:
                    result2 = await self.db.execute(
                        select(Message).where(Message.id == assistant_msg.id)
                    )
                    msg = result2.scalar_one_or_none()
                    if msg:
                        msg.content = collected
                        await self.db.commit()
                except Exception:
                    pass

        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,
    ) -> tuple[str, str, str]:
        """
        简单同步版对话（无流式）

        Returns:
            (conversation_id, message_id, response_content)
        """
        # 获取或创建对话
        if conversation_id:
            result = await self.db.execute(
                select(Conversation).where(Conversation.id == conversation_id)
            )
            conv = result.scalar_one_or_none()
        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)

        # 加载记忆上下文
        memory_ctx = await memory_service.build_memory_context(
            self.db, user_id, conversation_id, message
        )

        # 调用 LangGraph Agent
        set_current_user(user_id)
        graph = get_agent_graph()
        langgraph_state = {
            "messages": [HumanMessage(content=full_message)],  # type: ignore[arg-type]
            "user_id": user_id,
            "conversation_id": conversation_id,
            "current_agent": "master",
            "active_agents": ["master"],
            "pending_tasks": [],
            "completed_tasks": [],
            "tool_calls": [],
            "last_tool_result": None,
            "knowledge_context": None,
            "graph_context": None,
            "plan": None,
            "plan_steps": [],
            "analysis_report": None,
            "final_response": None,
            "should_respond": True,
            "memory_context": memory_ctx,
        }

        try:
            result_state = await graph.ainvoke(langgraph_state)
            response_content = result_state.get("final_response", "抱歉，我无法处理这个请求。")
        except Exception as e:
            response_content = f"抱歉，发生错误: {str(e)}"
        finally:
            clear_current_user()
            # 异步触发自动摘要
            import asyncio
            try:
                asyncio.get_running_loop().create_task(
                    memory_service.try_auto_summarize(self.db, user_id, conversation_id)
                )
            except Exception:
                pass

        # 保存助手消息
        assistant_msg = Message(
            conversation_id=conversation_id,
            role="assistant",
            content=response_content,
            model="jarvis",
        )
        self.db.add(assistant_msg)
        await self.db.commit()
        await self.db.refresh(assistant_msg)

        return conversation_id, assistant_msg.id, response_content