KiloStar/kilostar/api/chat.py

#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from fastapi import APIRouter, Depends
from pydantic import BaseModel
from kilostar.utils.ray_hook import ray_actor_hook
from kilostar.utils.access import Accessor, TokenData

chat_router = APIRouter(prefix="/api/v1/chat", tags=["chat"])


class CreateChatRequest(BaseModel):
    title: str = "新对话"
    initial_message: str


class SendMessageRequest(BaseModel):
    message: str


@chat_router.post("")
async def create_chat_session(
    request: CreateChatRequest,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    chat = await postgres_database.create_chat_session.remote(
        user_id=token_data.user_id, title=request.title
    )

    # 存入用户消息
    await postgres_database.add_chat_message.remote(
        chat_id=chat.chat_id, message=request.initial_message, message_owner="user"
    )

    # 调用监管节点处理简单任务/交流
    regulatory_node = ray_actor_hook("regulatory_node").regulatory_node
    # 在此发起任务并等待或异步返回结果
    response_msg = await regulatory_node.handle_chat_message.remote(
        user_id=token_data.user_id,
        chat_id=chat.chat_id,
        message=request.initial_message,
    )

    # 存入回复消息
    if response_msg:
        await postgres_database.add_chat_message.remote(
            chat_id=chat.chat_id, message=response_msg, message_owner="regulatory_node"
        )

    return {"chat_id": chat.chat_id, "reply": response_msg}


@chat_router.get("")
async def list_chat_sessions(
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    sessions = await postgres_database.list_chat_sessions.remote(
        user_id=token_data.user_id
    )
    return {"sessions": sessions}


@chat_router.get("/{chat_id}")
async def get_chat_history(
    chat_id: str, token_data: TokenData = Depends(Accessor.get_current_user)
):
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    messages = await postgres_database.list_chat_messages.remote(chat_id=chat_id)
    return {"messages": messages}


@chat_router.post("/{chat_id}/reply")
async def send_chat_message(
    chat_id: str,
    request: SendMessageRequest,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    # 存用户消息
    await postgres_database.add_chat_message.remote(
        chat_id=chat_id, message=request.message, message_owner="user"
    )

    # 调用监管节点
    regulatory_node = ray_actor_hook("regulatory_node").regulatory_node
    response_msg = await regulatory_node.handle_chat_message.remote(
        user_id=token_data.user_id, chat_id=chat_id, message=request.message
    )

    # 存回复
    if response_msg:
        await postgres_database.add_chat_message.remote(
            chat_id=chat_id, message=response_msg, message_owner="regulatory_node"
        )

    return {"reply": response_msg}