summaryrefslogtreecommitdiff
path: root/tests/test_clear_messages.py
blob: a0653ca31aa798e980e2c5e658047a29d7770a50 (plain)
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
86
87
88
89
90
91
92
93
94
"""Tests for clearing session messages and files."""

from __future__ import annotations

import tempfile

from ollama_web.sessions import SessionStore, is_valid_id


def make_store() -> SessionStore:
    tmp = tempfile.mkdtemp()
    return SessionStore(data_dir=tmp)


def test_clear_messages_removes_history_and_files():
    store = make_store()
    session = store.create(title="test")
    session_id = session["id"]
    assert is_valid_id(session_id)

    chat_file = store.add_file(
        session_id=session_id,
        name="notes.txt",
        data=b"hello",
        mime="text/plain",
        text="hello",
    )
    store.add_message(
        session_id=session_id,
        role="user",
        content="hi",
        attachments=[chat_file.id],
    )
    store.add_message(session_id=session_id, role="assistant", content="hello")

    files_dir = store._files_dir(session_id)
    stored_path = store.data_dir / chat_file.path
    assert stored_path.exists()
    assert (files_dir / f"{chat_file.id}.txt").exists()

    cleared = store.clear_messages(session_id)
    assert cleared is not None
    assert cleared["messages"] == []
    assert cleared["files"] == []
    assert not stored_path.exists()
    assert files_dir.exists()

    reloaded = store.get(session_id)
    assert reloaded is not None
    assert reloaded["messages"] == []
    assert reloaded["files"] == []


def test_clear_messages_invalid_session_id():
    store = make_store()
    assert store.clear_messages("not-a-valid-id") is None


def test_clear_messages_missing_session():
    store = make_store()
    assert store.clear_messages("0" * 32) is None


def test_add_message_stores_token_count_and_clears_it():
    store = make_store()
    session = store.create(title="test")
    session_id = session["id"]
    store.add_message(
        session_id=session_id,
        role="assistant",
        content="hello",
        token_count={
            "prompt_eval_count": 120,
            "eval_count": 80,
            "num_ctx": 8192,
            "total_count": 200,
        },
    )
    reloaded = store.get(session_id)
    assert reloaded is not None
    assert reloaded["messages"][0]["token_count"]["prompt_eval_count"] == 120
    assert reloaded["token_count"] == {
        "prompt_eval_count": 120,
        "eval_count": 80,
        "num_ctx": 8192,
        "total_count": 200,
    }

    cleared = store.clear_messages(session_id)
    assert cleared is not None
    assert "token_count" not in cleared
    reloaded_after = store.get(session_id)
    assert reloaded_after is not None
    assert "token_count" not in reloaded_after