arsam
6865ac0272
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import pytest
|
|
|
|
|
|
def test_chunk_text_single_short_text():
|
|
from axon.services.reranker import chunk_text
|
|
chunks = chunk_text("hello world", max_tokens=512)
|
|
assert chunks == ["hello world"]
|
|
|
|
|
|
def test_chunk_text_splits_long_text():
|
|
from axon.services.reranker import chunk_text
|
|
# 384 words per chunk at max_tokens=512 (512 * 0.75)
|
|
text = " ".join([f"word{i}" for i in range(500)])
|
|
chunks = chunk_text(text, max_tokens=512)
|
|
assert len(chunks) == 2
|
|
assert chunks[0].startswith("word0")
|
|
|
|
|
|
def test_chunk_text_empty_string():
|
|
from axon.services.reranker import chunk_text
|
|
chunks = chunk_text("", max_tokens=512)
|
|
assert chunks == [""]
|
|
|
|
|
|
def test_chunk_text_exact_boundary():
|
|
from axon.services.reranker import chunk_text
|
|
# 384 words should be exactly 1 chunk
|
|
text = " ".join(["x"] * 384)
|
|
chunks = chunk_text(text, max_tokens=512)
|
|
assert len(chunks) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rerank_results_sorts_by_score(monkeypatch):
|
|
import numpy as np
|
|
|
|
call_count = [0]
|
|
|
|
class FakeModel:
|
|
def predict(self, pairs):
|
|
# First pair gets low score, second gets high score
|
|
return np.array([0.1, 0.9])
|
|
|
|
monkeypatch.setattr("axon.services.reranker._model", FakeModel())
|
|
|
|
from axon.services.reranker import rerank_results
|
|
contents = [("https://low.com", "low relevance text"), ("https://high.com", "high relevance text")]
|
|
ranked = await rerank_results("query", contents)
|
|
assert ranked[0][0] == "https://high.com"
|
|
assert ranked[1][0] == "https://low.com"
|
|
assert ranked[0][2] > ranked[1][2]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rerank_results_empty_input():
|
|
from axon.services.reranker import rerank_results
|
|
result = await rerank_results("query", [])
|
|
assert result == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rerank_results_handles_empty_content(monkeypatch):
|
|
import numpy as np
|
|
|
|
class FakeModel:
|
|
def predict(self, pairs):
|
|
return np.array([0.5] * len(pairs))
|
|
|
|
monkeypatch.setattr("axon.services.reranker._model", FakeModel())
|
|
|
|
from axon.services.reranker import rerank_results
|
|
contents = [("https://a.com", ""), ("https://b.com", "some text")]
|
|
ranked = await rerank_results("query", contents)
|
|
assert len(ranked) == 2
|