mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-06-30 16:40:05 +08:00
Andreas Köpf
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from fastapi import APIRouter, Depends
|
|
from oasst_backend.api import deps
|
|
from oasst_backend.models import ApiClient
|
|
from oasst_backend.prompt_repository import PromptRepository
|
|
from oasst_backend.tree_manager import TreeManager, TreeManagerStats, TreeMessageCountStats
|
|
from oasst_shared.schemas import protocol
|
|
from sqlmodel import Session
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.get("/", response_model=protocol.SystemStats)
|
|
def get_message_stats(
|
|
db: Session = Depends(deps.get_db),
|
|
api_client: ApiClient = Depends(deps.get_trusted_api_client),
|
|
):
|
|
pr = PromptRepository(db, api_client)
|
|
return pr.get_stats()
|
|
|
|
|
|
@router.get("/tree_manager/state_counts", response_model=dict[str, int])
|
|
def get_tree_manager__state_counts(
|
|
db: Session = Depends(deps.get_db),
|
|
api_client: ApiClient = Depends(deps.get_trusted_api_client),
|
|
):
|
|
pr = PromptRepository(db, api_client)
|
|
tm = TreeManager(db, pr)
|
|
return tm.tree_counts_by_state()
|
|
|
|
|
|
@router.get("/tree_manager/message_counts", response_model=list[TreeMessageCountStats])
|
|
def get_tree_manager__message_counts(
|
|
only_active: bool = True,
|
|
db: Session = Depends(deps.get_db),
|
|
api_client: ApiClient = Depends(deps.get_trusted_api_client),
|
|
):
|
|
pr = PromptRepository(db, api_client)
|
|
tm = TreeManager(db, pr)
|
|
return tm.tree_message_count_stats(only_active=only_active)
|
|
|
|
|
|
@router.get("/tree_manager", response_model=TreeManagerStats)
|
|
def get_tree_manager__stats(
|
|
db: Session = Depends(deps.get_db),
|
|
api_client: ApiClient = Depends(deps.get_trusted_api_client),
|
|
):
|
|
pr = PromptRepository(db, api_client)
|
|
tm = TreeManager(db, pr)
|
|
return tm.stats()
|