auto_agent/backend/app/api/demo/chat.py

from __future__ import annotations

from typing import Annotated
from uuid import uuid4

from fastapi import APIRouter, Depends, Header, HTTPException, status
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.core.constants import ERROR_CODE_OK
from app.core.time import format_now
from app.db.session import get_db
from app.schemas.chat import (
    ChatConfirmTaskData,
    ChatConfirmTaskRequest,
    ChatSendMessageData,
    ChatSendMessageRequest,
    ChatSessionCreateRequest,
    ChatSessionData,
)
from app.schemas.common import ApiResponse
from app.services.chat_service import ChatService, ChatSessionNotFoundError

router = APIRouter(prefix="/api/demo/chat", tags=["demo-chat"])


def build_request_id(header_value: str | None) -> str:
    return header_value or f"req-{uuid4().hex[:12]}"


@router.post("/sessions", response_model=ApiResponse[ChatSessionData])
def create_session(
    payload: ChatSessionCreateRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[ChatSessionData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    service = ChatService(db, settings.default_timezone)
    session = service.create_session(payload.tenant_id, payload.channel)
    messages = [service.to_message_item(item) for item in service.list_messages(session.session_id)]
    return ApiResponse[ChatSessionData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=ChatSessionData(
            session_id=session.session_id,
            tenant_id=session.tenant_id,
            channel=session.channel,
            title=session.title,
            last_task_id=session.last_task_id,
            sample_prompts=ChatService.SAMPLE_PROMPTS,
            messages=messages,
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.get("/sessions/{session_id}", response_model=ApiResponse[ChatSessionData])
def get_session(
    session_id: str,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[ChatSessionData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    service = ChatService(db, settings.default_timezone)
    try:
        session = service.get_session(session_id)
    except ChatSessionNotFoundError as exc:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail={"code": "NOT_FOUND", "message": "chat session not found"}) from exc
    messages = [service.to_message_item(item) for item in service.list_messages(session.session_id)]
    return ApiResponse[ChatSessionData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=ChatSessionData(
            session_id=session.session_id,
            tenant_id=session.tenant_id,
            channel=session.channel,
            title=session.title,
            last_task_id=session.last_task_id,
            sample_prompts=ChatService.SAMPLE_PROMPTS,
            messages=messages,
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.post("/sessions/{session_id}/messages", response_model=ApiResponse[ChatSendMessageData])
def send_message(
    session_id: str,
    payload: ChatSendMessageRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[ChatSendMessageData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    service = ChatService(db, settings.default_timezone)
    try:
        _, assistant_message, task_data = service.handle_user_message(session_id, payload.content, payload.context)
    except ChatSessionNotFoundError as exc:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail={"code": "NOT_FOUND", "message": "chat session not found"}) from exc

    return ApiResponse[ChatSendMessageData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=ChatSendMessageData(
            session_id=session_id,
            task_id=task_data["task_id"],
            task_status=task_data["task_status"],
            parsed_intent=task_data["parsed_intent"],
            missing_slots=task_data["missing_slots"],
            risk_level=task_data["risk_level"],
            next_action=task_data["next_action"],
            assistant_message=service.to_message_item(assistant_message),
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.post("/sessions/{session_id}/tasks/{task_id}/confirm", response_model=ApiResponse[ChatConfirmTaskData])
def confirm_task_from_chat(
    session_id: str,
    task_id: str,
    payload: ChatConfirmTaskRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[ChatConfirmTaskData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    service = ChatService(db, settings.default_timezone)
    try:
        _, assistant_message, task_data = service.confirm_task(session_id, task_id, payload.comment)
    except ChatSessionNotFoundError as exc:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail={"code": "NOT_FOUND", "message": "chat session not found"}) from exc

    return ApiResponse[ChatConfirmTaskData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=ChatConfirmTaskData(
            session_id=session_id,
            task_id=task_data["task_id"],
            task_status=task_data["task_status"],
            approval_status=task_data["approval_status"],
            assistant_message=service.to_message_item(assistant_message),
        ),
        timestamp=format_now(settings.default_timezone),
    )