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.
107 lines
3.9 KiB
Python
107 lines
3.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
|
|
|
|
"""Caption figures with the loaded vision model and splice the text into the page
|
|
so images are searchable via the normal FTS5 + dense path. No-op (never raises)
|
|
without a vision model or on failure; gated by ``config.CAPTION_IMAGES``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import logging
|
|
|
|
from . import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CAPTION_PROMPT = (
|
|
"Describe this figure or image from a document in one or two concise "
|
|
"sentences, for search indexing. State what it depicts (e.g. a diagram, "
|
|
"chart, table or photo) and its key content. Do not add commentary."
|
|
)
|
|
|
|
|
|
def vision_endpoint() -> tuple[str, str] | None:
|
|
"""``(base_url, model)`` for a loaded vision GGUF model, else None."""
|
|
try:
|
|
from routes.inference import get_llama_cpp_backend
|
|
backend = get_llama_cpp_backend()
|
|
if getattr(backend, "is_loaded", False) and getattr(backend, "is_vision", False):
|
|
return backend.base_url, "local"
|
|
except Exception: # noqa: BLE001 - never let discovery break ingestion
|
|
return None
|
|
return None
|
|
|
|
|
|
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
|
|
import httpx
|
|
|
|
data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": _CAPTION_PROMPT},
|
|
{"type": "image_url", "image_url": {"url": data_url}},
|
|
],
|
|
}
|
|
],
|
|
"max_tokens": 200,
|
|
"temperature": 0.2,
|
|
"stream": False,
|
|
# Off: thinking models would spend the budget reasoning, returning "".
|
|
"chat_template_kwargs": {"enable_thinking": False},
|
|
}
|
|
try:
|
|
r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout)
|
|
r.raise_for_status()
|
|
text = r.json()["choices"][0]["message"]["content"]
|
|
return text.strip() or None
|
|
except Exception: # noqa: BLE001 - a failed caption is non-fatal
|
|
logger.debug("caption request failed", exc_info = True)
|
|
return None
|
|
|
|
|
|
def caption_images(
|
|
images: list, *, endpoint: tuple[str, str] | None = None
|
|
) -> dict[int, list[str]]:
|
|
"""Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when
|
|
disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``."""
|
|
if not config.CAPTION_IMAGES or not images:
|
|
return {}
|
|
ep = endpoint or vision_endpoint()
|
|
if ep is None:
|
|
return {}
|
|
base_url, model = ep
|
|
|
|
out: dict[int, list[str]] = {}
|
|
for img in images[: config.CAPTION_MAX_IMAGES]:
|
|
image_bytes = getattr(img, "image_bytes", None)
|
|
if not image_bytes:
|
|
continue
|
|
caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S)
|
|
if caption:
|
|
page = getattr(img, "page_number", None) or 0
|
|
out.setdefault(int(page), []).append(caption)
|
|
return out
|
|
|
|
|
|
def splice_captions(pages: list, captions: dict[int, list[str]]) -> list:
|
|
"""Append captions to their page's text so the chunker indexes them, keeping
|
|
figures attributable in retrieved chunks. Returns new ``Page`` objects."""
|
|
if not captions:
|
|
return pages
|
|
from .parsers import Page
|
|
|
|
out: list = []
|
|
for page in pages:
|
|
caps = captions.get(page.page_number or 0)
|
|
if not caps:
|
|
out.append(page)
|
|
continue
|
|
extra = "".join(f"\n\n[Figure on page {page.page_number}: {c}]" for c in caps)
|
|
text = page.text + extra
|
|
out.append(Page(text = text, page_number = page.page_number, char_count = len(text)))
|
|
return out
|