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

Studio: per-thread RAG chunking/mode overrides

Browse source 
Threads can now opt into late chunking or multimodal mode independently
of the KBs they reference. Per-thread settings persist in chat_settings
under thread:<id>:rag and fall back to the app-level defaults when the
thread hasn't set anything explicitly.

Backend (routes/rag.py)
- ThreadRagSettings / UpdateThreadRagSettingsRequest Pydantic models.
- GET/PUT /api/rag/threads/{thread_id}/settings backed by
  chat_settings (upsert_chat_settings_merge). Same (multimodal, late)
  constraint enforcement as the create + defaults endpoints.
- POST /api/rag/threads/{thread_id}/reingest now accepts the same
  body shape — if any field is set, the new settings are persisted
  via set_thread_rag_settings BEFORE the reingest, so subsequent
  uploads pick up the change too.
- upload_thread_document reads the per-thread settings and passes
  them through to _start_ingestion, replacing the previous hard-coded
  ('standard', 'text', RAG_EMBEDDING_MODEL) defaults.

Frontend
- rag-api.ts: ThreadRagSettings type + getThreadRagSettings /
  setThreadRagSettings wrappers. reingestThreadDocuments now accepts
  optional UpdateThreadRagSettingsRequest opts.
- rag-store.ts: threadSettings map keyed by threadId, plus
  loadThreadSettings / updateThreadSettings actions. reingestThread
  refreshes the local settings copy when opts were supplied.
- chat-settings-sheet.tsx Retrieval section: when source = thread,
  shows side-by-side Mode + Chunking selects above the documents
  list. Selecting a different value:
    - persists immediately if the thread has no docs
    - prompts "Re-index N documents?" if docs exist; on Yes calls
      reingestThread with the new opts, on No reverts the select
  The (multimodal, late) constraint is enforced via per-option
  disabled + tooltip, matching the KB create dialog.
This commit is contained in:
Roland Tannous 2026-05-24 12:53:28 +04:00
commit ee1ff2bb50
4 changed files with 325 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -480,6 +480,89 @@ def set_rag_defaults(
)
class ThreadRagSettings(BaseModel):
chunking_strategy: ChunkingStrategy = "standard"
mode: KBMode = "text"
embedding_model: str | None = None
class UpdateThreadRagSettingsRequest(BaseModel):
chunking_strategy: ChunkingStrategy | None = None
mode: KBMode | None = None
embedding_model: str | None = None
def _thread_settings_key(thread_id: str) -> str:
return f"thread:{thread_id}:rag"
def _load_thread_settings(thread_id: str) -> ThreadRagSettings:
"""Per-thread RAG settings, falling back to app-level defaults.
Stored in chat_settings under "thread:<id>:rag" as a nested JSON
dict — same shape as RagDefaults.
"""
settings = list_chat_settings()
raw = settings.get(_thread_settings_key(thread_id)) or {}
if not isinstance(raw, dict):
raw = {}
fallback = _load_rag_defaults()
return ThreadRagSettings(
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,
)
@router.get(
"/threads/{thread_id}/settings",
response_model = ThreadRagSettings,
)
def get_thread_rag_settings(
thread_id: str,
current_subject: str = Depends(get_current_subject),
) -> ThreadRagSettings:
return _load_thread_settings(thread_id)
@router.put(
"/threads/{thread_id}/settings",
response_model = ThreadRagSettings,
)
def set_thread_rag_settings(
thread_id: str,
payload: UpdateThreadRagSettingsRequest,
current_subject: str = Depends(get_current_subject),
) -> ThreadRagSettings:
current = _load_thread_settings(thread_id)
new_strategy = payload.chunking_strategy or current.chunking_strategy
new_mode = payload.mode or current.mode
if payload.embedding_model is None:
new_embedder = current.embedding_model
elif payload.embedding_model.strip() == "":
new_embedder = None
else:
new_embedder = payload.embedding_model.strip()
_validate_mode_combo(new_mode, new_strategy)
upsert_chat_settings_merge(
{
_thread_settings_key(thread_id): {
"chunking_strategy": new_strategy,
"mode": new_mode,
"embedding_model": new_embedder,
}
}
)
return ThreadRagSettings(
chunking_strategy = new_strategy,
mode = new_mode,
embedding_model = new_embedder,
)
class ReingestKBRequest(BaseModel):
"""All fields optional — omitting one keeps the KB's current value."""
chunking_strategy: ChunkingStrategy | None = None
@ -619,20 +702,45 @@ def reingest_knowledge_base(
)
def reingest_thread_documents(
thread_id: str,
payload: UpdateThreadRagSettingsRequest | None = None,
current_subject: str = Depends(get_current_subject),
) -> ReingestResponse:
"""Rebuild a thread's RAG index using the current defaults.
"""Rebuild a thread's RAG index.
No body — per-thread strategy/mode overrides aren't exposed in v1.
Optional body lets the caller change the thread's chunking
strategy / mode / embedder at the same time — persisted into
chat_settings before re-ingestion so subsequent uploads pick up
the new values too. With an empty body, current settings are
reused.
"""
from utils.rag.config import RAG_EMBEDDING_MODEL
from utils.rag.config import resolve_embedder
if payload is None:
payload = UpdateThreadRagSettingsRequest()
if (
payload.chunking_strategy is not None
or payload.mode is not None
or payload.embedding_model is not None
):
# set_thread_rag_settings handles validation + persistence.
settings = set_thread_rag_settings(
thread_id,
payload,
current_subject = current_subject,
)
else:
settings = _load_thread_settings(thread_id)
embedder = settings.embedding_model or resolve_embedder(
settings.mode,
settings.chunking_strategy,
)
return _reingest_scope(
kb_id = None,
thread_id = thread_id,
chunking_strategy = "standard",
mode = "text",
embedding_model = RAG_EMBEDDING_MODEL,
chunking_strategy = settings.chunking_strategy,
mode = settings.mode,
embedding_model = embedder,
)
@ -696,12 +804,19 @@ async def upload_thread_document(
file: UploadFile,
current_subject: str = Depends(get_current_subject),
) -> UploadResponse:
from utils.rag.config import RAG_EMBEDDING_MODEL
from utils.rag.config import resolve_embedder
# Don't validate against chat_threads — a brand-new chat won't be
# persisted there until after the first runStart/runEnd. Users who
# attach a document on a fresh thread would otherwise hit a 404.
stored_path, filename, byte_size = await _save_upload(file)
# Per-thread settings fall back to app-level defaults inside the
# helper, so first-time-uploaded threads inherit user preferences.
settings = _load_thread_settings(thread_id)
embedder = settings.embedding_model or resolve_embedder(
settings.mode,
settings.chunking_strategy,
)
return _start_ingestion(
filename = filename,
stored_path = stored_path,
@ -709,7 +824,9 @@ async def upload_thread_document(
content_type = file.content_type,
kb_id = None,
thread_id = thread_id,
embedding_model = RAG_EMBEDDING_MODEL,
embedding_model = embedder,
chunking_strategy = settings.chunking_strategy,
mode = settings.mode,
)

121
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -96,6 +96,10 @@ import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog";
import { useKnowledgeBases } from "@/features/rag/hooks/use-knowledge-bases";
import { useThreadDocuments } from "@/features/rag/hooks/use-kb-documents";
import { useRagStore } from "@/features/rag/stores/rag-store";
import type {
ChunkingStrategy as RagChunkingStrategy,
KBMode,
} from "@/features/rag/api/rag-api";
import { Add01Icon, Delete02Icon } from "@hugeicons/core-free-icons";
function ragSourceLabel(
@ -448,6 +452,49 @@ export function ChatSettingsPanel({
);
const clearThreadIndex = useRagStore((s) => s.clearThreadIndex);
const reingestThread = useRagStore((s) => s.reingestThread);
const threadSettings = useRagStore((s) =>
activeThreadId ? s.threadSettings[activeThreadId] : undefined,
);
const loadThreadSettings = useRagStore((s) => s.loadThreadSettings);
const updateThreadSettings = useRagStore((s) => s.updateThreadSettings);
const ragDefaults = useRagStore((s) => s.defaults);
// Load this thread's RAG settings once when the sheet sees a thread
// for the first time. Updates re-render automatically via the store.
useEffect(() => {
if (ragSource.kind === "thread" && activeThreadId && !threadSettings) {
void loadThreadSettings(activeThreadId);
}
}, [ragSource.kind, activeThreadId, threadSettings, loadThreadSettings]);
const effectiveThreadChunking: RagChunkingStrategy =
threadSettings?.chunking_strategy ??
ragDefaults?.chunking_strategy ??
"standard";
const effectiveThreadMode: KBMode =
threadSettings?.mode ?? ragDefaults?.mode ?? "text";
const applyThreadSettingChange = (
patch: { chunking_strategy?: RagChunkingStrategy; mode?: KBMode },
) => {
if (!activeThreadId) return;
if (threadDocs.length === 0) {
// No existing chunks to invalidate — just persist.
void updateThreadSettings(activeThreadId, patch);
return;
}
const ok = window.confirm(
`Re-index ${threadDocs.length} document${threadDocs.length === 1 ? "" : "s"} ` +
`with the new settings? Existing chunks will be deleted and rebuilt.`,
);
if (ok) {
void reingestThread(activeThreadId, patch);
} else {
// User declined — refresh the store so the select snaps back
// to the unchanged settings.
void loadThreadSettings(activeThreadId);
}
};
const [kbCreateOpen, setKbCreateOpen] = useState(false);
const ragEnabled = ragSource.kind !== "off";
const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null;
@ -1311,7 +1358,76 @@ export function ChatSettingsPanel({
onCreated={(kb) => setRagSource({ kind: "kb", kbId: kb.id })}
/>
{ragSource.kind === "thread" && activeThreadId ? (
<div className="flex flex-col gap-1.5">
<>
<div className="grid grid-cols-2 gap-2">
<div className="flex flex-col gap-1">
<label className="text-[11px] font-medium text-muted-foreground">
Mode
</label>
<Select
value={effectiveThreadMode}
onValueChange={(v) => {
const next = v as KBMode;
if (next === effectiveThreadMode) return;
applyThreadSettingChange({ mode: next });
}}
>
<SelectTrigger className="h-8 text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="text">Text only</SelectItem>
<SelectItem
value="multimodal"
disabled={effectiveThreadChunking === "late"}
title={
effectiveThreadChunking === "late"
? "Multimodal cannot be combined with late chunking"
: undefined
}
>
Multimodal
</SelectItem>
</SelectContent>
</Select>
</div>
<div className="flex flex-col gap-1">
<label className="text-[11px] font-medium text-muted-foreground">
Chunking
</label>
<Select
value={effectiveThreadChunking}
onValueChange={(v) => {
const next = v as RagChunkingStrategy;
if (next === effectiveThreadChunking) return;
applyThreadSettingChange({ chunking_strategy: next });
}}
>
<SelectTrigger className="h-8 text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="standard">Standard</SelectItem>
<SelectItem
value="late"
disabled={effectiveThreadMode === "multimodal"}
title={
effectiveThreadMode === "multimodal"
? "Late chunking cannot be combined with multimodal mode"
: undefined
}
>
Late
</SelectItem>
</SelectContent>
</Select>
</div>
</div>
<p className="text-[11px] text-muted-foreground">
Changing either setting will re-index this thread's existing
documents.
</p>
<div className="flex flex-col gap-1.5">
<label className="text-[12px] font-medium text-muted-foreground">
Documents in this thread
</label>
@ -1368,7 +1484,8 @@ export function ChatSettingsPanel({
</div>
</>
)}
</div>
</div>
</>
) : null}
<div className="flex flex-col gap-1.5">
<div className="flex items-center justify-between">

39
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -232,15 +232,52 @@ export async function reingestKnowledgeBase(
return parseJsonOrThrow<ReingestResponse>(response);
}
export interface ThreadRagSettings {
chunking_strategy: ChunkingStrategy;
mode: KBMode;
embedding_model: string | null;
}
export interface UpdateThreadRagSettingsRequest {
chunking_strategy?: ChunkingStrategy;
mode?: KBMode;
embedding_model?: string | null;
}
export async function getThreadRagSettings(
threadId: string,
): Promise<ThreadRagSettings> {
const response = await authFetch(
`/api/rag/threads/${encodeURIComponent(threadId)}/settings`,
);
return parseJsonOrThrow<ThreadRagSettings>(response);
}
export async function setThreadRagSettings(
threadId: string,
payload: UpdateThreadRagSettingsRequest,
): Promise<ThreadRagSettings> {
const response = await authFetch(
`/api/rag/threads/${encodeURIComponent(threadId)}/settings`,
{
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
},
);
return parseJsonOrThrow<ThreadRagSettings>(response);
}
export async function reingestThreadDocuments(
threadId: string,
opts: UpdateThreadRagSettingsRequest = {},
): Promise<ReingestResponse> {
const response = await authFetch(
`/api/rag/threads/${encodeURIComponent(threadId)}/reingest`,
{
method: "POST",
headers: { "Content-Type": "application/json" },
body: "{}",
body: JSON.stringify(opts),
},
);
return parseJsonOrThrow<ReingestResponse>(response);

46
studio/frontend/src/features/rag/stores/rag-store.ts
View file

@ -9,6 +9,7 @@ import {
deleteDocument as apiDeleteDocument,
deleteKnowledgeBase as apiDeleteKB,
getRagDefaults as apiGetRagDefaults,
getThreadRagSettings as apiGetThreadSettings,
type JobEvent,
type KnowledgeBase,
listKBDocuments,
@ -21,9 +22,12 @@ import {
reingestKnowledgeBase as apiReingestKB,
reingestThreadDocuments as apiReingestThread,
setRagDefaults as apiSetRagDefaults,
setThreadRagSettings as apiSetThreadSettings,
subscribeToJobEvents,
type ThreadIndexSummary,
type ThreadRagSettings,
type UpdateRagDefaultsRequest,
type UpdateThreadRagSettingsRequest,
uploadKBDocument,
uploadThreadDocument,
} from "../api/rag-api";
@ -59,12 +63,22 @@ interface RagStoreState {
clearThreadIndex: (threadId: string) => Promise<void>;
reingestKB: (kbId: string, opts?: ReingestKBOptions) => Promise<string[]>;
reingestThread: (threadId: string) => Promise<string[]>;
reingestThread: (
threadId: string,
opts?: UpdateThreadRagSettingsRequest,
) => Promise<string[]>;
defaults: RagDefaults | null;
loadDefaults: () => Promise<void>;
updateDefaults: (patch: UpdateRagDefaultsRequest) => Promise<void>;
threadSettings: Record<string, ThreadRagSettings>;
loadThreadSettings: (threadId: string) => Promise<void>;
updateThreadSettings: (
threadId: string,
patch: UpdateThreadRagSettingsRequest,
) => Promise<ThreadRagSettings>;
subscribeJob: (jobId: string, onComplete?: () => void) => void;
}
@ -92,6 +106,8 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
defaults: null,
threadSettings: {},
async loadKnowledgeBases() {
set({ kbsLoading: true, kbsError: null });
try {
@ -269,10 +285,15 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
set({ defaults });
},
async reingestThread(threadId) {
const response = await apiReingestThread(threadId);
async reingestThread(threadId, opts) {
const response = await apiReingestThread(threadId, opts ?? {});
void get().loadThreadDocuments(threadId);
void get().loadThreadIndexes();
if (opts) {
// The reingest endpoint persists the new settings as a side
// effect; refresh the local copy so the UI reflects them.
void get().loadThreadSettings(threadId);
}
for (const jobId of response.job_ids) {
get().subscribeJob(jobId, () => {
void get().loadThreadDocuments(threadId);
@ -281,6 +302,25 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
return response.job_ids;
},
async loadThreadSettings(threadId) {
try {
const settings = await apiGetThreadSettings(threadId);
set((state) => ({
threadSettings: { ...state.threadSettings, [threadId]: settings },
}));
} catch {
// Best-effort — falls back to defaults UI-side when missing.
}
},
async updateThreadSettings(threadId, patch) {
const settings = await apiSetThreadSettings(threadId, patch);
set((state) => ({
threadSettings: { ...state.threadSettings, [threadId]: settings },
}));
return settings;
},
subscribeJob(jobId, onComplete) {
const existing = get().jobUnsubscribers[jobId];
if (existing) return;