86 lines
2.5 KiB
Python
86 lines
2.5 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 uuid
|
|
|
|
import storage.studio_db as studio_db
|
|
|
|
|
|
def _uid() -> str:
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
def test_locator_schema_is_additive_and_nullable(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
|
monkeypatch.setattr(studio_db, "_schema_ready", False)
|
|
|
|
with studio_db.get_connection() as conn:
|
|
chunk_cols = {
|
|
row["name"] for row in conn.execute("PRAGMA table_info(rag_chunks)")
|
|
}
|
|
assert {
|
|
"source_page_index",
|
|
"page_char_start",
|
|
"page_char_end",
|
|
"line_start",
|
|
"line_end",
|
|
}.issubset(chunk_cols)
|
|
|
|
page_cols = {
|
|
row["name"] for row in conn.execute("PRAGMA table_info(rag_document_pages)")
|
|
}
|
|
assert {
|
|
"document_id",
|
|
"page_index",
|
|
"page_number",
|
|
"text",
|
|
"char_count",
|
|
"line_count",
|
|
}.issubset(page_cols)
|
|
|
|
kb_id = _uid()
|
|
doc_id = _uid()
|
|
chunk_id = _uid()
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO rag_knowledge_bases
|
|
(id, name, embedding_model, owner_user_id, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(kb_id, "KB", "embedder", "alice", 1_700_000_000),
|
|
)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO rag_documents
|
|
(id, kb_id, thread_id, filename, content_type, stored_path, status,
|
|
num_chunks, byte_size, created_at)
|
|
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?)
|
|
""",
|
|
(doc_id, kb_id, "old.pdf", "application/pdf", "old.pdf", 1_700_000_001),
|
|
)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO rag_chunks
|
|
(id, document_id, chunk_index, text, token_count, page_number)
|
|
VALUES (?, ?, 0, ?, 3, 1)
|
|
""",
|
|
(chunk_id, doc_id, "legacy chunk"),
|
|
)
|
|
|
|
row = conn.execute(
|
|
"""
|
|
SELECT source_page_index, page_char_start, page_char_end,
|
|
line_start, line_end
|
|
FROM rag_chunks WHERE id = ?
|
|
""",
|
|
(chunk_id,),
|
|
).fetchone()
|
|
assert dict(row) == {
|
|
"source_page_index": None,
|
|
"page_char_start": None,
|
|
"page_char_end": None,
|
|
"line_start": None,
|
|
"line_end": None,
|
|
}
|