File size: 5,531 Bytes
7ab7df1
 
b39a3c2
7ab7df1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ed005f8
b39a3c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ed005f8
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
"""Tests for utils module."""

from langchain_core.documents import Document
from langchain_core.messages import AIMessage, HumanMessage

from utils import (
    format_chat_history,
    format_context_with_highlight,
    messages_to_tuples,
)


def test_format_chat_history_empty():
    """Test formatting empty chat history."""
    result = format_chat_history([])
    assert result == ""


def test_format_chat_history_tuples(sample_chat_history_tuples):
    """Test formatting chat history as tuples."""
    result = format_chat_history(sample_chat_history_tuples)
    assert "Human:" in result
    assert "Assistant:" in result
    assert "What is RAG?" in result


def test_format_chat_history_dicts(sample_chat_history):
    """Test formatting chat history as dicts."""
    result = format_chat_history(sample_chat_history)
    assert "Human:" in result
    assert "Assistant:" in result
    assert "What is RAG?" in result


def test_format_chat_history_messages():
    """Test formatting chat history as Message objects."""
    history = [
        HumanMessage(content="Hello"),
        AIMessage(content="Hi there"),
    ]
    result = format_chat_history(history)
    assert "Human:" in result
    assert "Assistant:" in result
    assert "Hello" in result
    assert "Hi there" in result


def test_format_chat_history_limit():
    """Test that chat history respects limit."""
    history = [
        {"role": "user", "content": f"Message {i}"}
        for i in range(10)
    ]
    result = format_chat_history(history, limit=5)
    # Should only include last 5 messages
    assert "Message 5" in result
    assert "Message 9" in result
    assert "Message 0" not in result


def test_messages_to_tuples_empty():
    """Test converting empty messages to tuples."""
    result = messages_to_tuples([])
    assert result == []


def test_messages_to_tuples(sample_chat_history):
    """Test converting messages to tuples."""
    result = messages_to_tuples(sample_chat_history)
    assert len(result) == 1  # Only one complete pair
    assert result[0][0] == "What is RAG?"
    assert result[0][1] == "RAG is Retrieval-Augmented Generation."


def test_messages_to_tuples_incomplete():
    """Test converting messages with incomplete pairs."""
    messages = [
        {"role": "user", "content": "Question 1"},
        {"role": "user", "content": "Question 2"},
    ]
    result = messages_to_tuples(messages)
    assert len(result) == 0  # No complete pairs


def test_format_context_with_highlight_empty():
    """Test formatting empty context."""
    result = format_context_with_highlight([])
    assert result == ""


def test_format_context_with_highlight(sample_documents):
    """Test formatting context with documents."""
    result = format_context_with_highlight(sample_documents[:2])
    assert "Sources:" in result
    assert "test1.pdf" in result
    assert "test2.pdf" in result
    assert "machine learning" in result


def test_format_context_with_scores(sample_documents):
    """Test formatting context with scores."""
    docs_with_scores = [
        (sample_documents[0], 0.9),
        (sample_documents[1], 0.8),
    ]
    result = format_context_with_highlight(
        sample_documents[:2], docs_with_scores=docs_with_scores
    )
    assert "⭐" in result  # Top chunk should be highlighted
    assert "rel:" in result or "dist:" in result  # Score info


def test_format_context_with_rewritten_query(sample_documents):
    """Test formatting context with rewritten query."""
    result = format_context_with_highlight(
        sample_documents[:2], rewritten_query="rewritten query"
    )
    assert "πŸ”„ Rewritten:" in result
    assert "rewritten query" in result


def test_format_context_with_hybrid_scores(sample_documents):
    """Test formatting context with hybrid scores."""
    hybrid_scores = [
        (sample_documents[0], 0.9, 0.85, 0.95),
        (sample_documents[1], 0.8, 0.75, 0.85),
    ]
    result = format_context_with_highlight(
        sample_documents[:2], hybrid_scores=hybrid_scores
    )
    assert "f:" in result  # Fused score
    assert "s:" in result  # Semantic score
    assert "k:" in result  # Keyword score


# Tests merged from test_utils_edge_cases.py


def test_format_context_score_edge_cases():
    """Test format_context_with_highlight with edge case scores."""
    docs = [
        Document(
            page_content="Test content",
            metadata={"source": "pdf/test.pdf", "page": 1},
        )
    ]

    # Test with score > 1.0 (distance score)
    docs_with_scores = [(docs[0], 1.5)]
    result = format_context_with_highlight(docs, docs_with_scores=docs_with_scores)
    assert "dist:" in result

    # Test with score exactly 1.0
    docs_with_scores = [(docs[0], 1.0)]
    result = format_context_with_highlight(docs, docs_with_scores=docs_with_scores)
    assert "rel:" in result

    # Test with None score
    docs_with_scores = [(docs[0], None)]
    result = format_context_with_highlight(docs, docs_with_scores=docs_with_scores)
    assert "⭐" in result  # Should still highlight first chunk


def test_format_context_no_scores():
    """Test format_context_with_highlight with no scores."""
    docs = [
        Document(
            page_content="Test content",
            metadata={"source": "pdf/test.pdf", "page": 1},
        )
    ]

    result = format_context_with_highlight(docs, docs_with_scores=None)
    assert "Test content" in result
    assert "⭐" in result  # First chunk should be highlighted by default