asklit/test_rag_scoping.py at main · SuffolkLITLab/asklit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
from asklit.db import get_connection, init_db
from asklit.rag import query_keyword_index, resolve_document_filter


def add_document(document_id, knowledgebase, filename, status, content):
    conn = get_connection()
    conn.execute(
        """
        INSERT INTO documents
            (id, knowledgebase, filename, file_path, file_type, file_size, status)
        VALUES (?, ?, ?, ?, '.txt', ?, ?)
        """,
        (
            document_id,
            knowledgebase,
            filename,
            f"data/uploads/{filename}",
            len(content),
            status,
        ),
    )
    conn.execute(
        """
        INSERT INTO document_chunks
            (document_id, chunk_index, content, page_number)
        VALUES (?, 0, ?, 1)
        """,
        (document_id, content),
    )
    conn.commit()
    conn.close()


def test_keyword_retrieval_is_scoped_by_knowledgebase_and_file(monkeypatch, tmp_path):
    db_path = tmp_path / "app.sqlite3"
    monkeypatch.setenv("ASKLIT_DB_PATH", str(db_path))
    init_db(str(db_path))
    add_document("housing-a", "housing", "repairs.txt", "indexed", "repair rights")
    add_document("housing-b", "housing", "rent.txt", "indexed", "rent rights")
    add_document("benefits-a", "benefits", "snap.txt", "indexed", "benefit rights")
    add_document("pending", "housing", "pending.txt", "indexing", "repair rights")

    housing_results = query_keyword_index("rights", knowledgebase="housing")
    selected_file_results = query_keyword_index(
        "rights",
        knowledgebase="housing",
        connected_files=["rent.txt"],
    )

    assert {result["metadata"]["document_id"] for result in housing_results} == {
        "housing-a",
        "housing-b",
    }
    assert [result["metadata"]["document_id"] for result in selected_file_results] == [
        "housing-b"
    ]
    assert resolve_document_filter("housing") == {"housing-a", "housing-b"}


def test_keyword_retrieval_can_use_an_explicit_scaffold_database(tmp_path):
    db_path = tmp_path / "scaffold" / "app.sqlite3"
    init_db(str(db_path))
    conn = get_connection(str(db_path))
    conn.execute(
        """
        INSERT INTO documents
            (id, knowledgebase, filename, file_path, file_type, file_size, status)
        VALUES ('lab-doc', 'lab', 'lab.txt', 'data/uploads/lab.txt', '.txt', 20, 'indexed')
        """
    )
    conn.execute(
        """
        INSERT INTO document_chunks
            (document_id, chunk_index, content, page_number)
        VALUES ('lab-doc', 0, 'experiment laboratory rights', 1)
        """
    )
    conn.commit()
    conn.close()

    results = query_keyword_index(
        "laboratory rights", knowledgebase="lab", db_path=str(db_path)
    )

    assert [result["metadata"]["document_id"] for result in results] == ["lab-doc"]