app/api/v1/token_counter.py

from __future__ import annotations

import time

from fastapi import APIRouter, Depends

from app.api.deps import require_auth
from app.models.domain import count_tokens
from app.models.schemas import TokenCountRequest, TokenCountResponse

router = APIRouter()


@router.post(
    "/token/count",
    response_model=TokenCountResponse,
    summary="Count tokens in text using tiktoken",
)
async def count_text_tokens(
    body: TokenCountRequest,
    token: str = Depends(require_auth),
) -> TokenCountResponse:
    start = time.perf_counter()
    try:
        token_count = count_tokens(body.text, body.encoding)
        elapsed_ms = round((time.perf_counter() - start) * 1000, 3)
        return TokenCountResponse(
            success=True,
            time_ms=elapsed_ms,
            token_count=token_count,
            char_count=len(body.text),
            encoding=body.encoding,
        )
    except Exception as exc:
        elapsed_ms = round((time.perf_counter() - start) * 1000, 3)
        return TokenCountResponse(
            success=False,
            time_ms=elapsed_ms,
            error=str(exc),
        )