Llms¶

@router.post(
    "/downloads/{job_id}/cancel",
    status_code=200,
)
def cancel_download(
    job_id: int,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    job_repo: Download_Job_Repository = Depends(get_download_job_repository),
    db: Session = Depends(get_db),
):
    """Cancel an active download job and cleanup partial files."""
    try:
        return cancel_download_job(job_id, job_repo, llm_repo, db)
    except (
        DownloadJobNotFoundException,
        ModelNotFoundException,
        InvalidInputException,
        StateConflictException,
    ):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to cancel download", trace=str(e))

@router.delete("/{llm_id}")
async def delete_llm(
    llm_id: int,
    orphan_dependents: bool = False,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    db: Session = Depends(get_db),
):
    """Delete local LLM (permanent deletion).

    Regular local model: model files are removed from disk along with the
    database record. If the model has dependent KB assistants (rows sharing its
    COPIED ``link``, #209), the delete is guarded: without ``orphan_dependents``
    it raises 409 carrying the dependents payload so the client can offer to
    orphan them and proceed. With ``orphan_dependents=true`` the base is deleted,
    its assistants REMAIN (their link now dangles, re-bindable via /rebind), and
    every conversation bound to the base is nulled server-side (SET NULL, #225).

    Specialized KB assistant: its ``link`` is a COPY of the base model's (set at
    creation) — the files belong to the base model and are left untouched; the
    assistant's KnowledgeBase is deleted instead, and server-side cascades sweep
    its documents and ``rag.kb_chunks``. Its own conversations survive (SET NULL).

    Args:
        llm_id: ID of the LLM to delete.
        orphan_dependents: When true, delete a base model that has dependent KB
            assistants anyway (they become orphans). Ignored for assistants and
            for base models with no dependents.
        llm_repo: Injected LLM repository.
        db: Database session for transaction control.

    Returns:
        dict: Success message.

    Raises:
        ModelNotFoundException: If LLM not found.
        StateConflictException: 400 if currently downloading; 409 (with the
            dependents payload) if a base model has dependents and
            ``orphan_dependents`` was not set.
        DatabaseException: If deletion fails.

    Warning:
        Deletes model files from disk (regular models). Cannot be undone.
    """
    try:
        llm = llm_repo.get_by_id(llm_id)
        if not llm:
            raise ModelNotFoundException(f"LLM {llm_id}")

        if llm.local == 2:
            raise StateConflictException("Cannot delete LLM while downloading")

        if llm.is_attached_to_kb:
            # KB assistant: never touch the (shared) model files. Deleting
            # the KB cascades to its documents/chunks, and the ORM cascade
            # (KnowledgeBase.llm, delete-orphan) removes the assistant row.
            # The assistant's conversations survive server-side (llm_id SET NULL).
            from src.entities.KnowledgeBase import KnowledgeBase

            await _unload_if_resident(llm.id)
            kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == llm.kb_id).first()
            if kb:
                db.delete(kb)
            else:
                # Inconsistent state (flag set, KB gone): drop the row alone.
                llm_repo.delete(llm)
            db.commit()
            logger.info(f"Deleted KB assistant {llm_id} (KB {llm.kb_id}, files kept)")
            return {"message": "LLM deleted successfully"}

        # Base model: guard against silently orphaning KB assistants that share
        # its weights. Without opt-in, surface the dependents (409) so the client
        # can confirm; with orphan_dependents=true, proceed and let them dangle.
        dependents = _dependents_payload(llm_repo, llm)
        if dependents["assistants"] and not orphan_dependents:
            raise StateConflictException(
                f"Base model '{llm.name}' has {len(dependents['assistants'])} "
                "dependent KB assistant(s). Retry with orphan_dependents=true to "
                "delete it anyway; the assistants remain and can be rebound to "
                "another base.",
                status_code=http_status.HTTP_409_CONFLICT,
                detail=dependents,
            )

        # Unload before the weights go, and only now: the 409 above means the
        # delete was refused, and a refused delete must leave the engine alone.
        await _unload_if_resident(llm.id)

        # Delete files from disk if they exist
        if llm.link and os.path.exists(llm.link):
            shutil.rmtree(llm.link, ignore_errors=True)
            logger.info(f"Deleted model files: {llm.link}")

            # Check and delete residual temp files (e.g., temp_36 for llm.link = data/models/36)
            temp_path = config.LLM_DIR / f"temp_{llm.id}"
            if os.path.exists(str(temp_path)):
                shutil.rmtree(str(temp_path), ignore_errors=True)
                logger.warning(f"Deleting residual temp files associated to llm: {temp_path}")

        # Delete database record
        llm_repo.delete(llm)
        db.commit()

        return {"message": "LLM deleted successfully"}

    except (ModelNotFoundException, StateConflictException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to delete LLM", trace=str(e))

@router.post(
    "/download/huggingface",
    response_model=DownloadJobResponse,
    status_code=http_status.HTTP_200_OK,
)
async def download_huggingface_route(
    payload: HFDownloadRequest,
    background_tasks: BackgroundTasks,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    job_repo: Download_Job_Repository = Depends(get_download_job_repository),
    db: Session = Depends(get_db),
):
    """Download a model chosen from HF search by its repo id (not via the catalog).

    Mirrors POST /{llm_id}/download but takes the HF repo id directly, so search
    results never need to be persisted into the catalog. Creates the local=2
    placeholder + DownloadJob and runs the download in the background.

    Args:
        payload: The chosen HF repo id plus display metadata from the search hit.

    Returns:
        DownloadJobResponse: Job record to poll via GET /downloads/{job_id}/status.
    """
    try:
        name = payload.name or humanize_model_name(payload.link)
        # Family type is metadata-only for a local row; fall back to the category.
        type_ = payload.type or payload.category or "community"
        return _start_download(
            remote_model_id=payload.link,
            remote_link=payload.link,
            name=name,
            type=type_,
            description=None,
            model_metadata=None,
            quantized=payload.quantized,
            param_size=payload.param_size,
            category=payload.category,
            llm_repo=llm_repo,
            job_repo=job_repo,
            db=db,
            background_tasks=background_tasks,
        )
    except (FileSystemException, InvalidInputException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to start download", trace=str(e))

@router.post(
    "/{llm_id}/download",
    response_model=DownloadJobResponse,
    status_code=http_status.HTTP_200_OK,
)
async def download_llm_route(
    llm_id: int,
    background_tasks: BackgroundTasks,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    job_repo: Download_Job_Repository = Depends(get_download_job_repository),
    db: Session = Depends(get_db),
):
    """Start background download of LLM from HuggingFace.

    Creates a DownloadJob and runs download in background. The catalog link is
    already a pre-built quant for this engine (MLX repo on Apple Silicon, GGUF
    file elsewhere); nothing is converted locally.

    Args:
        llm_id: ID of the remote LLM to download.
        background_tasks: FastAPI background tasks manager.
        llm_repo: Injected LLM repository.
        job_repo: Injected download job repository.
        db: Database session for transaction control.

    Returns:
        DownloadJobResponse: Job record with job_id, status, progress.

    Raises:
        HTTPException: 404 if LLM not found, 500 if paths already exist.

    Note:
        Download progress can be polled via GET /downloads/{job_id}/status.
        Large models (7B+) take 5-30 minutes.
    """
    try:
        # Get remote LLM metadata
        remote_llm = llm_repo.get_by_id(llm_id)
        if not remote_llm:
            raise ModelNotFoundException(f"LLM {llm_id}")

        return _start_download(
            remote_model_id=str(llm_id),
            remote_link=remote_llm.link,
            name=remote_llm.name,
            type=remote_llm.type,
            description=remote_llm.description,
            model_metadata=remote_llm.model_metadata,
            quantized=remote_llm.quantized,
            param_size=remote_llm.param_size,
            category=getattr(remote_llm, "category", "general"),
            llm_repo=llm_repo,
            job_repo=job_repo,
            db=db,
            background_tasks=background_tasks,
            generation_hints=getattr(remote_llm, "generation_hints", None),
            artifact_size_bytes=remote_llm.artifact_size_bytes,
        )

    except (ModelNotFoundException, InvalidInputException, FileSystemException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to start download", trace=str(e))

@router.get("/", response_model=List[LLMResponse])
async def get_all_llms(llm_repo: Llm_Repository = Depends(get_llm_repository)):
    """List all LLMs (local and remote).

    Args:
        llm_repo: Injected LLM repository.

    Returns:
        List[LLMResponse]: All LLM models with metadata.
    """
    # Read-only operation, no commit needed
    llms = llm_repo.get_all()
    return llms

def get_download_job_repository(db: Session = Depends(get_db)) -> Download_Job_Repository:
    """Dependency injection for Download_Job_Repository.

    Args:
        db: Database session from FastAPI.

    Returns:
        Download_Job_Repository: Repository instance with injected session.
    """
    return Download_Job_Repository(db)

@router.get(
    "/downloads/{job_id}/status",
    response_model=DownloadJobResponse,
    status_code=200,
)
def get_download_status_by_jobId(
    job_id: int,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    job_repo: Download_Job_Repository = Depends(get_download_job_repository),
    db: Session = Depends(get_db),
):
    """Get download job status by ID with automatic cleanup for failed/cancelled jobs.

    Polls the download job status and performs cleanup operations for terminal states
    (failed/cancelled/completed). Failed/cancelled jobs delete the temp LLM entry and
    clean up temporary files. Completed jobs mark the LLM as local=1 (ready).

    Args:
        job_id: The database ID of the download job to query.
        llm_repo: Injected LLM repository.
        job_repo: Injected download job repository.
        db: Database session for transaction control.

    Returns:
        DownloadJobResponse: Download job with current status, progress, ETA, and file paths.

    Raises:
        DownloadJobNotFoundException: If job_id not found.
        ModelNotFoundException: If a completed job's associated LLM is missing.
        DatabaseException: If status fetch fails.

    Example:
        GET /llms/downloads/42/status
        Response: {"id": 42, "status": "running", "progress": 65.0, "eta_seconds": 120, ...}
    """
    try:
        job = job_repo.get_by_id(job_id)
        if not job:
            raise DownloadJobNotFoundException(job_id)

        # Handle failed/cancelled jobs: cleanup temp files and LLM entry.
        # `cancel_download_job` already deletes the temp LLM (and Postgres
        # nulls local_model_id through the FK) before the job reaches
        # `cancelled`, so a missing row here is the common case, not an
        # error - only run cleanup when there is still something to clean.
        if job.status in ["failed", "cancelled"]:
            llm = llm_repo.get_by_id(job.local_model_id)
            if llm:
                # Delete temp LLM entry
                llm_repo.delete(llm)

                # Clean up temp files using repository method
                job_repo.cleanup_job_files(job)

                db.commit()
                db.refresh(job)
                logger.info(f"Cleaned up {job.status} download job {job_id}")

        # Handle completed jobs: mark LLM as ready, ONCE.
        #
        # The frontend polls this endpoint every 2s while a download is in
        # flight, and keeps polling a completed job until it navigates away. The
        # write below used to run on every one of those ticks: a GET handler
        # issuing an UPDATE plus a COMMIT, and an INFO log line, for a row that
        # already said local = 1. Guarding on the current value makes the write
        # idempotent in the real sense (it happens once) rather than merely
        # convergent (it happens forever but lands on the same value).
        elif job.status == "completed":
            llm = llm_repo.get_by_id(job.local_model_id)
            if not llm:
                raise ModelNotFoundException(f"LLM {job.local_model_id}")
            if llm.local != 1:
                llm_repo.update(llm, local=1)
                db.commit()
                db.refresh(llm)
                logger.info(f"Marked LLM {llm.id} as ready (download job {job_id} completed)")

        return job

    except (DownloadJobNotFoundException, ModelNotFoundException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to get download status", trace=str(e))

@router.get(
    "/downloads/status",
    response_model=DownloadJobResponse,
    status_code=200,
)
def get_download_status_without_jobId(
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    job_repo: Download_Job_Repository = Depends(get_download_job_repository),
    db: Session = Depends(get_db),
):
    """Get most recent active download job (for single-download UI polling).

    Finds the most recently updated download job that is still running/pending within
    the last 60 seconds. Useful for UIs that only support one active download at a time
    and need to poll without tracking job IDs. Performs same cleanup as get_by_id.

    Args:
        llm_repo: Injected LLM repository.
        job_repo: Injected download job repository.
        db: Database session for transaction control.

    Returns:
        DownloadJobResponse: The most recent active job with status and progress.

    Raises:
        DownloadJobNotFoundException: If no active job found in last 60 seconds.
        ModelNotFoundException: If a completed job's associated LLM is missing.
        DatabaseException: If status fetch fails.

    Example:
        GET /llms/downloads/status
        Response: {"id": 42, "status": "running", "progress": 65.0, ...}
    """
    try:
        # Get most recent active job
        job = job_repo.get_most_recent_active()
        if not job:
            raise DownloadJobNotFoundException("recent active")

        # Handle failed jobs: cleanup temp files and LLM entry. Same
        # already-cleaned-up guard as get_download_status_by_jobId - see #509.
        if job.status == "failed":
            llm = llm_repo.get_by_id(job.local_model_id)
            if llm:
                # Delete temp LLM entry
                llm_repo.delete(llm)

                # Clean up temp files using repository method
                job_repo.cleanup_job_files(job)

                db.commit()
                db.refresh(job)
                logger.info(f"Cleaned up failed download job {job.id}")

        # Handle completed jobs: mark LLM as ready, ONCE. Same idempotency
        # guard as get_download_status_by_jobId - see the comment there.
        elif job.status == "completed":
            llm = llm_repo.get_by_id(job.local_model_id)
            if not llm:
                raise ModelNotFoundException(f"LLM {job.local_model_id}")
            if llm.local != 1:
                llm_repo.update(llm, local=1)
                db.commit()
                db.refresh(llm)
                logger.info(f"Marked LLM {llm.id} as ready (download job {job.id} completed)")

        return job

    except (DownloadJobNotFoundException, ModelNotFoundException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to get download status", trace=str(e))

@router.get("/{llm_id}", response_model=LLMResponse)
async def get_llm_by_id(llm_id: int, llm_repo: Llm_Repository = Depends(get_llm_repository)):
    """Get LLM details by ID.

    Args:
        llm_id: ID of the LLM to retrieve.
        llm_repo: Injected LLM repository.

    Returns:
        LLMResponse: LLM metadata.

    Raises:
        ModelNotFoundException: If LLM not found.
    """
    # Read-only operation, no commit needed
    llm = llm_repo.get_by_id(llm_id)
    if not llm:
        raise ModelNotFoundException(f"LLM {llm_id}")
    return llm

@router.get("/{llm_id}/dependents", response_model=DependentsResponse)
async def get_llm_dependents(
    llm_id: int,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
):
    """List the KB assistants that depend on this model's weights, plus the
    conversation counts a delete would orphan.

    Dependents share the model's COPIED ``link`` (#209): deleting the base
    leaves their weights dangling. A model with no dependents returns an empty
    list with zero-or-more conversation counts. The frontend calls this before a
    base-model delete to render the choice dialog.

    A KB assistant NEVER has dependents (#317): deleting it is a direct 200
    that keeps the shared weights (they belong to the base), so sibling
    assistants copying the same link are not reported — this endpoint answers
    with the exact truth the guarded DELETE's 409 uses.

    Args:
        llm_id: ID of the (base) model to inspect.
        llm_repo: Injected LLM repository.

    Returns:
        DependentsResponse: assistants + own/total conversation counts.

    Raises:
        ModelNotFoundException: If the model is not found.
    """
    llm = llm_repo.get_by_id(llm_id)
    if not llm:
        raise ModelNotFoundException(f"LLM {llm_id}")
    if llm.is_attached_to_kb:
        own = llm_repo.count_conversations(llm.id)
        return {
            "assistants": [],
            "own_conversation_count": own,
            "total_conversation_count": own,
        }
    return _dependents_payload(llm_repo, llm)

def get_llm_repository(db: Session = Depends(get_db)) -> Llm_Repository:
    """Dependency injection for Llm_Repository.

    Args:
        db: Database session from FastAPI.

    Returns:
        Llm_Repository: Repository instance with injected session.
    """
    return Llm_Repository(db)

@router.get("/local", response_model=List[LLMResponse])
async def get_local_llms(llm_repo: Llm_Repository = Depends(get_llm_repository)):
    """List only local (downloaded) LLMs ready for inference.

    Args:
        llm_repo: Injected LLM repository.

    Returns:
        List[LLMResponse]: LLMs with local=1.
    """
    # Read-only operation, no commit needed
    llms = llm_repo.get_all_local()
    return llms

@router.get("/remote", response_model=List[LLMResponse])
async def get_remote_llms(llm_repo: Llm_Repository = Depends(get_llm_repository)):
    """List only remote (HuggingFace) LLMs available for download.

    Args:
        llm_repo: Injected LLM repository.

    Returns:
        List[LLMResponse]: LLMs with local=0.
    """
    # Read-only operation, no commit needed
    llms = llm_repo.get_all_remote()
    return llms

@router.post("/{assistant_id}/rebind", response_model=LLMResponse)
async def rebind_assistant(
    assistant_id: int,
    payload: RebindRequest,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    db: Session = Depends(get_db),
):
    """Re-point an orphaned KB assistant at a new base model.

    Re-copies the base's weights ``link`` plus every descriptive COPIED_FIELDS
    entry (#209) onto the assistant, keeping its own name/description and KB
    wiring (kb_id, is_attached_to_kb). Used to recover an assistant whose base
    was deleted (weights_available False). The target must be a local,
    non-assistant model whose weights exist on disk.

    Args:
        assistant_id: ID of the KB assistant to rebind.
        payload: RebindRequest with ``new_base_llm_id``.
        llm_repo: Injected LLM repository.
        db: Database session for transaction control.

    Returns:
        LLMResponse: The updated assistant (weights_available now True).

    Raises:
        ModelNotFoundException: 404 if the assistant or target base is missing.
        StateConflictException: 409 if the assistant is not a KB assistant, or
            the target is itself an assistant / not downloaded / has no weights.
        DatabaseException: If the rebind fails.
    """
    try:
        assistant = llm_repo.get_by_id(assistant_id)
        if not assistant:
            raise ModelNotFoundException(f"LLM {assistant_id}")
        if not assistant.is_attached_to_kb:
            raise StateConflictException(
                f"Model '{assistant.name}' is not a KB assistant; only KB "
                "assistants can be rebound to a new base.",
                status_code=http_status.HTTP_409_CONFLICT,
            )

        new_base = llm_repo.get_by_id(payload.new_base_llm_id)
        if not new_base:
            raise ModelNotFoundException(f"LLM {payload.new_base_llm_id}")
        if new_base.is_attached_to_kb:
            raise StateConflictException(
                f"Target model '{new_base.name}' is itself a KB assistant; pick "
                "a standalone base model.",
                status_code=http_status.HTTP_409_CONFLICT,
            )
        if new_base.local != 1:
            raise StateConflictException(
                f"Target model '{new_base.name}' is not downloaded; its weights "
                "must exist to rebind onto it.",
                status_code=http_status.HTTP_409_CONFLICT,
            )
        if not new_base.link or not os.path.exists(new_base.link):
            raise StateConflictException(
                f"Target model '{new_base.name}' weights are missing on disk; "
                "cannot rebind onto it.",
                status_code=http_status.HTTP_409_CONFLICT,
            )

        # Re-copy link + descriptive columns; name/description/KB wiring untouched.
        updated = llm_repo.update(
            assistant,
            **{field: getattr(new_base, field) for field in COPIED_FIELDS},
        )
        db.commit()
        db.refresh(updated)
        logger.info(f"Rebound assistant {assistant_id} onto base {new_base.id}")
        return updated

    except (ModelNotFoundException, StateConflictException):
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to rebind assistant", trace=str(e))

@router.get("/search/huggingface", response_model=List[HFSearchResult])
async def search_huggingface_route(q: str, limit: int = 30):
    """Live HuggingFace search beyond the curated catalog (#122 follow-up).

    Searches HF directly for runnable models matching ``q`` in the active engine's
    format, filtered to chat/vision LLMs (so a query like "french" doesn't return
    token-classification repos). Results are ephemeral — they are NOT added to the
    catalog; download a chosen one via POST /download/huggingface.

    Args:
        q: Free-text query (model name, family, trait like "uncensored").
        limit: Max results to return (default 30).

    Returns:
        List[HFSearchResult]: Downloadable matches with metadata + category.
    """
    return search_huggingface(q, limit=limit)

@router.get("/search", response_model=List[LLMResponse])
async def search_llms(name: str, llm_repo: Llm_Repository = Depends(get_llm_repository)):
    """Search LLMs by name (case-insensitive partial match).

    Args:
        name: Search query string.
        llm_repo: Injected LLM repository.

    Returns:
        List[LLMResponse]: Matching LLMs.
    """
    # Read-only operation, no commit needed
    llms = llm_repo.search_by_name(name)
    return llms

@router.put("/{llm_id}", response_model=LLMResponse)
async def update_llm(
    llm_id: int,
    llm_data: LLMCreate,
    llm_repo: Llm_Repository = Depends(get_llm_repository),
    db: Session = Depends(get_db),
):
    """Update LLM metadata (name, description, etc.).

    Args:
        llm_id: ID of the LLM to update.
        llm_data: LLMCreate schema with new values.
        llm_repo: Injected LLM repository.
        db: Database session for transaction control.

    Returns:
        LLMResponse: Updated LLM.

    Raises:
        ModelNotFoundException: If LLM not found.
        DatabaseException: If update fails.
    """
    try:
        llm = llm_repo.get_by_id(llm_id)
        if not llm:
            raise ModelNotFoundException(f"LLM {llm_id}")

        # Update fields from request
        updated_llm = llm_repo.update(llm, **llm_data.dict())
        db.commit()
        return updated_llm
    except ModelNotFoundException:
        raise
    except Exception as e:
        db.rollback()
        raise DatabaseException("Failed to update LLM", trace=str(e))