harness_zbt/backend/app/api/v1/agent.py

from typing import Annotated, Any

from fastapi import APIRouter, Depends, Request, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel, Field
from starlette.responses import StreamingResponse

from app.api.responses import success_response
from app.core.errors import AppError
from app.domains.agent.service import AgentThreadService
from app.domains.identity.models import AdminUser, H5User
from app.domains.identity.service import AdminAuthService, H5AuthService

agent_bearer_scheme = HTTPBearer(auto_error=False)


class CreateThreadRequest(BaseModel):
    title: str | None = Field(default=None, max_length=128)


class TextContent(BaseModel):
    type: str = Field(pattern="^text$")
    text: str = Field(min_length=1, max_length=4000)


class SendMessageRequest(BaseModel):
    content: TextContent
    page_context: dict[str, Any] | None = None
    client_message_id: str = Field(min_length=1, max_length=64)


def create_agent_router(
    thread_service: AgentThreadService,
    h5_auth_service: H5AuthService,
    admin_auth_service: AdminAuthService,
) -> APIRouter:
    router = APIRouter(tags=["agent"])

    def current_h5_user(
        credentials: Annotated[
            HTTPAuthorizationCredentials | None,
            Depends(agent_bearer_scheme),
        ],
    ) -> H5User:
        if credentials is None:
            raise AppError("AUTH_REQUIRED", "请先登录", 401)
        return h5_auth_service.authenticate_access(credentials.credentials)

    def current_admin(
        credentials: Annotated[
            HTTPAuthorizationCredentials | None,
            Depends(agent_bearer_scheme),
        ],
    ) -> AdminUser:
        if credentials is None:
            raise AppError("AUTH_REQUIRED", "请先登录后台", 401)
        return admin_auth_service.authenticate_access(credentials.credentials)

    @router.post("/agent/threads", status_code=status.HTTP_201_CREATED)
    def create_customer_thread(
        payload: CreateThreadRequest,
        request: Request,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.create_customer_thread(user, title=payload.title),
        )

    @router.post(
        "/agent/threads/{thread_id}/messages",
        status_code=status.HTTP_202_ACCEPTED,
    )
    def send_customer_message(
        thread_id: str,
        payload: SendMessageRequest,
        request: Request,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.send_customer_message(
                user,
                thread_id=thread_id,
                text=payload.content.text,
            ),
        )

    @router.post("/agent/threads/{thread_id}/messages/stream")
    def stream_customer_message(
        thread_id: str,
        payload: SendMessageRequest,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> StreamingResponse:
        events = thread_service.stream_customer_message(
            user,
            thread_id=thread_id,
            text=payload.content.text,
        )
        return StreamingResponse(
            (event.to_sse() for event in events),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "X-Accel-Buffering": "no",
            },
        )

    @router.get("/agent/threads/{thread_id}/messages")
    def list_customer_messages(
        thread_id: str,
        request: Request,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.list_customer_messages(user, thread_id=thread_id),
        )

    @router.get("/agent/runs/{run_id}")
    def get_customer_run(
        run_id: str,
        request: Request,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.get_customer_run(user, run_id=run_id),
        )

    @router.get("/agent/runs/{run_id}/events")
    def replay_customer_run_events(
        run_id: str,
        user: Annotated[H5User, Depends(current_h5_user)],
    ) -> StreamingResponse:
        events = thread_service.list_customer_run_events(user, run_id=run_id)
        return StreamingResponse(
            (event.to_sse() for event in events),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache"},
        )

    @router.post(
        "/admin/agent/threads",
        status_code=status.HTTP_201_CREATED,
    )
    def create_operation_thread(
        payload: CreateThreadRequest,
        request: Request,
        user: Annotated[AdminUser, Depends(current_admin)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.create_operation_thread(user, title=payload.title),
        )

    @router.post(
        "/admin/agent/threads/{thread_id}/messages",
        status_code=status.HTTP_202_ACCEPTED,
    )
    def send_operation_message(
        thread_id: str,
        payload: SendMessageRequest,
        request: Request,
        user: Annotated[AdminUser, Depends(current_admin)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.send_operation_message(
                user,
                thread_id=thread_id,
                text=payload.content.text,
            ),
        )

    @router.get("/admin/agent/threads/{thread_id}/messages")
    def list_operation_messages(
        thread_id: str,
        request: Request,
        user: Annotated[AdminUser, Depends(current_admin)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.list_operation_messages(user, thread_id=thread_id),
        )

    @router.get("/admin/agent/runs/{run_id}")
    def get_operation_run(
        run_id: str,
        request: Request,
        user: Annotated[AdminUser, Depends(current_admin)],
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.get_operation_run(user, run_id=run_id),
        )

    @router.get("/admin/agent/runs")
    def list_operation_runs(
        request: Request,
        user: Annotated[AdminUser, Depends(current_admin)],
        limit: int = 20,
    ) -> dict[str, Any]:
        return success_response(
            request,
            thread_service.list_operation_runs(user, limit=min(max(limit, 1), 100)),
        )

    return router