Spaces:
Runtime error
Runtime error
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
|