chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,267 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from opensquilla.memory.retrieval import MemoryRetriever
|
||||
from opensquilla.memory.session_source import (
|
||||
SessionSourceIndexer,
|
||||
build_session_source_document,
|
||||
)
|
||||
from opensquilla.memory.store import LongTermMemoryStore
|
||||
from opensquilla.memory.sync_manager import MemorySyncManager
|
||||
from opensquilla.memory.types import MemorySearchOpts, MemorySearchResult, MemorySource
|
||||
from opensquilla.session.models import SessionNode, TranscriptEntry
|
||||
from opensquilla.session.storage import SessionStorage
|
||||
from opensquilla.tools.builtin.memory_tools import create_memory_tools
|
||||
from opensquilla.tools.registry import ToolRegistry
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_indexes_sessions_source_and_reports_counts(tmp_path):
|
||||
store = LongTermMemoryStore(tmp_path / "memory.db")
|
||||
await store.initialize()
|
||||
try:
|
||||
chunks = await store.index_file(
|
||||
path="sessions/main/session-1.md",
|
||||
content="User: remember the launch flag\nAssistant: noted",
|
||||
source=MemorySource.sessions,
|
||||
)
|
||||
|
||||
results, mode = await store.search("launch flag", min_score=0.0)
|
||||
counts = await store.source_counts()
|
||||
|
||||
assert chunks > 0
|
||||
assert mode.value == "fts-only"
|
||||
assert results[0].source is MemorySource.sessions
|
||||
assert counts["sessions"]["files"] == 1
|
||||
assert counts["sessions"]["chunks"] >= 1
|
||||
finally:
|
||||
await store.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retriever_keeps_sessions_source_and_filters_hidden_memory_paths():
|
||||
class FakeStore:
|
||||
async def search(self, **_kwargs):
|
||||
return (
|
||||
[
|
||||
MemorySearchResult(
|
||||
chunk_id="hidden",
|
||||
path="memory/.hidden.md",
|
||||
source=MemorySource.memory,
|
||||
start_line=1,
|
||||
end_line=1,
|
||||
snippet="hidden",
|
||||
score=0.99,
|
||||
),
|
||||
MemorySearchResult(
|
||||
chunk_id="session",
|
||||
path="sessions/main/session-1.md",
|
||||
source=MemorySource.sessions,
|
||||
start_line=1,
|
||||
end_line=2,
|
||||
snippet="session recall",
|
||||
score=0.90,
|
||||
),
|
||||
MemorySearchResult(
|
||||
chunk_id="curated",
|
||||
path="memory/a.md",
|
||||
source=MemorySource.memory,
|
||||
start_line=1,
|
||||
end_line=1,
|
||||
snippet="curated recall",
|
||||
score=0.86,
|
||||
),
|
||||
],
|
||||
"fts_only",
|
||||
)
|
||||
|
||||
retriever = MemoryRetriever(FakeStore()) # type: ignore[arg-type]
|
||||
|
||||
results = await retriever.search("recall", MemorySearchOpts(max_results=5, min_score=0.0))
|
||||
|
||||
assert [r.chunk_id for r in results] == ["curated", "session"]
|
||||
assert results[1].metadata["search_intent"] == "tool"
|
||||
|
||||
|
||||
def test_session_source_document_renders_transcript_with_citable_path():
|
||||
session = SessionNode(
|
||||
session_key="direct:user:thread",
|
||||
session_id="session-1",
|
||||
agent_id="main",
|
||||
updated_at=1_700_000_000_000,
|
||||
)
|
||||
entries = [
|
||||
TranscriptEntry(
|
||||
id=10,
|
||||
session_id="session-1",
|
||||
session_key=session.session_key,
|
||||
role="user",
|
||||
content="Please remember the launch flag.",
|
||||
created_at=1_700_000_000_000,
|
||||
),
|
||||
TranscriptEntry(
|
||||
id=11,
|
||||
session_id="session-1",
|
||||
session_key=session.session_key,
|
||||
role="assistant",
|
||||
content="Launch flag recorded.",
|
||||
created_at=1_700_000_001_000,
|
||||
),
|
||||
]
|
||||
|
||||
document = build_session_source_document(session, entries)
|
||||
|
||||
assert document.path == "sessions/main/session-1.md"
|
||||
assert document.mtime == 1_700_000_000.0
|
||||
assert "session_id: session-1" in document.content
|
||||
assert "[entry 10] User: Please remember the launch flag." in document.content
|
||||
assert "[entry 11] Assistant: Launch flag recorded." in document.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_source_indexer_syncs_transcript_documents(tmp_path):
|
||||
storage = SessionStorage(":memory:")
|
||||
await storage.connect()
|
||||
store = LongTermMemoryStore(tmp_path / "memory.db")
|
||||
await store.initialize()
|
||||
try:
|
||||
session = SessionNode(
|
||||
session_key="direct:user:thread",
|
||||
session_id="session-1",
|
||||
agent_id="main",
|
||||
updated_at=1_700_000_000_000,
|
||||
)
|
||||
await storage.upsert_session(session)
|
||||
await storage.append_transcript_entry(
|
||||
TranscriptEntry(
|
||||
session_id=session.session_id,
|
||||
session_key=session.session_key,
|
||||
role="user",
|
||||
content="The migration keyword is cerulean.",
|
||||
created_at=1_700_000_000_000,
|
||||
)
|
||||
)
|
||||
|
||||
indexer = SessionSourceIndexer(storage=storage, store=store, agent_id="main")
|
||||
result = await indexer.sync(force=True)
|
||||
results, _mode = await store.search("cerulean", min_score=0.0)
|
||||
|
||||
assert result.indexed == 1
|
||||
assert results[0].path == "sessions/main/session-1.md"
|
||||
assert results[0].source is MemorySource.sessions
|
||||
|
||||
await storage.delete_session(session.session_key)
|
||||
cleanup = await indexer.sync(force=True)
|
||||
stale_results, _mode = await store.search("cerulean", min_score=0.0)
|
||||
|
||||
assert cleanup.removed == 1
|
||||
assert stale_results == []
|
||||
finally:
|
||||
await store.close()
|
||||
await storage.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_manager_indexes_session_source_on_session_delta(tmp_path):
|
||||
class NoopStore:
|
||||
async def index_file(self, **_kwargs):
|
||||
return 0
|
||||
|
||||
async def remove_file(self, _path):
|
||||
return None
|
||||
|
||||
class FakeSessionIndexer:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[bool] = []
|
||||
|
||||
async def sync(self, *, force: bool = False):
|
||||
self.calls.append(force)
|
||||
return SimpleNamespace(indexed=1, removed=0)
|
||||
|
||||
indexer = FakeSessionIndexer()
|
||||
manager = MemorySyncManager(
|
||||
store=NoopStore(), # type: ignore[arg-type]
|
||||
workspace_dir=tmp_path,
|
||||
memory_dir=tmp_path / "memory",
|
||||
session_indexer=indexer,
|
||||
)
|
||||
manager.notify_message(150_000)
|
||||
|
||||
await manager.sync(reason="session-delta")
|
||||
|
||||
assert indexer.calls == [False]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_manager_search_indexes_pending_session_delta(tmp_path):
|
||||
class NoopStore:
|
||||
async def index_file(self, **_kwargs):
|
||||
return 0
|
||||
|
||||
async def remove_file(self, _path):
|
||||
return None
|
||||
|
||||
class FakeSessionIndexer:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[bool] = []
|
||||
|
||||
async def sync(self, *, force: bool = False):
|
||||
self.calls.append(force)
|
||||
return SimpleNamespace(indexed=1, removed=0)
|
||||
|
||||
indexer = FakeSessionIndexer()
|
||||
manager = MemorySyncManager(
|
||||
store=NoopStore(), # type: ignore[arg-type]
|
||||
workspace_dir=tmp_path,
|
||||
memory_dir=tmp_path / "memory",
|
||||
session_indexer=indexer,
|
||||
)
|
||||
manager.notify_message(20)
|
||||
|
||||
await manager.sync(reason="search:tool")
|
||||
await manager.sync(reason="search:tool")
|
||||
|
||||
assert indexer.calls == [False]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_search_tool_outputs_sessions_source(tmp_path):
|
||||
registry = ToolRegistry()
|
||||
|
||||
class FakeRetriever:
|
||||
def __init__(self) -> None:
|
||||
self.opts = None
|
||||
|
||||
async def search(self, _query, opts, **_kwargs):
|
||||
self.opts = opts
|
||||
return [
|
||||
MemorySearchResult(
|
||||
chunk_id="session",
|
||||
path="sessions/main/session-1.md",
|
||||
source=MemorySource.sessions,
|
||||
start_line=1,
|
||||
end_line=2,
|
||||
snippet="session recall",
|
||||
score=0.8,
|
||||
text="session recall",
|
||||
)
|
||||
]
|
||||
|
||||
retriever = FakeRetriever()
|
||||
create_memory_tools(
|
||||
stores=SimpleNamespace(),
|
||||
retrievers=retriever, # type: ignore[arg-type]
|
||||
memory_dir=str(tmp_path),
|
||||
registry=registry,
|
||||
)
|
||||
|
||||
registered = registry.get("memory_search")
|
||||
assert registered is not None
|
||||
output = await registered.handler(query="recall", source="sessions")
|
||||
|
||||
assert retriever.opts.source is MemorySource.sessions
|
||||
assert "source: sessions" in output
|
||||
assert "sessions/main/session-1.md" in output
|
||||
Reference in New Issue
Block a user