mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-05 07:39:47 +00:00
17fb9960e1
* fix(search): keep zero float16 cosine scores finite Updated the handling of zero-norm rows to use a representable denominator of 1 instead of 1e-10. * test(search): cover zero cosine vectors and batch boundaries Added tests for batch cosine similarity with zero vectors and empty batches. * test(search): keep zero embedding rankings finite Add tests for zero embeddings in vector search and legacy JSON loading.
297 lines
9.8 KiB
Python
297 lines
9.8 KiB
Python
"""Tests for small utilities in ``reme.utils``."""
|
|
|
|
import asyncio
|
|
import sys
|
|
import threading
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from reme.utils import common_utils
|
|
from reme.utils.counter import (
|
|
COUNTER_LOCK_KEY,
|
|
COUNTER_TREE_KEY,
|
|
global_counter_add,
|
|
global_counter_add_many,
|
|
global_counter_get,
|
|
global_counter_get_all,
|
|
global_counter_inc,
|
|
)
|
|
from reme.utils.similarity_utils import batch_cosine_similarity, cosine_similarity
|
|
|
|
|
|
def test_batch_cosine_similarity_rejects_1d_inputs():
|
|
"""1D vectors should fail with a clear validation error, not IndexError."""
|
|
with pytest.raises(ValueError, match="Expected 2D arrays"):
|
|
batch_cosine_similarity(np.array([1.0, 0.0]), np.array([[1.0, 0.0]]))
|
|
|
|
|
|
def test_batch_cosine_similarity_pairwise_matrix():
|
|
"""Batch cosine returns the full pairwise matrix for valid 2D inputs."""
|
|
result = batch_cosine_similarity(
|
|
np.array([[1.0, 0.0], [0.0, 1.0]]),
|
|
np.array([[1.0, 0.0], [1.0, 1.0]]),
|
|
)
|
|
|
|
assert result.shape == (2, 2)
|
|
np.testing.assert_allclose(result[0], [1.0, 2**-0.5])
|
|
np.testing.assert_allclose(result[1], [0.0, 2**-0.5])
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
|
|
def test_batch_cosine_similarity_zero_vectors(dtype):
|
|
"""Zero rows and columns have finite zero scores, including stored float16 embeddings."""
|
|
left = np.array([[0, 0], [1, 0], [0, 1]], dtype=dtype)
|
|
right = np.array([[1, 0], [0, 0], [-1, 0]], dtype=dtype)
|
|
|
|
with np.errstate(invalid="raise", divide="raise"):
|
|
result = batch_cosine_similarity(left, right)
|
|
reverse = batch_cosine_similarity(right, left)
|
|
zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), right)
|
|
all_zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), np.zeros((3, 2), dtype=dtype))
|
|
|
|
assert result.dtype == dtype
|
|
np.testing.assert_array_equal(result, [[0, 0, 0], [1, 0, -1], [0, 0, 0]])
|
|
np.testing.assert_array_equal(reverse, result.T)
|
|
np.testing.assert_array_equal(zeros, np.zeros((2, 3)))
|
|
np.testing.assert_array_equal(all_zeros, np.zeros((2, 3)))
|
|
|
|
|
|
@pytest.mark.parametrize("left_rows,right_rows", [(0, 2), (2, 0), (0, 0)])
|
|
def test_batch_cosine_similarity_empty_batches(left_rows, right_rows):
|
|
"""Empty batches retain the pairwise output shape."""
|
|
result = batch_cosine_similarity(
|
|
np.zeros((left_rows, 2), dtype=np.float16),
|
|
np.zeros((right_rows, 2), dtype=np.float16),
|
|
)
|
|
assert result.shape == (left_rows, right_rows)
|
|
assert result.dtype == np.float16
|
|
|
|
|
|
def test_batch_cosine_similarity_rejects_mismatched_dimensions():
|
|
"""Batched embeddings must share their feature dimension."""
|
|
with pytest.raises(ValueError, match="Embedding dimensions must match"):
|
|
batch_cosine_similarity(np.zeros((2, 2)), np.zeros((2, 3)))
|
|
|
|
|
|
def test_cosine_similarity_rejects_mismatched_lengths():
|
|
"""Single-vector cosine validates dimensions before computing."""
|
|
with pytest.raises(ValueError, match="Vectors must have same length"):
|
|
cosine_similarity([1.0], [1.0, 2.0])
|
|
|
|
|
|
def test_mock_reme_server_uses_reme_entrypoint(monkeypatch):
|
|
"""The test server helper should spawn the reme CLI module, not legacy reme."""
|
|
captured: dict[str, list[str]] = {}
|
|
|
|
class DummyProcess:
|
|
"""Process stub returned by the patched Popen."""
|
|
|
|
stdout = None
|
|
|
|
def poll(self):
|
|
"""Return a successful process status."""
|
|
return 0
|
|
|
|
def fake_popen(cmd, **_kwargs):
|
|
"""Capture the spawned command."""
|
|
captured["cmd"] = cmd
|
|
return DummyProcess()
|
|
|
|
async def fake_wait_ready(_host, _port, _timeout):
|
|
return None
|
|
|
|
monkeypatch.setattr(common_utils.subprocess, "Popen", fake_popen)
|
|
monkeypatch.setattr(common_utils, "_wait_reme_ready", fake_wait_ready)
|
|
|
|
async def run():
|
|
async with common_utils.mock_reme_server(port=45678, log_to_file=False, enable_logo=False):
|
|
pass
|
|
|
|
asyncio.run(run())
|
|
|
|
assert captured["cmd"][:4] == [sys.executable, "-m", "reme.reme", "start"]
|
|
|
|
|
|
def test_inc_returns_old_value_starting_at_zero():
|
|
"""First call returns 0, then values increase by 1 per call."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_inc(metadata, ["a"]) == 0
|
|
assert global_counter_inc(metadata, ["a"]) == 1
|
|
assert global_counter_inc(metadata, ["a"]) == 2
|
|
|
|
|
|
def test_add_returns_old_value_and_adds_val():
|
|
"""``add`` is fetch-and-add: old value out, ``val`` added in."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_add(metadata, ["a"], 10) == 0
|
|
assert global_counter_add(metadata, ["a"], 5) == 10
|
|
assert global_counter_inc(metadata, ["a"]) == 15
|
|
assert global_counter_get(metadata, ["a"]) == 16
|
|
|
|
|
|
def test_add_many_returns_old_values_and_updates_all_paths():
|
|
"""``add_many`` updates sibling metrics under one counter-tree lock."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["usage", "input"], 2)
|
|
|
|
previous = global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): 4,
|
|
("usage", "total"): 7,
|
|
},
|
|
)
|
|
|
|
assert previous == {
|
|
("usage", "input"): 2,
|
|
("usage", "output"): 0,
|
|
("usage", "total"): 0,
|
|
}
|
|
assert global_counter_get(metadata, ["usage", "input"]) == 5
|
|
assert global_counter_get(metadata, ["usage", "output"]) == 4
|
|
assert global_counter_get(metadata, ["usage", "total"]) == 7
|
|
|
|
|
|
def test_add_many_holds_the_counter_lock_once_for_the_batch():
|
|
"""A multi-metric update is one critical section, not several writes."""
|
|
|
|
class CountingLock:
|
|
"""Lock test double that counts entered critical sections."""
|
|
|
|
def __init__(self):
|
|
self.lock = threading.Lock()
|
|
self.entries = 0
|
|
|
|
def __enter__(self):
|
|
self.lock.acquire()
|
|
self.entries += 1
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
self.lock.release()
|
|
|
|
lock = CountingLock()
|
|
metadata = {COUNTER_LOCK_KEY: lock}
|
|
|
|
global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): 4,
|
|
("usage", "total"): 7,
|
|
},
|
|
)
|
|
|
|
assert lock.entries == 1
|
|
|
|
|
|
def test_add_many_validates_every_update_before_mutating():
|
|
"""One invalid update leaves every valid counter unchanged."""
|
|
metadata: dict = {}
|
|
|
|
with pytest.raises(TypeError, match="increments must be integers"):
|
|
global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): "invalid",
|
|
},
|
|
)
|
|
|
|
assert COUNTER_TREE_KEY not in metadata
|
|
|
|
|
|
def test_counters_are_isolated_by_key_path():
|
|
"""Sibling and nested keys, plus the root, hold independent counters."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_inc(metadata, ["a"]) == 0
|
|
assert global_counter_inc(metadata, ["b"]) == 0
|
|
assert global_counter_inc(metadata, ["a", "child"]) == 0
|
|
assert global_counter_inc(metadata, []) == 0
|
|
|
|
assert global_counter_get(metadata, ["a"]) == 1
|
|
assert global_counter_get(metadata, ["b"]) == 1
|
|
assert global_counter_get(metadata, ["a", "child"]) == 1
|
|
assert global_counter_get(metadata, []) == 1
|
|
|
|
|
|
def test_get_does_not_create_missing_nodes():
|
|
"""``get`` reports 0 for missing paths and leaves the tree untouched."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_get(metadata, ["missing"]) == 0
|
|
assert COUNTER_TREE_KEY not in metadata
|
|
|
|
global_counter_inc(metadata, ["a"])
|
|
assert global_counter_get(metadata, ["a", "missing"]) == 0
|
|
assert "missing" not in metadata[COUNTER_TREE_KEY]["children"]["a"]["children"]
|
|
|
|
|
|
def test_get_all_returns_none_for_missing_key():
|
|
"""``get_all`` returns None when the tree or the path does not exist."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_get_all(metadata, []) is None
|
|
assert global_counter_get_all(metadata, ["missing"]) is None
|
|
|
|
global_counter_inc(metadata, ["a"])
|
|
assert global_counter_get_all(metadata, ["missing"]) is None
|
|
assert global_counter_get_all(metadata, ["a", "missing"]) is None
|
|
|
|
|
|
def test_get_all_returns_subtree_and_whole_tree():
|
|
"""``get_all`` returns the node at ``key``; an empty key returns the root."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["a"], 2)
|
|
global_counter_add(metadata, ["a", "child"], 3)
|
|
|
|
subtree = global_counter_get_all(metadata, ["a"])
|
|
assert subtree == {"value": 2, "children": {"child": {"value": 3, "children": {}}}}
|
|
|
|
root = global_counter_get_all(metadata, [])
|
|
assert root["value"] == 0
|
|
assert root["children"]["a"] == subtree
|
|
|
|
|
|
def test_get_all_returns_deep_copy():
|
|
"""Mutating the returned subtree must not affect the live counter tree."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["a", "child"], 3)
|
|
|
|
subtree = global_counter_get_all(metadata, ["a"])
|
|
subtree["value"] = 999
|
|
subtree["children"]["child"]["value"] = 999
|
|
subtree["children"]["extra"] = {"value": 1, "children": {}}
|
|
|
|
assert global_counter_get(metadata, ["a"]) == 0
|
|
assert global_counter_get(metadata, ["a", "child"]) == 3
|
|
assert global_counter_get_all(metadata, ["a", "extra"]) is None
|
|
|
|
|
|
def test_concurrent_inc_yields_unique_values():
|
|
"""Parallel increments on one key never return duplicate values."""
|
|
metadata: dict = {}
|
|
results: list[int] = []
|
|
results_lock = threading.Lock()
|
|
calls_per_thread = 200
|
|
|
|
def worker():
|
|
for _ in range(calls_per_thread):
|
|
value = global_counter_inc(metadata, ["shared"])
|
|
with results_lock:
|
|
results.append(value)
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(8)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
total = len(threads) * calls_per_thread
|
|
assert sorted(results) == list(range(total))
|
|
assert global_counter_get(metadata, ["shared"]) == total
|