mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-31 04:45:23 +08:00
368 lines
13 KiB
Python
368 lines
13 KiB
Python
import asyncio
|
|
import unittest
|
|
|
|
from benchmarks.sdk_stubs import install_stubs
|
|
|
|
install_stubs()
|
|
|
|
from components.knowledge_engine.langrag import LangRAG, PARSED_TEXT_LIMIT_BYTES
|
|
from components.observability.installation import InstallationTelemetry
|
|
from components.offload import BoundedOffload
|
|
from components.shared_state import FENCE_KEY_PREFIX, FenceStateUnavailableError
|
|
from benchmarks.state_fixture import attach_installation_state
|
|
from langbot_plugin.api.entities.builtin.rag import (
|
|
DocumentStatus,
|
|
FileMetadata,
|
|
FileObject,
|
|
IngestionContext,
|
|
ParseResult,
|
|
RetrievalContext,
|
|
)
|
|
|
|
|
|
class ControllableStoragePlugin:
|
|
"""Host storage that can be taken offline, with observable remote calls."""
|
|
|
|
def __init__(self):
|
|
self.data = {}
|
|
self.reads_offline = False
|
|
self.embeddings = 0
|
|
self.upserts = []
|
|
self.deletes = []
|
|
self.delete_entered = asyncio.Event()
|
|
self.delete_release = asyncio.Event()
|
|
self.delete_release.set()
|
|
self.offload = BoundedOffload()
|
|
self.telemetry = InstallationTelemetry(self)
|
|
|
|
async def get_plugin_storage_keys(self):
|
|
if self.reads_offline:
|
|
raise RuntimeError("fixture fence store offline")
|
|
return list(self.data)
|
|
|
|
async def get_plugin_storage(self, key):
|
|
if self.reads_offline:
|
|
raise RuntimeError("fixture fence store offline")
|
|
return self.data[key]
|
|
|
|
async def set_plugin_storage(self, key, value):
|
|
self.data[key] = value
|
|
|
|
async def invoke_embedding(self, embedding_model_uuid, texts):
|
|
self.embeddings += 1
|
|
return [[float(i)] for i in range(len(texts))]
|
|
|
|
async def vector_upsert(self, **kwargs):
|
|
self.upserts.append(kwargs)
|
|
|
|
async def vector_delete(self, **kwargs):
|
|
self.deletes.append(kwargs)
|
|
|
|
async def commit():
|
|
# A detached Host commit outlives cancellation of the caller.
|
|
self.delete_entered.set()
|
|
await self.delete_release.wait()
|
|
return 0
|
|
|
|
return await asyncio.shield(asyncio.create_task(commit()))
|
|
|
|
|
|
def ingest_context(text="external parser text"):
|
|
"""One external-parser ingestion for knowledge base ``kb1``."""
|
|
return IngestionContext(
|
|
file_object=FileObject(
|
|
metadata=FileMetadata(
|
|
filename="sample.txt",
|
|
file_size=len(text),
|
|
mime_type="text/plain",
|
|
document_id="doc1",
|
|
knowledge_base_id="kb1",
|
|
),
|
|
storage_path="/missing/sample.txt",
|
|
),
|
|
knowledge_base_id="kb1",
|
|
creation_settings={
|
|
"embedding_model_uuid": "emb1",
|
|
"chunk_size": 100,
|
|
"overlap": 0,
|
|
},
|
|
parsed_content=ParseResult(text=text),
|
|
)
|
|
|
|
|
|
async def bound_engine(plugin):
|
|
engine = LangRAG()
|
|
engine.plugin = plugin
|
|
await plugin.telemetry.initialize()
|
|
return engine
|
|
|
|
|
|
class RecordingIngestPlugin:
|
|
def __init__(self):
|
|
self.read_called = False
|
|
self.upserts = []
|
|
|
|
async def get_knowledge_file_stream(self, storage_path):
|
|
self.read_called = True
|
|
raise AssertionError("external parser content should skip file reads")
|
|
|
|
async def invoke_embedding(self, embedding_model_uuid, texts):
|
|
return [[float(i)] for i, _ in enumerate(texts)]
|
|
|
|
async def vector_upsert(self, **kwargs):
|
|
self.upserts.append(kwargs)
|
|
|
|
|
|
class RecordingVectorPlugin:
|
|
def __init__(self, items):
|
|
self.items = items
|
|
self.requested_ids = []
|
|
|
|
async def vector_get_by_ids(self, collection_id, ids):
|
|
self.requested_ids = ids
|
|
return [self.items[item_id] for item_id in ids if item_id in self.items]
|
|
|
|
|
|
class RecordingRetrievePlugin:
|
|
def __init__(self):
|
|
self.search_top_k = None
|
|
self.rerank_documents = []
|
|
self.rerank_top_k = None
|
|
|
|
async def invoke_embedding(self, embedding_model_uuid, texts):
|
|
return [[0.0] for _ in texts]
|
|
|
|
async def vector_search(self, **kwargs):
|
|
self.search_top_k = kwargs["top_k"]
|
|
return [
|
|
{
|
|
"id": f"chunk-{i}",
|
|
"distance": float(i),
|
|
"metadata": {
|
|
"text": f"candidate {i}",
|
|
"document_name": "doc.txt",
|
|
"chunk_index": i,
|
|
"index_type": "chunk",
|
|
},
|
|
}
|
|
for i in range(kwargs["top_k"])
|
|
]
|
|
|
|
async def invoke_rerank(self, rerank_model_uuid, query, documents, top_k=None):
|
|
self.rerank_documents = list(documents)
|
|
self.rerank_top_k = top_k
|
|
return [
|
|
{"index": 2, "relevance_score": 0.95},
|
|
{"index": 0, "relevance_score": 0.5},
|
|
]
|
|
|
|
|
|
class LangRAGTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_ingest_uses_external_parse_result_without_reading_file(self):
|
|
engine = LangRAG()
|
|
plugin = RecordingIngestPlugin()
|
|
await attach_installation_state(plugin)
|
|
engine.plugin = plugin
|
|
|
|
context = IngestionContext(
|
|
file_object=FileObject(
|
|
metadata=FileMetadata(
|
|
filename="sample.txt",
|
|
file_size=12,
|
|
mime_type="text/plain",
|
|
document_id="doc1",
|
|
knowledge_base_id="kb1",
|
|
),
|
|
storage_path="/missing/sample.txt",
|
|
),
|
|
knowledge_base_id="kb1",
|
|
creation_settings={
|
|
"embedding_model_uuid": "emb1",
|
|
"chunk_size": 100,
|
|
"overlap": 0,
|
|
},
|
|
parsed_content=ParseResult(text="external parser text"),
|
|
)
|
|
|
|
result = await engine.ingest(context)
|
|
|
|
self.assertEqual(result.status, DocumentStatus.COMPLETED)
|
|
self.assertFalse(plugin.read_called)
|
|
self.assertEqual(result.chunks_created, 1)
|
|
self.assertEqual(plugin.upserts[0]["documents"], ["external parser text"])
|
|
|
|
async def test_context_window_uses_parent_child_id_scheme(self):
|
|
engine = LangRAG()
|
|
engine.plugin = RecordingVectorPlugin(
|
|
{
|
|
"doc1_p1_c0": {
|
|
"id": "doc1_p1_c0",
|
|
"metadata": {"text": "previous parent"},
|
|
},
|
|
"doc1_p3_c0": {
|
|
"id": "doc1_p3_c0",
|
|
"metadata": {"text": "next parent"},
|
|
},
|
|
}
|
|
)
|
|
results = [
|
|
{
|
|
"metadata": {
|
|
"document_id": "doc1",
|
|
"index_type": "parent_child",
|
|
"parent_index": 2,
|
|
}
|
|
}
|
|
]
|
|
|
|
await engine._expand_context(results, "kb1", 1)
|
|
|
|
self.assertEqual(set(engine.plugin.requested_ids), {"doc1_p1_c0", "doc1_p3_c0"})
|
|
self.assertEqual(results[0]["metadata"]["context_before"], "previous parent")
|
|
self.assertEqual(results[0]["metadata"]["context_after"], "next parent")
|
|
|
|
def test_neighbor_id_supports_all_index_types(self):
|
|
self.assertEqual(
|
|
LangRAG._neighbor_id(
|
|
{"document_id": "doc1", "index_type": "chunk", "chunk_index": 2},
|
|
-1,
|
|
),
|
|
"doc1_1",
|
|
)
|
|
self.assertEqual(
|
|
LangRAG._neighbor_id(
|
|
{"document_id": "doc1", "index_type": "qa", "chunk_index": "2"},
|
|
1,
|
|
),
|
|
"doc1_3_qa0",
|
|
)
|
|
self.assertEqual(
|
|
LangRAG._neighbor_id(
|
|
{
|
|
"document_id": "doc1",
|
|
"index_type": "parent_child",
|
|
"parent_index": 2,
|
|
},
|
|
1,
|
|
),
|
|
"doc1_p3_c0",
|
|
)
|
|
|
|
async def test_retrieve_uses_host_rerank_model_with_overfetched_candidates(self):
|
|
engine = LangRAG()
|
|
plugin = RecordingRetrievePlugin()
|
|
await attach_installation_state(plugin)
|
|
engine.plugin = plugin
|
|
|
|
response = await engine.retrieve(
|
|
RetrievalContext(
|
|
query="candidate",
|
|
knowledge_base_id="kb1",
|
|
creation_settings={
|
|
"embedding_model_uuid": "emb1",
|
|
"index_type": "chunk",
|
|
},
|
|
retrieval_settings={
|
|
"top_k": 2,
|
|
"rerank": "rerank_model",
|
|
"rerank_model_uuid": "rerank1",
|
|
},
|
|
)
|
|
)
|
|
|
|
self.assertEqual(plugin.search_top_k, 6)
|
|
self.assertEqual(len(plugin.rerank_documents), 6)
|
|
self.assertEqual(plugin.rerank_top_k, 2)
|
|
self.assertEqual([entry.id for entry in response.results], ["chunk-2", "chunk-0"])
|
|
self.assertEqual(response.results[0].score, 0.95)
|
|
self.assertTrue(response.metadata["trace_id"].startswith("retrieval-"))
|
|
self.assertTrue(response.metadata["trace_spans"])
|
|
self.assertTrue(
|
|
all(span["trace_id"] == response.metadata["trace_id"] for span in response.metadata["trace_spans"])
|
|
)
|
|
self.assertIn("vector_search", {span["name"] for span in response.metadata["trace_spans"]})
|
|
|
|
recent = (await plugin.telemetry.snapshot())["recent"]["retrieval"][0]
|
|
self.assertEqual(recent["trace_id"], response.metadata["trace_id"])
|
|
self.assertEqual(recent["trace_spans"], response.metadata["trace_spans"])
|
|
|
|
|
|
class LangRAGFenceTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_unreadable_fence_store_refuses_mutations_and_recovers(self):
|
|
plugin = ControllableStoragePlugin()
|
|
engine = await bound_engine(plugin)
|
|
|
|
plugin.reads_offline = True
|
|
# An unreadable fence store is not "no fence": nothing may be dispatched,
|
|
# and the failure reports the underlying storage error for diagnosis.
|
|
with self.assertRaisesRegex(FenceStateUnavailableError, "fence store offline"):
|
|
await engine.ingest(ingest_context())
|
|
with self.assertRaisesRegex(FenceStateUnavailableError, "fence store offline"):
|
|
await engine.delete_document("kb1", "doc1")
|
|
self.assertEqual(plugin.embeddings, 0)
|
|
self.assertEqual(plugin.upserts, [])
|
|
self.assertEqual(plugin.deletes, [])
|
|
|
|
# The failed read was not cached, so the retry proceeds normally.
|
|
plugin.reads_offline = False
|
|
result = await engine.ingest(ingest_context())
|
|
self.assertEqual(result.status, DocumentStatus.COMPLETED)
|
|
self.assertEqual(len(plugin.upserts), 1)
|
|
self.assertTrue(await engine.delete_document("kb1", "doc1"))
|
|
self.assertEqual(len(plugin.deletes), 1)
|
|
|
|
async def test_cancelled_delete_fences_the_knowledge_base(self):
|
|
plugin = ControllableStoragePlugin()
|
|
engine = await bound_engine(plugin)
|
|
|
|
task = asyncio.create_task(engine.delete_document("kb1", "doc1"))
|
|
await plugin.delete_entered.wait()
|
|
task.cancel()
|
|
await asyncio.sleep(0.01)
|
|
task.cancel()
|
|
try:
|
|
# The dispatched delete's outcome is unknown, so no retry may claim
|
|
# deletion while it is still in flight.
|
|
with self.assertRaisesRegex(RuntimeError, "fenced"):
|
|
await engine.delete_document("kb1", "doc1")
|
|
finally:
|
|
plugin.delete_release.set()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
|
|
# The dispatched delete ran exactly once, and the fence is durable: a
|
|
# restarted worker refuses the same knowledge base.
|
|
self.assertEqual(len(plugin.deletes), 1)
|
|
self.assertTrue(any(key.startswith(FENCE_KEY_PREFIX) for key in plugin.data))
|
|
restarted = LangRAG()
|
|
restarted.plugin = plugin
|
|
with self.assertRaisesRegex(RuntimeError, "fenced"):
|
|
await restarted.delete_document("kb1", "doc1")
|
|
self.assertEqual(len(plugin.deletes), 1)
|
|
|
|
# The caller never observed a confirmed removal, so no delete record
|
|
# claims one; LangRAG holds no document mapping to mark deleted because
|
|
# the Host vector store is the record of what exists.
|
|
recent = (await plugin.telemetry.snapshot())["recent"]["delete"]
|
|
self.assertNotIn(True, [event.get("deleted") for event in recent])
|
|
|
|
async def test_deterministic_validation_error_does_not_fence(self):
|
|
plugin = ControllableStoragePlugin()
|
|
engine = await bound_engine(plugin)
|
|
|
|
oversized = await engine.ingest(
|
|
ingest_context(text="x" * (PARSED_TEXT_LIMIT_BYTES + 1))
|
|
)
|
|
self.assertEqual(oversized.status, DocumentStatus.FAILED)
|
|
self.assertEqual(plugin.upserts, [])
|
|
|
|
# Rejected before any dispatch: the knowledge base stays usable.
|
|
result = await engine.ingest(ingest_context())
|
|
self.assertEqual(result.status, DocumentStatus.COMPLETED)
|
|
self.assertEqual(len(plugin.upserts), 1)
|
|
self.assertFalse(engine._state.is_fenced(engine._state.binding(plugin), "kb1"))
|
|
self.assertFalse(any(key.startswith(FENCE_KEY_PREFIX) for key in plugin.data))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|