4c1ab745d6
Phase 3 lays two orthogonal per-KB knobs in the data and API layers so
follow-up commits (Phase 3B-late, Phase 3B-multimodal) only need to add
their code path and UI selector, not schema or types.
Schema (studio/backend/storage/studio_db.py)
- rag_knowledge_bases gains chunking_strategy ('standard'|'late', default
'standard') and mode ('text'|'multimodal', default 'text'). Both are
immutable after KB creation — changing either invalidates existing
chunks because they were ingested through a specific pipeline.
- rag_chunks gains kind ('text'|'image'|'caption', default 'text'),
image_path (NULLABLE), linked_chunk_id (NULLABLE) — used by Phase
3B-multimodal to pair image chunks with their captions.
- Idempotent ALTER TABLE additions for existing installs (mirrors the
chat_threads display_name / *_code_exec_container_id pattern earlier
in the file).
API (studio/backend/routes/rag.py)
- ChunkingStrategy + KBMode Literal aliases.
- CreateKBRequest accepts both fields with backward-compat defaults.
- KBResponse exposes both.
- _validate_mode_combo rejects (multimodal, late) with 400 — no public
open-weight embedder supports both at once. Surface the constraint
early rather than failing silently during ingestion.
Config (studio/backend/utils/rag/config.py)
- RAG_EMBEDDER_MATRIX dict keyed by (mode, strategy) → embedder name.
- resolve_embedder() helper falls back to RAG_EMBEDDING_MODEL for legacy
KBs that pre-date the columns.
- (multimodal, late) intentionally absent.
Frontend (studio/frontend/src/features/rag/)
- api/rag-api.ts: ChunkingStrategy + KBMode types; KnowledgeBase
interface and createKnowledgeBase request type updated.
- stores/rag-store.ts: createKB signature uses the shared request type.
No user-visible UI changes yet — the only currently-usable combination
is (text, standard), so adding one-option selectors would be UX noise.
Phase 3B-late and Phase 3B-multimodal each add the relevant selector
option as part of shipping the code path.
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
|
|
def _env_int(name: str, default: int) -> int:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
return int(raw)
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _env_float(name: str, default: float) -> float:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
return float(raw)
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
RAG_EMBEDDING_MODEL: str = (
|
|
os.environ.get("UNSLOTH_RAG_EMBEDDING_MODEL", "").strip()
|
|
or "BAAI/bge-small-en-v1.5"
|
|
)
|
|
|
|
# Phase 3: default embedders per (mode, chunking_strategy). Ingestion in
|
|
# Phase 3B-late and Phase 3B-multimodal looks the embedder up here at job
|
|
# start, falling back to RAG_EMBEDDING_MODEL (above) for legacy KBs that
|
|
# pre-date the columns. The (multimodal, late) combo is intentionally
|
|
# absent — no public open-weight embedder supports both at once, and
|
|
# routes/rag.py rejects the combo with a 400 at KB create time.
|
|
RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = {
|
|
("text", "standard"): "BAAI/bge-small-en-v1.5",
|
|
("text", "late"): "nomic-ai/nomic-embed-text-v1.5",
|
|
("multimodal", "standard"): "BAAI/BGE-VL-base",
|
|
}
|
|
|
|
|
|
def resolve_embedder(mode: str, chunking_strategy: str) -> str:
|
|
"""Look up the default embedder for a (mode, chunking_strategy) pair.
|
|
|
|
Unknown combos fall back to the legacy single default so old KBs
|
|
keep working. Callers that explicitly require the new matrix
|
|
behaviour (Phase 3B paths) should validate the inputs before
|
|
calling.
|
|
"""
|
|
return RAG_EMBEDDER_MATRIX.get(
|
|
(mode, chunking_strategy),
|
|
RAG_EMBEDDING_MODEL,
|
|
)
|
|
|
|
RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512)
|
|
RAG_CHUNK_OVERLAP: int = _env_int("UNSLOTH_RAG_CHUNK_OVERLAP", 64)
|
|
|
|
RAG_TOP_K_BM25: int = _env_int("UNSLOTH_RAG_TOP_K_BM25", 30)
|
|
RAG_TOP_K_DENSE: int = _env_int("UNSLOTH_RAG_TOP_K_DENSE", 30)
|
|
RAG_TOP_K_HYBRID: int = _env_int("UNSLOTH_RAG_TOP_K_HYBRID", 10)
|
|
|
|
RAG_RRF_K: int = _env_int("UNSLOTH_RAG_RRF_K", 60)
|
|
|
|
RAG_MAX_UPLOAD_MB: int = _env_int("UNSLOTH_RAG_MAX_UPLOAD_MB", 50)
|
|
|
|
RAG_EMBED_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_EMBED_BATCH_SIZE", 32)
|
|
|
|
# Reranking is off by default. The CrossEncoder runs on GPU and competes
|
|
# with the active chat model — callers opt in per-request via
|
|
# `enable_rerank` on SearchRequest.
|
|
RAG_RERANKER_MODEL: str = (
|
|
os.environ.get("UNSLOTH_RAG_RERANKER_MODEL", "").strip()
|
|
or "BAAI/bge-reranker-base"
|
|
)
|
|
RAG_RERANK_CANDIDATE_K: int = _env_int("UNSLOTH_RAG_RERANK_CANDIDATE_K", 50)
|
|
RAG_RERANK_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_RERANK_BATCH_SIZE", 16)
|
|
|
|
RAG_UPLOAD_EXTS: frozenset[str] = frozenset(
|
|
{".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
|
|
)
|