mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-16 01:10:06 +08:00
208 lines
7.2 KiB
Python
208 lines
7.2 KiB
Python
from datetime import datetime
|
|
from typing import Optional
|
|
from uuid import UUID
|
|
|
|
from fastapi import APIRouter, Depends, Query
|
|
from oasst_backend.api import deps
|
|
from oasst_backend.api.v1 import utils
|
|
from oasst_backend.models import ApiClient
|
|
from oasst_backend.prompt_repository import PromptRepository
|
|
from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode
|
|
from oasst_shared.schemas import protocol
|
|
from sqlmodel import Session
|
|
from starlette.status import HTTP_204_NO_CONTENT
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.get("/", response_model=list[protocol.Message])
|
|
def query_messages(
|
|
auth_method: Optional[str] = None,
|
|
username: Optional[str] = None,
|
|
api_client_id: Optional[str] = None,
|
|
max_count: Optional[int] = Query(10, gt=0, le=1000),
|
|
start_date: Optional[datetime] = None,
|
|
end_date: Optional[datetime] = None,
|
|
only_roots: Optional[bool] = False,
|
|
desc: Optional[bool] = True,
|
|
allow_deleted: Optional[bool] = False,
|
|
api_client: ApiClient = Depends(deps.get_api_client),
|
|
db: Session = Depends(deps.get_db),
|
|
):
|
|
"""
|
|
Query messages.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
messages = pr.query_messages_ordered_by_created_date(
|
|
auth_method=auth_method,
|
|
username=username,
|
|
api_client_id=api_client_id,
|
|
desc=desc,
|
|
limit=max_count,
|
|
gte_created_date=start_date,
|
|
lte_created_date=end_date,
|
|
only_roots=only_roots,
|
|
deleted=None if allow_deleted else False,
|
|
)
|
|
|
|
return utils.prepare_message_list(messages)
|
|
|
|
|
|
@router.get("/cursor", response_model=protocol.MessagePage)
|
|
def get_messages_cursor(
|
|
lt: Optional[str] = None,
|
|
gt: Optional[str] = None,
|
|
user_id: Optional[UUID] = None,
|
|
auth_method: Optional[str] = None,
|
|
username: Optional[str] = None,
|
|
api_client_id: Optional[str] = None,
|
|
only_roots: Optional[bool] = False,
|
|
include_deleted: Optional[bool] = False,
|
|
max_count: Optional[int] = Query(10, gt=0, le=1000),
|
|
desc: Optional[bool] = False,
|
|
api_client: ApiClient = Depends(deps.get_api_client),
|
|
db: Session = Depends(deps.get_db),
|
|
):
|
|
def split_cursor(x: str | None) -> tuple[datetime, UUID]:
|
|
if not x:
|
|
return None, None
|
|
try:
|
|
m = utils.split_uuid_pattern.match(x)
|
|
if m:
|
|
return datetime.fromisoformat(m[2]), UUID(m[1])
|
|
return datetime.fromisoformat(x), None
|
|
except ValueError:
|
|
raise OasstError("Invalid cursor value", OasstErrorCode.INVALID_CURSOR_VALUE)
|
|
|
|
lte_created_date, lt_id = split_cursor(lt)
|
|
gte_created_date, gt_id = split_cursor(gt)
|
|
|
|
pr = PromptRepository(db, api_client)
|
|
messages = pr.query_messages_ordered_by_created_date(
|
|
user_id=user_id,
|
|
auth_method=auth_method,
|
|
username=username,
|
|
api_client_id=api_client_id,
|
|
gte_created_date=gte_created_date,
|
|
gt_id=gt_id,
|
|
lte_created_date=lte_created_date,
|
|
lt_id=lt_id,
|
|
only_roots=only_roots,
|
|
deleted=None if include_deleted else False,
|
|
desc=desc,
|
|
limit=max_count,
|
|
)
|
|
|
|
items = utils.prepare_message_list(messages)
|
|
n, p = None, None
|
|
if len(items) > 0:
|
|
if len(items) == max_count or gte_created_date:
|
|
p = str(items[0].id) + "$" + items[0].created_date.isoformat()
|
|
if len(items) == max_count or lte_created_date:
|
|
n = str(items[-1].id) + "$" + items[-1].created_date.isoformat()
|
|
else:
|
|
if gte_created_date:
|
|
p = gte_created_date.isoformat()
|
|
if lte_created_date:
|
|
n = lte_created_date.isoformat()
|
|
|
|
order = "desc" if desc else "asc"
|
|
return protocol.MessagePage(prev=p, next=n, sort_key="created_date", order=order, items=items)
|
|
|
|
|
|
@router.get("/{message_id}", response_model=protocol.Message)
|
|
def get_message(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get a message by its internal ID.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
message = pr.fetch_message(message_id)
|
|
return utils.prepare_message(message)
|
|
|
|
|
|
@router.get("/{message_id}/conversation", response_model=protocol.Conversation)
|
|
def get_conv(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get a conversation from the tree root and up to the message with given internal ID.
|
|
"""
|
|
|
|
pr = PromptRepository(db, api_client)
|
|
messages = pr.fetch_message_conversation(message_id)
|
|
return utils.prepare_conversation(messages)
|
|
|
|
|
|
@router.get("/{message_id}/tree", response_model=protocol.MessageTree)
|
|
def get_tree(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get all messages belonging to the same message tree.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
message = pr.fetch_message(message_id)
|
|
tree = pr.fetch_message_tree(message.message_tree_id, reviewed=False)
|
|
return utils.prepare_tree(tree, message.message_tree_id)
|
|
|
|
|
|
@router.get("/{message_id}/children", response_model=list[protocol.Message])
|
|
def get_children(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get all messages belonging to the same message tree.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
messages = pr.fetch_message_children(message_id)
|
|
return utils.prepare_message_list(messages)
|
|
|
|
|
|
@router.get("/{message_id}/descendants", response_model=protocol.MessageTree)
|
|
def get_descendants(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get a subtree which starts with this message.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
message = pr.fetch_message(message_id)
|
|
descendants = pr.fetch_message_descendants(message)
|
|
return utils.prepare_tree(descendants, message.id)
|
|
|
|
|
|
@router.get("/{message_id}/longest_conversation_in_tree", response_model=protocol.Conversation)
|
|
def get_longest_conv(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get the longest conversation from the tree of the message.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
message = pr.fetch_message(message_id)
|
|
conv = pr.fetch_longest_conversation(message.message_tree_id)
|
|
return utils.prepare_conversation(conv)
|
|
|
|
|
|
@router.get("/{message_id}/max_children_in_tree", response_model=protocol.MessageTree)
|
|
def get_max_children(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
"""
|
|
Get message with the most children from the tree of the provided message.
|
|
"""
|
|
pr = PromptRepository(db, api_client)
|
|
message = pr.fetch_message(message_id)
|
|
message, children = pr.fetch_message_with_max_children(message.message_tree_id)
|
|
return utils.prepare_tree([message, *children], message.id)
|
|
|
|
|
|
@router.delete("/{message_id}", status_code=HTTP_204_NO_CONTENT)
|
|
def mark_message_deleted(
|
|
message_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db)
|
|
):
|
|
pr = PromptRepository(db, api_client)
|
|
pr.mark_messages_deleted(message_id)
|