8848a310dfAdds a self-contained RAG stack to Studio: knowledge bases with chunked indexing, hybrid (dense + lexical) retrieval, and an automatic first-pass context inject into chat. Embeddings run through a local llama-server GGUF backend (default unsloth/bge-small-en-v1.5-GGUF) with a sentence-transformers fallback. The chat tool loop gains a search_knowledge_base tool, a per-turn re-search cap, and source citation, layered on top of the shared ToolLoopController.
96 lines
3.3 KiB
Python
96 lines
3.3 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
|
|
|
|
"""Lexical (FTS5) + dense (vec0 cosine) retrieval fused via Reciprocal Rank
|
|
Fusion. ``dense_score`` is carried so callers can apply a similarity floor."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
|
|
from . import config, embeddings, store
|
|
|
|
|
|
@dataclass
|
|
class Hit:
|
|
chunk_id: str
|
|
score: float
|
|
lexical_score: float | None = None
|
|
dense_score: float | None = None
|
|
|
|
|
|
def retrieve_lexical(
|
|
conn: sqlite3.Connection,
|
|
scope: str,
|
|
query: str,
|
|
k: int | None = None,
|
|
) -> list[Hit]:
|
|
k = k or config.TOP_K_LEXICAL
|
|
return [Hit(cid, s, lexical_score = s) for cid, s in store.search_lexical(conn, scope, query, k)]
|
|
|
|
|
|
def retrieve_dense(
|
|
conn: sqlite3.Connection,
|
|
scope: str,
|
|
query: str,
|
|
k: int | None = None,
|
|
*,
|
|
model_name: str | None = None,
|
|
) -> list[Hit]:
|
|
k = k or config.TOP_K_DENSE
|
|
vec = embeddings.encode([query], model_name = model_name, normalize = True)[0]
|
|
return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)]
|
|
|
|
|
|
def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]:
|
|
fused: dict[str, float] = {}
|
|
best: dict[str, Hit] = {}
|
|
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)
|
|
cur = best.get(hit.chunk_id)
|
|
if cur is None:
|
|
best[hit.chunk_id] = Hit(hit.chunk_id, 0.0, hit.lexical_score, hit.dense_score)
|
|
else:
|
|
cur.lexical_score = (
|
|
cur.lexical_score if cur.lexical_score is not None else hit.lexical_score
|
|
)
|
|
cur.dense_score = (
|
|
cur.dense_score if cur.dense_score is not None else hit.dense_score
|
|
)
|
|
out: list[Hit] = []
|
|
for cid, s in sorted(fused.items(), key = lambda kv: kv[1], reverse = True)[:top_k]:
|
|
h = best[cid]
|
|
h.score = s
|
|
out.append(h)
|
|
return out
|
|
|
|
|
|
def retrieve_hybrid(
|
|
conn: sqlite3.Connection,
|
|
scope: str,
|
|
query: str,
|
|
*,
|
|
k: int | None = None,
|
|
model_name: str | None = None,
|
|
mode: str = "hybrid",
|
|
) -> list[Hit]:
|
|
"""``mode`` picks the backend: lexical-only, dense-only, or RRF of both
|
|
(default). Pool sizes and the RRF constant come from config."""
|
|
k = k if k is not None else config.TOP_K_HYBRID
|
|
k = int(k) # tool-call / scope top_k may arrive as a float; LIMIT + slice need int
|
|
if mode == "lexical":
|
|
return retrieve_lexical(conn, scope, query, k)
|
|
if mode == "dense":
|
|
return retrieve_dense(conn, scope, query, k, model_name = model_name)
|
|
lexical = retrieve_lexical(conn, scope, query, config.TOP_K_LEXICAL)
|
|
dense = retrieve_dense(conn, scope, query, config.TOP_K_DENSE, model_name = model_name)
|
|
return _rrf([lexical, dense], config.RRF_K, k)
|
|
|
|
|
|
def filter_min_score(hits: list[Hit], min_score: float) -> list[Hit]:
|
|
"""Cosine floor; gates only hits with a dense_score (lexical-only pass)."""
|
|
if min_score <= 0:
|
|
return hits
|
|
return [h for h in hits if h.dense_score is None or h.dense_score >= min_score]
|