DEMO-AGENT/review_agent/services.py

from __future__ import annotations

import json

from django.db.models import Q, QuerySet
from django.conf import settings
from django.utils import timezone

from .file_summary.workflow import create_file_summary_batch, start_file_summary_workflow
from .file_summary.workflow_trigger import evaluate_file_summary_trigger
from .llm import LLMConfigurationError, LLMRequestError, generate_reply, stream_reply
from .models import Conversation, Message


def list_conversations(user, search: str = "") -> QuerySet[Conversation]:
    """Returns a user's conversations, optionally filtered by title or content."""

    conversations = Conversation.objects.filter(user=user)
    if not search:
        return conversations

    return conversations.filter(
        Q(title__icontains=search) | Q(messages__content__icontains=search)
    ).distinct()


def get_conversation_for_user(user, conversation_id: int | None) -> Conversation | None:
    """Loads a conversation only when it belongs to the current user."""

    if not conversation_id:
        return None
    return Conversation.objects.filter(user=user, pk=conversation_id).first()


def create_conversation(user) -> Conversation:
    """Creates an empty conversation that can immediately accept messages."""

    now = timezone.localtime()
    return Conversation.objects.create(
        user=user,
        title=f"新对话 {now.strftime('%m-%d %H:%M')}",
    )


def append_user_message(conversation: Conversation, content: str) -> Message:
    """Appends a user message and updates the conversation title from the first prompt."""

    message = Message.objects.create(
        conversation=conversation,
        role=Message.Role.USER,
        content=content.strip(),
    )

    if conversation.messages.filter(role=Message.Role.USER).count() == 1:
        conversation.title = build_conversation_title(content)
        conversation.save(update_fields=["title", "updated_at"])

    return message


def append_assistant_message(conversation: Conversation, content: str) -> Message:
    """Appends the deterministic assistant reply."""

    return Message.objects.create(
        conversation=conversation,
        role=Message.Role.ASSISTANT,
        content=content,
    )


def send_message(conversation: Conversation, content: str) -> tuple[Message, Message]:
    """Stores one user message and one provider-backed assistant reply."""

    user_message = append_user_message(conversation, content)
    try:
        reply_content = generate_reply(conversation, content)
    except (LLMConfigurationError, LLMRequestError) as exc:
        reply_content = f"模型调用失败：{exc}"

    assistant_message = append_assistant_message(conversation, reply_content)

    if conversation.title.startswith("新对话"):
        conversation.title = build_conversation_title(content)
        conversation.save(update_fields=["title", "updated_at"])

    return user_message, assistant_message


def stream_message(conversation: Conversation, content: str):
    """Yields SSE events while collecting a streamed assistant reply."""

    user_message = append_user_message(conversation, content)
    assistant_parts: list[str] = []
    trigger = evaluate_file_summary_trigger(conversation, content)

    yield sse_event(
        "meta",
        {
            "conversation_id": conversation.pk,
            "title": conversation.title or build_conversation_title(content),
            "user_message_id": user_message.pk,
            "user_message": user_message.content,
        },
    )

    if trigger.reason == "missing_attachment":
        reply_content = "请先在当前对话右侧上传需要汇总的文件或压缩包，然后再发送自动汇总指令。"
        assistant_message = append_assistant_message(conversation, reply_content)
        yield sse_event("chunk", {"delta": reply_content})
        yield sse_event(
            "done",
            {
                "assistant_message_id": assistant_message.pk,
                "conversation_id": conversation.pk,
                "title": conversation.title,
            },
        )
        return

    if trigger.should_start:
        batch = create_file_summary_batch(
            conversation=conversation,
            user=conversation.user,
            trigger_message=user_message,
        )
        start_file_summary_workflow(
            batch,
            async_run=getattr(settings, "FILE_SUMMARY_ASYNC", True),
        )
        reply_content = f"已启动文件目录与页数自动汇总工作流，批次号：{batch.batch_no}。"
        assistant_message = append_assistant_message(conversation, reply_content)
        yield sse_event(
            "workflow_started",
            {
                "workflow_type": "file_summary",
                "batch_id": batch.pk,
                "batch_no": batch.batch_no,
            },
        )
        yield sse_event("chunk", {"delta": reply_content})
        yield sse_event(
            "done",
            {
                "assistant_message_id": assistant_message.pk,
                "conversation_id": conversation.pk,
                "title": conversation.title,
            },
        )
        return

    try:
        for chunk in stream_reply(conversation, content):
            assistant_parts.append(chunk)
            yield sse_event("chunk", {"delta": chunk})
    except (LLMConfigurationError, LLMRequestError) as exc:
        fallback = f"模型调用失败：{exc}"
        assistant_parts = [fallback]
        yield sse_event("error", {"message": fallback})

    assistant_message = append_assistant_message(conversation, "".join(assistant_parts).strip())

    if conversation.title.startswith("新对话"):
        conversation.title = build_conversation_title(content)
        conversation.save(update_fields=["title", "updated_at"])

    yield sse_event(
        "done",
        {
            "assistant_message_id": assistant_message.pk,
            "conversation_id": conversation.pk,
            "title": conversation.title,
        },
    )


def build_conversation_title(content: str) -> str:
    """Creates a concise title from the first user message."""

    normalized = " ".join(content.strip().split())
    if not normalized:
        return "新对话"
    return normalized[:24]


def sse_event(event_name: str, payload: dict[str, object]) -> str:
    """Formats one server-sent event frame."""

    return f"event: {event_name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"