ss-tools/backend/src/api/routes/llm.py

# #region LlmRoutes [C:3] [TYPE Module] [SEMANTICS fastapi, llm, api, provider]
# @defgroup Api Module group.
# @BRIEF API routes for LLM provider configuration and management.
# @LAYER API
# @RELATION DEPENDS_ON -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
# @RELATION DEPENDS_ON -> [get_current_user]
# @RELATION DEPENDS_ON -> [get_db]


from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from ...core.database import get_db
from ...core.logger import logger
from ...dependencies import get_current_user as get_current_active_user
from ...models.llm import ValidationPolicy
from ...plugins.llm_analysis.models import LLMProviderConfig, LLMProviderType
from ...schemas.auth import User
from ...services.llm_provider import LLMProviderService, is_masked_or_placeholder, mask_api_key


# #region FetchModelsRequest [C:1] [TYPE Class]
# @BRIEF Pydantic request model for the fetch-models endpoint.
class FetchModelsRequest(BaseModel):
    base_url: str | None = Field(None, description="LLM provider base URL")
    provider_type: str | None = Field(None, description="Provider type (openai, anthropic, etc)")
    api_key: str | None = Field(None, description="Direct API key (takes precedence over provider_id)")
    provider_id: str | None = Field(None, description="Saved provider ID to look up stored API key")


# #endregion FetchModelsRequest

# #region router [C:1] [TYPE Global]
# @BRIEF APIRouter instance for LLM routes.

router = APIRouter(tags=["LLM"])
# #endregion router


# #region _is_valid_runtime_api_key [C:4] [TYPE Function]
# @BRIEF Validate decrypted runtime API key presence/shape.
# @PRE  value can be None.
# @POST  Returns True only for non-placeholder key.
# @RELATION BINDS_TO -> [LlmRoutes]
def _is_valid_runtime_api_key(value: str | None) -> bool:
    key = (value or "").strip()
    if not key:
        return False
    if key in {"********", "EMPTY_OR_NONE"}:
        return False
    return len(key) >= 16


# #endregion _is_valid_runtime_api_key


# #region get_providers [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Retrieve all LLM provider configurations.
# @PRE  User is authenticated.
# @POST  Returns list of LLMProviderConfig.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
@router.get("/providers", response_model=list[LLMProviderConfig])
async def get_providers(
    current_user: User = Depends(get_current_active_user), db: Session = Depends(get_db)
):
    """
    Get all LLM provider configurations.
    """
    logger.reason(
        f"Fetching providers for user: {current_user.username}",
        extra={"src": "llm_routes.get_providers"},
    )
    service = LLMProviderService(db)
    providers = service.get_all_providers()
    return [
        LLMProviderConfig(
            id=p.id,
            provider_type=LLMProviderType(p.provider_type),
            name=p.name,
            base_url=p.base_url,
            api_key=mask_api_key(service.get_decrypted_api_key(p.id)) if p.api_key else "",
            default_model=p.default_model,
            is_active=p.is_active,
            is_multimodal=bool(p.is_multimodal) if p.is_multimodal is not None else False,
            max_images=p.max_images,
            context_window=p.context_window,
            max_output_tokens=p.max_output_tokens,
        )
        for p in providers
    ]


# #endregion get_providers


# #region fetch_models [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Fetch available models from an LLM provider by base_url+provider_type, or by provider_id.
# @PRE  User is authenticated. Either provider_id or base_url+provider_type must be provided.
# @POST  Returns a list of available model IDs.
# @SIDE_EFFECT  Makes HTTP call to LLM provider API; closes DB connection during network wait.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION CALLS -> [LLMClient]
@router.post("/providers/fetch-models")
async def fetch_models(
    payload: FetchModelsRequest,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    from ...plugins.llm_analysis.models import LLMProviderType
    from ...plugins.llm_analysis.service import LLMClient

    base_url = (payload.base_url or "").strip()
    provider_type_str = (payload.provider_type or "").strip()
    api_key = payload.api_key or ""
    provider_id = payload.provider_id or ""

    # Treat masked/placeholder API keys as if no api_key was provided;
    # they are display-only and cannot be used for real API calls.
    if is_masked_or_placeholder(api_key):
        api_key = ""

    # Early validation: at least one of api_key or provider_id is required
    if not api_key and not provider_id:
        raise HTTPException(
            status_code=400,
            detail="api_key or provider_id is required",
        )

    # Resolve provider_id to stored credentials if no direct api_key given
    if not api_key and provider_id:
        service = LLMProviderService(db)
        db_provider = service.get_provider(provider_id)
        if not db_provider:
            raise HTTPException(status_code=404, detail="Provider not found")
        base_url = base_url or db_provider.base_url
        provider_type_str = provider_type_str or db_provider.provider_type
        stored_key = service.get_decrypted_api_key(provider_id)
        if stored_key:
            api_key = stored_key

    if not base_url:
        raise HTTPException(status_code=400, detail="base_url is required")
    if not provider_type_str:
        raise HTTPException(status_code=400, detail="provider_type is required")

    try:
        provider_type = LLMProviderType(provider_type_str)
    except ValueError:
        raise HTTPException(status_code=400, detail=f"Invalid provider_type: {provider_type_str}")

    # Release DB connection before the network call to avoid idle-in-transaction
    # blocking other queries while we wait for the LLM API response.
    db.close()

    client = LLMClient(
        provider_type=provider_type,
        api_key=api_key or "sk-placeholder",
        base_url=base_url,
        default_model="",
    )

    try:
        models = await client.fetch_models()
        return {"models": models}
    except Exception as e:
        logger.error(
            f"[llm_routes.fetch_models] Failed to fetch models: {e}",
            extra={"src": "llm_routes.fetch_models"},
        )
        raise HTTPException(status_code=502, detail=str(e))


# #endregion fetch_models


# #region get_llm_status [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Returns whether LLM runtime is configured for dashboard validation.
# @PRE  User is authenticated.
# @POST  configured=true only when an active provider with valid decrypted key exists.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION CALLS -> [_is_valid_runtime_api_key]
@router.get("/status")
async def get_llm_status(
    current_user: User = Depends(get_current_active_user), db: Session = Depends(get_db)
):
    service = LLMProviderService(db)
    providers = service.get_all_providers()
    active_provider = next((p for p in providers if p.is_active), None)

    if not providers:
        return {
            "configured": False,
            "reason": "no_providers_configured",
            "provider_count": 0,
            "active_provider_count": 0,
        }

    if not active_provider:
        return {
            "configured": False,
            "reason": "no_active_provider",
            "provider_count": len(providers),
            "active_provider_count": 0,
            "providers": [
                {
                    "id": provider.id,
                    "name": provider.name,
                    "provider_type": provider.provider_type,
                    "is_active": bool(provider.is_active),
                }
                for provider in providers
            ],
        }

    api_key = service.get_decrypted_api_key(active_provider.id)
    if not _is_valid_runtime_api_key(api_key):
        return {
            "configured": False,
            "reason": "invalid_api_key",
            "provider_count": len(providers),
            "active_provider_count": len(
                [provider for provider in providers if provider.is_active]
            ),
            "provider_id": active_provider.id,
            "provider_name": active_provider.name,
            "provider_type": active_provider.provider_type,
            "default_model": active_provider.default_model,
        }

    return {
        "configured": True,
        "reason": "ok",
        "provider_count": len(providers),
        "active_provider_count": len(
            [provider for provider in providers if provider.is_active]
        ),
        "provider_id": active_provider.id,
        "provider_name": active_provider.name,
        "provider_type": active_provider.provider_type,
        "default_model": active_provider.default_model,
    }


# #endregion get_llm_status


# #region create_provider [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Create a new LLM provider configuration.
# @PRE  User is authenticated and has admin permissions.
# @POST  Returns the created LLMProviderConfig.
# @SIDE_EFFECT  Creates a new DB row for the LLM provider.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
@router.post(
    "/providers", response_model=LLMProviderConfig, status_code=status.HTTP_201_CREATED
)
async def create_provider(
    config: LLMProviderConfig,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    """
    Create a new LLM provider configuration.
    """
    service = LLMProviderService(db)
    provider = service.create_provider(config)
    return LLMProviderConfig(
        id=provider.id,
        provider_type=LLMProviderType(provider.provider_type),
        name=provider.name,
        base_url=provider.base_url,
        api_key=mask_api_key(config.api_key),
        default_model=provider.default_model,
        is_active=provider.is_active,
        is_multimodal=bool(provider.is_multimodal) if provider.is_multimodal is not None else False,
        max_images=provider.max_images,
        context_window=provider.context_window,
        max_output_tokens=provider.max_output_tokens,
    )


# #endregion create_provider


# #region update_provider [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Update an existing LLM provider configuration.
# @PRE  User is authenticated and has admin permissions.
# @POST  Returns the updated LLMProviderConfig.
# @SIDE_EFFECT  Updates a DB row for the LLM provider.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
@router.put("/providers/{provider_id}", response_model=LLMProviderConfig)
async def update_provider(
    provider_id: str,
    config: LLMProviderConfig,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    """
    Update an existing LLM provider configuration.
    """
    service = LLMProviderService(db)
    provider = service.update_provider(provider_id, config)
    if not provider:
        raise HTTPException(status_code=404, detail="Provider not found")

    return LLMProviderConfig(
        id=provider.id,
        provider_type=LLMProviderType(provider.provider_type),
        name=provider.name,
        base_url=provider.base_url,
        api_key=mask_api_key(service.get_decrypted_api_key(provider.id)) if provider.api_key else "",
        default_model=provider.default_model,
        is_active=provider.is_active,
        is_multimodal=bool(provider.is_multimodal) if provider.is_multimodal is not None else False,
        max_images=provider.max_images,
        context_window=provider.context_window,
        max_output_tokens=provider.max_output_tokens,
    )


# #endregion update_provider


# #region delete_provider [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Delete an LLM provider configuration.
# @PRE  User is authenticated and has admin permissions.
# @POST  Returns success status.
# @SIDE_EFFECT  Deletes a DB row for the LLM provider; checks active validation tasks.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION CALLS -> [ValidationPolicy]
@router.delete("/providers/{provider_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_provider(
    provider_id: str,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    """
    Delete an LLM provider configuration.
    """
    # Check if any active validation tasks reference this provider
    active_tasks = db.query(ValidationPolicy).filter(
        ValidationPolicy.provider_id == provider_id,
        ValidationPolicy.is_active == True,
    ).all()
    if active_tasks:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail={
                "error": f"Provider is used by {len(active_tasks)} active validation task(s)",
                "blocking_tasks": [{"id": t.id, "name": t.name} for t in active_tasks],
            },
        )

    service = LLMProviderService(db)
    if not service.delete_provider(provider_id):
        raise HTTPException(status_code=404, detail="Provider not found")
    return


# #endregion delete_provider


# #region test_connection [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Test connection to an LLM provider.
# @PRE  User is authenticated.
# @POST  Returns success status and message.
# @SIDE_EFFECT  Makes HTTP call to LLM provider API; decrypts stored API key.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [LLMClient]
@router.post("/providers/{provider_id}/test")
async def test_connection(
    provider_id: str,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    logger.reason(
        f"Testing connection for provider_id: {provider_id}",
        extra={"src": "llm_routes.test_connection"},
    )
    """
    Test connection to an LLM provider.
    """
    from ...plugins.llm_analysis.service import LLMClient

    service = LLMProviderService(db)
    db_provider = service.get_provider(provider_id)
    if not db_provider:
        raise HTTPException(status_code=404, detail="Provider not found")

    api_key = service.get_decrypted_api_key(provider_id)

    # Check if API key was successfully decrypted
    if not api_key:
        logger.error(
            f"[llm_routes][test_connection] Failed to decrypt API key for provider {provider_id}"
        )
        raise HTTPException(
            status_code=500,
            detail="Failed to decrypt API key. The provider may have been encrypted with a different encryption key. Please update the provider with a new API key.",
        )

    client = LLMClient(
        provider_type=LLMProviderType(db_provider.provider_type),
        api_key=api_key,
        base_url=db_provider.base_url,
        default_model=db_provider.default_model,
    )

    try:
        await client.test_runtime_connection()
        return {"success": True, "message": "Connection successful"}
    except Exception as e:
        return {"success": False, "error": str(e)}


# #endregion test_connection


# #region test_provider_config [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Test connection with a provided configuration (not yet saved).
# @PRE  User is authenticated.
# @POST  Returns success status and message.
# @SIDE_EFFECT  Makes HTTP call to LLM provider API using provided config.
# @RELATION DEPENDS_ON -> [LLMClient]
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
@router.post("/providers/test")
async def test_provider_config(
    config: LLMProviderConfig, current_user: User = Depends(get_current_active_user)
):
    """
    Test connection with a provided configuration.
    """
    from ...plugins.llm_analysis.service import LLMClient

    logger.reason(
        f"Testing config for {config.name}",
        extra={"src": "llm_routes.test_provider_config"},
    )

    # Check if API key is provided (reject empty, placeholder, or masked keys)
    if is_masked_or_placeholder(config.api_key):
        raise HTTPException(
            status_code=400, detail="API key is required for testing connection"
        )

    client = LLMClient(
        provider_type=config.provider_type,
        api_key=config.api_key,
        base_url=config.base_url,
        default_model=config.default_model,
    )

    try:
        await client.test_runtime_connection()
        return {"success": True, "message": "Connection successful"}
    except Exception as e:
        return {"success": False, "error": str(e)}


# #endregion test_provider_config


# #region ProbeMaxImagesResponse [C:1] [TYPE Class]
# @BRIEF Response model for probe-max-images endpoint.
class ProbeMaxImagesResponse(BaseModel):
    max_images: int | None
    method: str = "binary_search"


# #endregion ProbeMaxImagesResponse


# #region probe_max_images [C:4] [TYPE Function]
# @ingroup Api
# @BRIEF Probe an LLM provider to discover its max images per request limit.
# @PRE  Provider exists and has valid API key.
# @POST  Returns the detected max_images limit and caches it on the provider.
# @SIDE_EFFECT  Makes multiple HTTP calls to LLM provider API; updates provider max_images in DB.
# @RELATION CALLS -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [EXT:Library:openai]
@router.post("/providers/{provider_id}/probe-max-images", response_model=ProbeMaxImagesResponse)
async def probe_max_images(
    provider_id: str,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db),
):
    """
    Probe an LLM provider to discover the maximum number of images
    allowed per request. Uses binary search with a minimal 1x1 JPEG
    to find the limit without consuming significant tokens.

    The result is cached on the provider's max_images field.
    """
    from openai import AsyncOpenAI

    # Minimal 1x1 white JPEG in base64 (~840 chars, PIL-generated)
    # NOTE: Uses raw AsyncOpenAI client — works for OpenAI/OpenRouter providers.
    # Kilo API gateway does NOT support the OpenAI image content format;
    # probes against Kilo providers will return max_images=0 (not a bug).
    PROBE_IMAGE_B64 = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAFA3PEY8MlBGQUZaVVBfeMiCeG5uePWvuZHI////////////////////////////////////////////////////2wBDAVVaWnhpeOuCguv/////////////////////////////////////////////////////////////////////////wAARCAABAAEDASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwC7RRRQB//Z"

    service = LLMProviderService(db)
    db_provider = service.get_provider(provider_id)
    if not db_provider:
        raise HTTPException(status_code=404, detail="Provider not found")

    api_key = service.get_decrypted_api_key(provider_id)
    if not api_key:
        raise HTTPException(
            status_code=400,
            detail="Provider has no valid API key. Decryption failed or key is missing.",
        )

    model = db_provider.default_model
    if not model:
        raise HTTPException(status_code=400, detail="Provider has no default model set")

    # Build the client
    client = AsyncOpenAI(
        api_key=api_key,
        base_url=db_provider.base_url,
    )

    def build_content(n_images: int) -> list[dict]:
        """Build a minimal content array with N probe images."""
        content: list[dict] = [{"type": "text", "text": "OK"}]
        for _ in range(n_images):
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{PROBE_IMAGE_B64}"},
            })
        return content

    async def try_n(n: int):
        """Try sending N images. Returns True if OK, integer limit if parsed from error, or False if unknown failure."""
        try:
            await client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": build_content(n)}],
                max_tokens=1,
            )
            logger.reason(f"[probe_max_images] {n} images: OK", extra={"src": "probe_max_images"})
            return True
        except Exception as e:
            msg = str(e)
            logger.reason(
                f"[probe_max_images] {n} images: FAILED — {msg[:200]}",
                extra={"src": "probe_max_images"},
            )
            # Try to parse "At most 8 image(s)" from error
            import re
            match = re.search(r"At most (\d+) image", msg)
            if match:
                return int(match.group(1))
            return False

    # Phase 1: Exponential growth to find upper bound
    last_ok = 0
    first_fail = None
    limit_parsed = None

    for n in [1, 2, 4, 8, 16, 32, 64]:
        result = await try_n(n)
        if result is True:
            last_ok = n
        elif type(result) is int:
            # Parsed exact limit from error message
            limit_parsed = result
            first_fail = n
            break
        else:
            first_fail = n
            break

    if limit_parsed is not None:
        detected_limit = limit_parsed
        method = "parse_error"
    elif first_fail is None:
        # 64 still OK — no practical limit detected
        detected_limit = None
        method = "no_limit_detected"
    elif last_ok == 0:
        # Even 1 image failed — provider may not support multimodal
        # Set a minimal safe limit (0 = don't use images) but log the issue
        logger.warning(
            f"[probe_max_images] Even 1 image failed for provider {provider_id}. "
            f"Provider may not support multimodal input or credentials are invalid."
        )
        detected_limit = 0
        method = "binary_search"
    else:
        # Phase 2: Binary search between last_ok and first_fail
        lo, hi = last_ok, first_fail
        while lo < hi:
            mid = (lo + hi + 1) // 2
            result = await try_n(mid)
            if result is True:
                lo = mid
            else:
                hi = mid - 1
        detected_limit = lo
        method = "binary_search"

    # Store the result
    service.set_max_images(provider_id, detected_limit)

    return ProbeMaxImagesResponse(
        max_images=detected_limit,
        method=method,
    )


# #endregion probe_max_images

# #endregion LlmRoutes