ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-05-25 06:38:58 +00:00
commit b931b0039b
18 changed files with 107 additions and 74 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/auth/authentication.py
View file

@ -168,7 +168,7 @@ async def get_current_subject_sse(
"""
raw = token
if not raw and authorization and authorization.lower().startswith("bearer "):
raw = authorization[len("bearer "):].strip()
raw = authorization[len("bearer ") :].strip()
if not raw:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,

1
studio/backend/core/inference/tools.py
View file

@ -502,6 +502,7 @@ TERMINAL_TOOL = {
},
}
# Lazy import — keeps studio.db init lazy so tools.py doesn't pull in
# the whole rag stack on inference paths that never see RAG.
def _get_rag_tool_spec():

2
studio/backend/core/rag/chunking.py
View file

@ -57,7 +57,7 @@ def _atomic_split(
if count(piece) <= max_tokens:
out.append(piece)
else:
tail = separators[separators.index(sep) + 1:]
tail = separators[separators.index(sep) + 1 :]
out.extend(_atomic_split(piece, tail, max_tokens, count))
return out
# No separator made progress — hard-slice by characters.

12
studio/backend/core/rag/embeddings.py
View file

@ -109,10 +109,14 @@ class _BGEVLAdapter:
from PIL import Image
if inputs is None or len(inputs) == 0:
return np.zeros((0, self.get_sentence_embedding_dimension()), dtype = np.float32)
return np.zeros(
(0, self.get_sentence_embedding_dimension()), dtype = np.float32
)
sample = inputs[0]
is_image = isinstance(sample, Image.Image) or isinstance(sample, (bytes, bytearray))
is_image = isinstance(sample, Image.Image) or isinstance(
sample, (bytes, bytearray)
)
chunks_out = []
for start in range(0, len(inputs), batch_size):
@ -513,9 +517,7 @@ def _windowed_late_chunk_encode(
ws, _we = windows[best_window]
emb = _window_embeddings(best_window)
local_indices = [
ti - ws
for ti in chunk_token_indices
if ws <= ti < ws + emb.shape[0]
ti - ws for ti in chunk_token_indices if ws <= ti < ws + emb.shape[0]
]
if not local_indices:
vec = model.encode(

13
studio/backend/core/rag/ingestion.py
View file

@ -78,7 +78,9 @@ def _subprocess_worker(
parsed = parse(Path(stored_path), want_images = (mode == "multimodal"))
pages = parsed.pages
if not pages and not parsed.images:
out_queue.put({"type": "error", "error": "no extractable content in document"})
out_queue.put(
{"type": "error", "error": "no extractable content in document"}
)
return
out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1})
@ -249,9 +251,7 @@ def _stream_image_chunks(
image_vectors = encode_images(bytes_for_encoding, model_name = model_name)
# Embed only the non-empty captions; track which images they map to.
caption_to_image: list[int] = [
i for i, cap in enumerate(captions) if cap.strip()
]
caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()]
if caption_to_image:
caption_vectors_arr = encode(
[captions[i] for i in caption_to_image],
@ -369,6 +369,7 @@ def _run_late_chunking(
# Job manager (parent side)
# ------------------------------------------------------------------
class _JobState:
def __init__(self, job_id: str, document_id: str, scope: str) -> None:
self.job_id = job_id
@ -531,7 +532,7 @@ def _insert_chunks_and_collect_for_bm25(
def _all_scope_chunks(scope: str) -> list[dict]:
if scope.startswith("kb_"):
kb_id = scope[len("kb_"):]
kb_id = scope[len("kb_") :]
sql = (
"SELECT c.id, c.text FROM rag_chunks c "
"JOIN rag_documents d ON d.id = c.document_id "
@ -539,7 +540,7 @@ def _all_scope_chunks(scope: str) -> list[dict]:
)
bind = (kb_id,)
elif scope.startswith("thread_"):
thread_id = scope[len("thread_"):]
thread_id = scope[len("thread_") :]
sql = (
"SELECT c.id, c.text FROM rag_chunks c "
"JOIN rag_documents d ON d.id = c.document_id "

3
studio/backend/core/rag/parsers/__init__.py
View file

@ -18,6 +18,7 @@ class ParsedPage:
pipe-tables, and list bullets survive extraction so the chunker can
split on them. Parsers MUST emit Markdown, not bare plain text.
"""
text: str
page_number: int | None = None
@ -32,6 +33,7 @@ class ParsedImage:
when no caption could be paired (the image still ingests, just
without the paired-caption chunk).
"""
image_bytes: bytes
mime_type: str
page_number: int | None = None
@ -46,6 +48,7 @@ class ParseResult:
passed `want_images=True`. Iteration aliases for `pages` so legacy
code that did `for page in parse(path)` keeps working.
"""
pages: list[ParsedPage] = field(default_factory = list)
images: list[ParsedImage] = field(default_factory = list)

4
studio/backend/core/rag/parsers/pdf.py
View file

@ -56,7 +56,9 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult:
def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]:
"""Pull embedded images and pair each with the nearest text on the same page."""
captions_by_page: dict[int, str] = {p.page_number: p.text for p in pages if p.page_number}
captions_by_page: dict[int, str] = {
p.page_number: p.text for p in pages if p.page_number
}
out: list[ParsedImage] = []
for page_index in range(len(doc)):
page_number = page_index + 1

4
studio/backend/core/rag/retrieval.py
View file

@ -101,7 +101,9 @@ def _rrf_fuse(
dense_scores: dict[str, float] = {}
for ranking in rankings:
for rank, hit in enumerate(ranking):
fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (rrf_k + rank + 1)
fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (
rrf_k + rank + 1
)
if hit.chunk_id not in seen:
seen[hit.chunk_id] = hit
if hit.dense_score is not None:

4
studio/backend/core/rag/tool.py
View file

@ -114,9 +114,7 @@ def search_knowledge_base(
from core.rag.vector_store import kb_scope, thread_scope
from storage.studio_db import get_connection
scope = (
kb_scope(scope_kb_id) if scope_kb_id else thread_scope(scope_thread_id)
)
scope = kb_scope(scope_kb_id) if scope_kb_id else thread_scope(scope_thread_id)
k = top_k if top_k is not None else default_top_k
if enable_rerank:

8
studio/backend/routes/inference.py
View file

@ -2484,9 +2484,7 @@ async def openai_chat_completions(
else 300,
session_id = payload.session_id,
tool_context = (
{"rag_scope": payload.rag_scope}
if payload.rag_scope
else None
{"rag_scope": payload.rag_scope} if payload.rag_scope else None
),
)
@ -2956,9 +2954,7 @@ async def openai_chat_completions(
def sf_generate_with_tools():
return backend.generate_chat_completion_with_tools(
tool_context = (
{"rag_scope": payload.rag_scope}
if payload.rag_scope
else None
{"rag_scope": payload.rag_scope} if payload.rag_scope else None
),
messages = _sf_chat_messages,
tools = _sf_tools_to_use,

62
studio/backend/routes/rag.py
View file

@ -28,7 +28,15 @@ from pathlib import Path
from typing import Any, Literal, Optional
from uuid import uuid4
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, UploadFile
from fastapi import (
APIRouter,
Depends,
Header,
HTTPException,
Query,
Request,
UploadFile,
)
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
@ -40,6 +48,8 @@ async def _sse_auth(
authorization: str | None = Header(None),
) -> str:
return await get_current_subject_sse(token, authorization)
from core.rag import embeddings, ingestion, reranker, retrieval, vector_store
from core.rag.vector_store import kb_scope, thread_scope
from loggers import get_logger
@ -161,6 +171,7 @@ class SearchResponse(BaseModel):
# Helpers
# ------------------------------------------------------------------
def _sanitize_filename(filename: str) -> str:
name = Path(filename).name.strip().replace("\x00", "")
return name or "document"
@ -182,7 +193,7 @@ def _resolve_scope_embedder(scope: str) -> str | None:
from utils.rag.config import resolve_embedder
if scope.startswith("kb_"):
kb_id = scope[len("kb_"):]
kb_id = scope[len("kb_") :]
with get_connection() as conn:
row = conn.execute(
"SELECT embedding_model FROM rag_knowledge_bases WHERE id = ?",
@ -190,7 +201,7 @@ def _resolve_scope_embedder(scope: str) -> str | None:
).fetchone()
return row["embedding_model"] if row else None
if scope.startswith("thread_"):
thread_id = scope[len("thread_"):]
thread_id = scope[len("thread_") :]
settings = _load_thread_settings(thread_id)
return settings.embedding_model or resolve_embedder(
settings.mode,
@ -204,9 +215,7 @@ def _row_to_kb(row: Any) -> KBResponse:
# pre-Phase-3 connection in tests; fall back to the schema defaults.
keys = row.keys() if hasattr(row, "keys") else ()
chunking_strategy = (
row["chunking_strategy"]
if "chunking_strategy" in keys
else "standard"
row["chunking_strategy"] if "chunking_strategy" in keys else "standard"
)
mode = row["mode"] if "mode" in keys else "text"
return KBResponse(
@ -378,6 +387,7 @@ def _unlink_if_under_uploads(path: Path) -> None:
# Knowledge bases
# ------------------------------------------------------------------
@router.post("/knowledge-bases", response_model = KBResponse)
def create_knowledge_base(
payload: CreateKBRequest,
@ -391,9 +401,8 @@ def create_knowledge_base(
# If the caller didn't override embedding_model, resolve from the
# Phase-3 matrix using their (mode, strategy) selection. Unknown
# combos fall back to the legacy default — see resolve_embedder.
embedding_model = (
payload.embedding_model
or resolve_embedder(payload.mode, payload.chunking_strategy)
embedding_model = payload.embedding_model or resolve_embedder(
payload.mode, payload.chunking_strategy
)
created_at = _now_ms()
with get_connection() as conn:
@ -452,6 +461,7 @@ class RagDefaults(BaseModel):
class UpdateRagDefaultsRequest(BaseModel):
"""Patch shape — only fields present overwrite stored values."""
chunking_strategy: ChunkingStrategy | None = None
mode: KBMode | None = None
embedding_model: str | None = None
@ -541,9 +551,7 @@ def _load_thread_settings(thread_id: str) -> ThreadRagSettings:
raw = {}
fallback = _load_rag_defaults()
return ThreadRagSettings(
chunking_strategy = (
raw.get("chunking_strategy") or fallback.chunking_strategy
),
chunking_strategy = (raw.get("chunking_strategy") or fallback.chunking_strategy),
mode = raw.get("mode") or fallback.mode,
embedding_model = raw.get("embedding_model") or fallback.embedding_model,
)
@ -598,6 +606,7 @@ def set_thread_rag_settings(
class ReingestKBRequest(BaseModel):
"""All fields optional — omitting one keeps the KB's current value."""
chunking_strategy: ChunkingStrategy | None = None
mode: KBMode | None = None
embedding_model: str | None = None
@ -689,9 +698,7 @@ def reingest_knowledge_base(
kb_row = _kb_or_404(kb_id)
keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
current_strategy = (
kb_row["chunking_strategy"]
if "chunking_strategy" in keys
else "standard"
kb_row["chunking_strategy"] if "chunking_strategy" in keys else "standard"
)
current_mode = kb_row["mode"] if "mode" in keys else "text"
current_embedder = kb_row["embedding_model"]
@ -700,13 +707,10 @@ def reingest_knowledge_base(
new_mode = payload.mode or current_mode
_validate_mode_combo(new_mode, new_strategy)
new_embedder = (
payload.embedding_model
or (
current_embedder
if (new_strategy == current_strategy and new_mode == current_mode)
else resolve_embedder(new_mode, new_strategy)
)
new_embedder = payload.embedding_model or (
current_embedder
if (new_strategy == current_strategy and new_mode == current_mode)
else resolve_embedder(new_mode, new_strategy)
)
with get_connection() as conn:
@ -800,6 +804,7 @@ def delete_knowledge_base(
# Document upload (KB and per-thread)
# ------------------------------------------------------------------
@router.post("/knowledge-bases/{kb_id}/documents", response_model = UploadResponse)
async def upload_kb_document(
kb_id: str,
@ -813,9 +818,7 @@ async def upload_kb_document(
# fall back to the same defaults as the column.
kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
chunking_strategy = (
kb_row["chunking_strategy"]
if "chunking_strategy" in kb_keys
else "standard"
kb_row["chunking_strategy"] if "chunking_strategy" in kb_keys else "standard"
)
mode = kb_row["mode"] if "mode" in kb_keys else "text"
return _start_ingestion(
@ -867,6 +870,7 @@ async def upload_thread_document(
# Document list / delete
# ------------------------------------------------------------------
@router.get("/knowledge-bases/{kb_id}/documents", response_model = DocumentListResponse)
def list_kb_documents(
kb_id: str,
@ -910,7 +914,7 @@ def get_rag_image(
if "/" in filename or "\\" in filename or filename.startswith("."):
raise HTTPException(status_code = 400, detail = "Invalid filename")
root = Path(os.path.realpath(rag_uploads_root() / "images"))
candidate = (rag_uploads_root() / "images" / document_id / filename)
candidate = rag_uploads_root() / "images" / document_id / filename
try:
real = Path(os.path.realpath(candidate))
real.relative_to(root)
@ -927,9 +931,7 @@ def delete_document(
current_subject: str = Depends(get_current_subject),
) -> dict:
row = _document_or_404(document_id)
scope = (
kb_scope(row["kb_id"]) if row["kb_id"] else thread_scope(row["thread_id"])
)
scope = kb_scope(row["kb_id"]) if row["kb_id"] else thread_scope(row["thread_id"])
with get_connection() as conn:
conn.execute("DELETE FROM rag_documents WHERE id = ?", (document_id,))
conn.commit()
@ -995,6 +997,7 @@ def clear_thread_documents(
# Ingestion job SSE
# ------------------------------------------------------------------
@router.get("/jobs/{job_id}/events")
async def job_events(
job_id: str,
@ -1065,6 +1068,7 @@ async def _replay_terminal_state(row: Any):
# Search
# ------------------------------------------------------------------
@router.post("/search", response_model = SearchResponse)
def search(
payload: SearchRequest,

4
studio/backend/utils/rag/config.py
View file

@ -62,6 +62,7 @@ def resolve_embedder(mode: str, chunking_strategy: str) -> str:
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)
@ -79,8 +80,7 @@ RAG_EMBED_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_EMBED_BATCH_SIZE", 32)
# 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"
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)

4
tests/python/test_rag_chunking.py
View file

@ -26,7 +26,9 @@ def test_chunk_pages_splits_long_text():
)
assert len(chunks) > 1
for chunk in chunks:
assert _wc_counter(chunk.text) <= 55 # max + small slack from atomic split granularity
assert (
_wc_counter(chunk.text) <= 55
) # max + small slack from atomic split granularity
def test_chunk_pages_short_text_is_one_chunk():

1
tests/python/test_rag_late_chunking.py
View file

@ -82,6 +82,7 @@ def test_late_chunk_encode_returns_one_vector_per_span():
pytest.importorskip("torch")
# all-MiniLM-L6-v2 is ~80MB and embeds at 384 dims.
import os
os.environ.setdefault(
"UNSLOTH_RAG_EMBEDDING_MODEL",
"sentence-transformers/all-MiniLM-L6-v2",

6
tests/python/test_rag_multimodal.py
View file

@ -34,10 +34,10 @@ def test_html_parser_returns_images_when_requested(tmp_path):
img_path.write_bytes(png_bytes)
html_path = tmp_path / "sample.html"
html_path.write_text(
f'<html><body><h1>Doc</h1>'
f'<p>Body text.</p>'
f"<html><body><h1>Doc</h1>"
f"<p>Body text.</p>"
f'<img src="tiny.png" alt="A tiny figure">'
f'</body></html>',
f"</body></html>",
encoding = "utf-8",
)

6
tests/python/test_rag_multimodal_integration.py
View file

@ -168,6 +168,6 @@ def test_text_and_image_vectors_share_dimension(monkeypatch):
image_vectors = embeddings_module.encode_images([buf.getvalue()])
text_vectors = embeddings_module.encode(["a blue square"])
assert image_vectors[0].shape == text_vectors[0].shape, (
f"text dim {text_vectors[0].shape} != image dim {image_vectors[0].shape}"
)
assert (
image_vectors[0].shape == text_vectors[0].shape
), f"text dim {text_vectors[0].shape} != image dim {image_vectors[0].shape}"

13
tests/python/test_rag_reranker.py
View file

@ -27,7 +27,9 @@ def test_rerank_empty_returns_empty():
@pytest.mark.server
def test_rerank_reorders_by_relevance(monkeypatch):
"""Hide the relevant chunk at the back of the input and check it bubbles up."""
monkeypatch.setenv("UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2")
monkeypatch.setenv(
"UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2"
)
from core.rag.reranker import rerank, unload
from core.rag.retrieval import Hit
@ -35,7 +37,10 @@ def test_rerank_reorders_by_relevance(monkeypatch):
(Hit("noise1", 0.0), "Cats are small carnivorous mammals."),
(Hit("noise2", 0.0), "The Eiffel Tower is in Paris, France."),
(Hit("noise3", 0.0), "Python is a programming language."),
(Hit("answer", 0.0), "The speed of light in vacuum is approximately 299792458 meters per second."),
(
Hit("answer", 0.0),
"The speed of light in vacuum is approximately 299792458 meters per second.",
),
]
try:
ranked = rerank("How fast does light travel?", pairs, top_k = 2)
@ -47,7 +52,9 @@ def test_rerank_reorders_by_relevance(monkeypatch):
@pytest.mark.server
def test_unload_clears_singleton(monkeypatch):
monkeypatch.setenv("UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2")
monkeypatch.setenv(
"UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2"
)
from core.rag import reranker
from core.rag.retrieval import Hit

32
tests/python/test_rag_tool_handler.py
View file

@ -14,8 +14,10 @@ if str(STUDIO_BACKEND) not in sys.path:
def _make_hit(chunk_id: str):
"""Minimal stand-in for retrieval.Hit — just needs .chunk_id."""
class _Hit:
pass
h = _Hit()
h.chunk_id = chunk_id
h.score = 1.0
@ -55,9 +57,11 @@ def test_kb_takes_precedence_over_thread():
captured["scope"] = scope
return []
with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve):
with patch.object(
tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve,
):
result = tool.search_knowledge_base(
query = "x",
scope_kb_id = "kb-abc",
@ -78,9 +82,11 @@ def test_thread_scope_when_only_thread_set():
captured["scope"] = scope
return []
with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve):
with patch.object(
tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve,
):
tool.search_knowledge_base(
query = "x",
scope_thread_id = "thread-xyz",
@ -137,9 +143,17 @@ def test_execute_tool_dispatches_to_search_knowledge_base():
called = {}
def _stub(*, query, top_k = None, scope_kb_id = None, scope_thread_id = None,
enable_rerank = False, reranker_model = None, default_top_k = 5,
min_score = 0.0):
def _stub(
*,
query,
top_k = None,
scope_kb_id = None,
scope_thread_id = None,
enable_rerank = False,
reranker_model = None,
default_top_k = 5,
min_score = 0.0,
):
called["query"] = query
called["top_k"] = top_k
called["scope_kb_id"] = scope_kb_id