Backend (studio/backend/): - core/rag/: parsers (PDF/TXT/MD/DOCX/HTML via pypdf/python-docx/bs4), recursive token-aware chunker, embeddings singleton via FastSentenceTransformer.from_pretrained(for_inference=True), Qdrant local vector store, bm25s lexical index, RRF hybrid retrieval, spawn-subprocess ingestion job with SSE progress, optional CrossEncoder reranker (off-by-default). - routes/rag.py: KB CRUD, doc upload (KB + per-thread), doc list/delete, ingestion SSE, hybrid+rerank search, thread-index list/clear. - routes/chat_history.py: purge thread RAG artifacts on thread delete and clear-all (rag_documents has no FK cascade to chat_threads so uploads work on un-persisted threads). - studio.db gains 4 RAG tables; storage_roots gains rag_*() helpers. - auth/authentication.py: get_current_subject_sse accepts ?token=... so EventSource can stream ingestion progress. Frontend (studio/frontend/): - features/rag/: api client, Zustand store, hooks, dropzone, KB list, doc rows, ingestion-progress, thread-index list components. - Settings dialog gains a Knowledge Bases tab (master/detail + thread documents list); /knowledge-bases deep-links to it. - features/chat/: per-thread ragSource/enableRerank/ragTopK state in chat-runtime-store; Retrieval section in chat-settings-sheet with KB DropdownMenu (active highlight + per-row trash), thread doc list with Clear-thread-index button, RAG Top K slider, reranker toggle; chat-adapter retrieves before /v1/chat/completions and injects hits as a system block; shared-composer + button routes documents into pendingDocs (auto-uploads, send blocked while indexing).
59 lines
1.8 KiB
Python
59 lines
1.8 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"
|
|
)
|
|
|
|
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"}
|
|
)
|