free-claude-code/api/routes.py

"""FastAPI route handlers."""

import logging
import uuid

from fastapi import APIRouter, Request, Depends, HTTPException
from fastapi.responses import StreamingResponse

from .models import (
    MessagesRequest,
    MessagesResponse,
    TokenCountRequest,
    TokenCountResponse,
    Usage,
)
from .dependencies import get_provider, get_settings
from .request_utils import (
    is_quota_check_request,
    is_title_generation_request,
    is_prefix_detection_request,
    extract_command_prefix,
    get_token_count,
)
from config.settings import Settings
from providers.nvidia_nim import NvidiaNimProvider
from providers.exceptions import ProviderError
from providers.logging_utils import log_request_compact

logger = logging.getLogger(__name__)

router = APIRouter()


# =============================================================================
# Routes
# =============================================================================


@router.post("/v1/messages")
async def create_message(
    request_data: MessagesRequest,
    raw_request: Request,
    provider: NvidiaNimProvider = Depends(get_provider),
    settings: Settings = Depends(get_settings),
):
    """Create a message (streaming or non-streaming)."""

    try:
        if settings.fast_prefix_detection:
            is_prefix_req, command = is_prefix_detection_request(request_data)
            if is_prefix_req:
                return MessagesResponse(
                    id=f"msg_{uuid.uuid4()}",
                    model=request_data.model,
                    content=[{"type": "text", "text": extract_command_prefix(command)}],
                    stop_reason="end_turn",
                    usage=Usage(input_tokens=100, output_tokens=5),
                )

        # Optimization: Mock network probe/quota requests
        if settings.enable_network_probe_mock and is_quota_check_request(request_data):
            logger.info("Optimization: Intercepted and mocked quota probe")
            return MessagesResponse(
                id=f"msg_{uuid.uuid4()}",
                model=request_data.model,
                role="assistant",
                content=[{"type": "text", "text": "Quota check passed."}],
                stop_reason="end_turn",
                usage=Usage(input_tokens=10, output_tokens=5),
            )

        # Optimization: Skip title generation requests
        if settings.enable_title_generation_skip and is_title_generation_request(
            request_data
        ):
            logger.info("Optimization: Skipped title generation request")
            return MessagesResponse(
                id=f"msg_{uuid.uuid4()}",
                model=request_data.model,
                role="assistant",
                content=[{"type": "text", "text": "Conversation"}],
                stop_reason="end_turn",
                usage=Usage(input_tokens=100, output_tokens=5),
            )

        request_id = f"req_{uuid.uuid4().hex[:12]}"
        log_request_compact(logger, request_id, request_data)

        if request_data.stream:
            input_tokens = get_token_count(
                request_data.messages, request_data.system, request_data.tools
            )
            return StreamingResponse(
                provider.stream_response(request_data, input_tokens=input_tokens),
                media_type="text/event-stream",
                headers={
                    "X-Accel-Buffering": "no",
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                },
            )
        else:
            response_json = await provider.complete(request_data)
            return provider.convert_response(response_json, request_data)

    except ProviderError:
        raise
    except Exception as e:
        import traceback

        logger.error(f"Error: {str(e)}\n{traceback.format_exc()}")
        raise HTTPException(status_code=getattr(e, "status_code", 500), detail=str(e))


@router.post("/v1/messages/count_tokens")
async def count_tokens(request_data: TokenCountRequest):
    """Count tokens for a request."""
    try:
        return TokenCountResponse(
            input_tokens=get_token_count(
                request_data.messages, request_data.system, request_data.tools
            )
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/")
async def root(settings: Settings = Depends(get_settings)):
    """Root endpoint."""
    return {
        "status": "ok",
        "provider": "nvidia_nim",
        "model": settings.model,
    }


@router.get("/health")
async def health():
    """Health check endpoint."""
    return {"status": "healthy"}


@router.post("/stop")
async def stop_cli(request: Request):
    """Stop all CLI sessions and pending tasks."""
    handler = getattr(request.app.state, "message_handler", None)
    if not handler:
        # Fallback if messaging not initialized
        cli_manager = getattr(request.app.state, "cli_manager", None)
        if cli_manager:
            await cli_manager.stop_all()
            return {"status": "stopped", "source": "cli_manager"}
        return HTTPException(status_code=503, detail="Messaging system not initialized")

    count = await handler.stop_all_tasks()
    return {"status": "stopped", "cancelled_count": count}