mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
Prathap PandGitHub
c2756964a7
* fix(channels): wire channel→agent handler and fix Telegram send pipeline * format code * add supported tests
225 lines
6.9 KiB
Python
225 lines
6.9 KiB
Python
"""Tests for the Hybrid RRF memory backend."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from openjarvis.core.events import EventBus, EventType
|
|
from openjarvis.core.registry import MemoryRegistry
|
|
from openjarvis.tools.storage._stubs import MemoryBackend, RetrievalResult
|
|
from openjarvis.tools.storage.hybrid import (
|
|
HybridMemory,
|
|
reciprocal_rank_fusion,
|
|
)
|
|
|
|
# -- Fake in-memory backend for testing -----------------------------------
|
|
|
|
|
|
class _FakeBackend(MemoryBackend):
|
|
"""Minimal in-memory backend for testing hybrid logic."""
|
|
|
|
backend_id: str = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self._docs: Dict[str, tuple] = {}
|
|
|
|
def store(
|
|
self,
|
|
content: str,
|
|
*,
|
|
source: str = "",
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
doc_id = uuid.uuid4().hex
|
|
self._docs[doc_id] = (content, source, metadata or {})
|
|
return doc_id
|
|
|
|
def retrieve(
|
|
self,
|
|
query: str,
|
|
*,
|
|
top_k: int = 5,
|
|
**kwargs: Any,
|
|
) -> List[RetrievalResult]:
|
|
# Simple substring match with position-based scoring
|
|
results = []
|
|
for doc_id, (content, source, meta) in self._docs.items():
|
|
if query.lower() in content.lower():
|
|
results.append(
|
|
RetrievalResult(
|
|
content=content,
|
|
score=1.0,
|
|
source=source,
|
|
metadata=meta,
|
|
)
|
|
)
|
|
return results[:top_k]
|
|
|
|
def delete(self, doc_id: str) -> bool:
|
|
if doc_id in self._docs:
|
|
del self._docs[doc_id]
|
|
return True
|
|
return False
|
|
|
|
def clear(self) -> None:
|
|
self._docs.clear()
|
|
|
|
|
|
def _make_hybrid() -> HybridMemory:
|
|
"""Create a HybridMemory with two fake backends."""
|
|
if not MemoryRegistry.contains("hybrid"):
|
|
MemoryRegistry.register_value("hybrid", HybridMemory)
|
|
sparse = _FakeBackend()
|
|
dense = _FakeBackend()
|
|
return HybridMemory(sparse=sparse, dense=dense)
|
|
|
|
|
|
# -- RRF unit tests -------------------------------------------------------
|
|
|
|
|
|
def test_rrf_scoring_correctness():
|
|
"""Verify RRF formula: score = sum(w / (k + rank + 1))."""
|
|
list1 = [
|
|
RetrievalResult(content="A", score=10.0),
|
|
RetrievalResult(content="B", score=5.0),
|
|
]
|
|
list2 = [
|
|
RetrievalResult(content="B", score=8.0),
|
|
RetrievalResult(content="C", score=3.0),
|
|
]
|
|
fused = reciprocal_rank_fusion([list1, list2], k=60)
|
|
|
|
scores = {r.content: r.score for r in fused}
|
|
# A: only in list1 at rank 0 → 1/(60+1) ≈ 0.01639
|
|
assert abs(scores["A"] - 1 / 61) < 1e-6
|
|
# B: list1 rank 1 + list2 rank 0 → 1/(60+2) + 1/(60+1)
|
|
expected_b = 1 / 62 + 1 / 61
|
|
assert abs(scores["B"] - expected_b) < 1e-6
|
|
# C: only in list2 at rank 1 → 1/(60+2)
|
|
assert abs(scores["C"] - 1 / 62) < 1e-6
|
|
# B should be ranked first (appears in both lists)
|
|
assert fused[0].content == "B"
|
|
|
|
|
|
def test_rrf_with_disjoint_results():
|
|
"""Two lists with no overlap."""
|
|
list1 = [RetrievalResult(content="X", score=1.0)]
|
|
list2 = [RetrievalResult(content="Y", score=1.0)]
|
|
fused = reciprocal_rank_fusion([list1, list2], k=60)
|
|
contents = {r.content for r in fused}
|
|
assert contents == {"X", "Y"}
|
|
# Equal RRF scores (both at rank 0 in their list)
|
|
assert abs(fused[0].score - fused[1].score) < 1e-6
|
|
|
|
|
|
def test_rrf_with_overlapping_results():
|
|
"""Document appearing in both lists gets higher fused score."""
|
|
shared = RetrievalResult(content="shared", score=5.0)
|
|
unique = RetrievalResult(content="unique", score=10.0)
|
|
list1 = [shared, unique]
|
|
list2 = [RetrievalResult(content="shared", score=3.0)]
|
|
fused = reciprocal_rank_fusion([list1, list2], k=60)
|
|
scores = {r.content: r.score for r in fused}
|
|
assert scores["shared"] > scores["unique"]
|
|
|
|
|
|
def test_rrf_custom_weights():
|
|
"""Weights affect the contribution of each list."""
|
|
list1 = [RetrievalResult(content="A", score=1.0)]
|
|
list2 = [RetrievalResult(content="B", score=1.0)]
|
|
fused = reciprocal_rank_fusion(
|
|
[list1, list2],
|
|
k=60,
|
|
weights=[2.0, 1.0],
|
|
)
|
|
scores = {r.content: r.score for r in fused}
|
|
# A has weight 2, B has weight 1, both at rank 0
|
|
assert scores["A"] > scores["B"]
|
|
|
|
|
|
# -- HybridMemory integration tests ---------------------------------------
|
|
|
|
|
|
def test_registration():
|
|
MemoryRegistry.register_value("hybrid", HybridMemory)
|
|
assert MemoryRegistry.contains("hybrid")
|
|
|
|
|
|
def test_store_delegates_to_both():
|
|
hybrid = _make_hybrid()
|
|
doc_id = hybrid.store("test content", source="test.txt")
|
|
assert isinstance(doc_id, str)
|
|
# Both sub-backends should have the content
|
|
sparse_results = hybrid._sparse.retrieve("test content")
|
|
dense_results = hybrid._dense.retrieve("test content")
|
|
assert len(sparse_results) >= 1
|
|
assert len(dense_results) >= 1
|
|
|
|
|
|
def test_retrieve_fuses_results():
|
|
hybrid = _make_hybrid()
|
|
hybrid.store("machine learning algorithms")
|
|
hybrid.store("cooking recipes for dinner")
|
|
results = hybrid.retrieve("machine learning")
|
|
assert len(results) >= 1
|
|
assert "machine" in results[0].content.lower()
|
|
|
|
|
|
def test_retrieve_top_k():
|
|
hybrid = _make_hybrid()
|
|
for i in range(10):
|
|
hybrid.store(f"document {i} about testing topic")
|
|
results = hybrid.retrieve("testing", top_k=3)
|
|
assert len(results) <= 3
|
|
|
|
|
|
def test_delete_from_both():
|
|
hybrid = _make_hybrid()
|
|
doc_id = hybrid.store("deletable content")
|
|
assert hybrid.delete(doc_id) is True
|
|
# Should be gone from sparse
|
|
assert len(hybrid._sparse.retrieve("deletable")) == 0
|
|
|
|
|
|
def test_clear_both():
|
|
hybrid = _make_hybrid()
|
|
hybrid.store("doc one about topics")
|
|
hybrid.store("doc two about topics")
|
|
hybrid.clear()
|
|
assert len(hybrid._sparse.retrieve("topics")) == 0
|
|
assert len(hybrid._dense.retrieve("topics")) == 0
|
|
|
|
|
|
def test_event_bus_store():
|
|
bus = EventBus(record_history=True)
|
|
hybrid = _make_hybrid()
|
|
import openjarvis.tools.storage.hybrid as mod
|
|
|
|
original = mod.get_event_bus
|
|
mod.get_event_bus = lambda: bus
|
|
try:
|
|
hybrid.store("event test content")
|
|
events = [e for e in bus.history if e.event_type == EventType.MEMORY_STORE]
|
|
assert len(events) >= 1
|
|
assert any(e.data.get("backend") == "hybrid" for e in events)
|
|
finally:
|
|
mod.get_event_bus = original
|
|
|
|
|
|
def test_event_bus_retrieve():
|
|
bus = EventBus(record_history=True)
|
|
hybrid = _make_hybrid()
|
|
hybrid.store("retrievable content here")
|
|
import openjarvis.tools.storage.hybrid as mod
|
|
|
|
original = mod.get_event_bus
|
|
mod.get_event_bus = lambda: bus
|
|
try:
|
|
hybrid.retrieve("retrievable")
|
|
events = [e for e in bus.history if e.event_type == EventType.MEMORY_RETRIEVE]
|
|
assert len(events) >= 1
|
|
assert any(e.data.get("backend") == "hybrid" for e in events)
|
|
finally:
|
|
mod.get_event_bus = original
|