feat: 添加知识库检索功能和改进模型路由错误处理
- 在ChatRequest中添加Citation事件类型，支持引用来源展示 - 实现聊天上下文准备服务，构建带源元数据的受限聊天上下文 - 添加ThreadedProcess类以支持Windows平台的子进程操作 - 改进检索引擎中的错误处理和向量搜索逻辑 - 实现严格的嵌入模型验证和索引重建机制 - 添加前端聊天界面的知识库检索开关 - 实现搜索历史记录功能和错误降级处理 - 更新模型路由设置提示信息以反映索引重建需求
This commit is contained in:
@@ -0,0 +1,56 @@
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from app.contracts import ChatRequest, Message, ModelEvent, ModelEventType, SearchRequest
|
||||
from app.routes import chat, utc_now
|
||||
from app.services import note_service
|
||||
from app.services.chat_context import prepare
|
||||
|
||||
|
||||
@pytest.mark.parametrize('enabled', [True, False])
|
||||
def test_chat_stream_retrieves_real_notes_and_emits_sources(monkeypatch, enabled):
|
||||
received = []
|
||||
|
||||
class Adapter:
|
||||
async def stream(self, request):
|
||||
received.append(request)
|
||||
yield ModelEvent(event=ModelEventType.text_delta, sequence=0, data={'text': 'answer [1]'}, timestamp=utc_now())
|
||||
yield ModelEvent(event=ModelEventType.done, sequence=1, data={}, timestamp=utc_now())
|
||||
|
||||
monkeypatch.setattr('app.routes.provider_or_404', lambda _: SimpleNamespace(adapter=Adapter()))
|
||||
|
||||
async def scenario():
|
||||
note = await note_service.create_note(title='Orchard', markdown='apple orchard knowledge', folder=None, tags=[])
|
||||
request = ChatRequest(provider_id='test', model='test', use_rag=enabled,
|
||||
system='Keep original instructions',
|
||||
messages=[Message(role='user', content='apple')],
|
||||
retrieval=SearchRequest(query='apple', mode='fts'))
|
||||
response = await chat(request)
|
||||
chunks = [chunk async for chunk in response.body_iterator]
|
||||
events = [json.loads(chunk.split('data: ', 1)[1]) for chunk in chunks]
|
||||
assert [e['sequence'] for e in events] == list(range(len(events)))
|
||||
assert events[-1]['event'] == 'Done'
|
||||
assert received[0].messages == request.messages
|
||||
if enabled:
|
||||
assert events[0]['event'] == 'Citation'
|
||||
assert events[0]['data']['note_id'] == note.note_id
|
||||
assert 'apple orchard knowledge' in received[0].system
|
||||
assert 'Keep original instructions' in received[0].system
|
||||
else:
|
||||
assert all(e['event'] != 'Citation' for e in events)
|
||||
assert received[0].system == request.system
|
||||
assert request.system == 'Keep original instructions'
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_empty_knowledge_base_has_no_invented_citations():
|
||||
async def scenario():
|
||||
request = ChatRequest(provider_id='test', model='test', messages=[Message(role='user', content='missing')])
|
||||
grounded, sources = await prepare(request)
|
||||
assert sources == []
|
||||
assert '不要编造' in grounded.system
|
||||
asyncio.run(scenario())
|
||||
@@ -84,3 +84,56 @@ def test_cancel_reaps_active_model_process(monkeypatch):
|
||||
await task
|
||||
assert process.killed and not runtime.active
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cancel", [False, True])
|
||||
def test_subprocess_fallback_runs_and_reaps_real_worker(monkeypatch, tmp_path, cancel):
|
||||
import app.local_models.runtime as module
|
||||
import app.local_models.process as process_module
|
||||
|
||||
monkeypatch.setattr(module, 'read_state', lambda key: {'status': 'installed'})
|
||||
monkeypatch.setattr(module, 'interpreter', lambda: Path(sys.executable))
|
||||
worker = tmp_path / 'worker.py'
|
||||
worker.write_text(
|
||||
'import json,sys,time\n'
|
||||
'request=json.load(sys.stdin)\n'
|
||||
'print(json.dumps({"progress": 1}),flush=True)\n'
|
||||
+ ('time.sleep(60)\n' if cancel else '')
|
||||
+ 'print(json.dumps({"result": [[1.0,0.0]], "usage": {"input_tokens": 2}}),flush=True)\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
processes = []
|
||||
original = process_module.ThreadedProcess
|
||||
|
||||
def spawn(args, **kwargs):
|
||||
process = original((sys.executable, str(worker)), **kwargs)
|
||||
processes.append(process)
|
||||
return process
|
||||
|
||||
async def unsupported(*args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
monkeypatch.setattr(module.asyncio, 'create_subprocess_exec', unsupported)
|
||||
monkeypatch.setattr(process_module, 'ThreadedProcess', spawn)
|
||||
|
||||
async def scenario():
|
||||
runtime = Runtime()
|
||||
started = asyncio.Event()
|
||||
token = module.runtime_progress.set(lambda message: started.set())
|
||||
try:
|
||||
task = asyncio.create_task(runtime.infer('bekko', 'embedding', {'texts': ['test']}))
|
||||
await asyncio.wait_for(started.wait(), 10)
|
||||
if cancel:
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
assert await task == [[1.0, 0.0]]
|
||||
assert not runtime.active and not runtime.active_files and not runtime.waiters
|
||||
assert processes[0].returncode is not None
|
||||
assert processes[0].process.stdin.closed
|
||||
assert processes[0].process.stdout.closed
|
||||
finally:
|
||||
module.runtime_progress.reset(token)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -464,3 +464,76 @@ def test_missing_runtime_uses_unchanged_local_retrieval(runtime, monkeypatch):
|
||||
assert runtime.calls == []
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def production_engine(monkeypatch):
|
||||
from app.local_models.runtime import LocalEmbedding
|
||||
embedding = LocalEmbedding()
|
||||
monkeypatch.setattr(note_service, "embedding", embedding)
|
||||
return RetrievalEngine(embedding, LexicalReranker(), SqliteVecStore(), route_embeddings=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["api", "local"])
|
||||
def test_real_embedding_route_rebuilds_missing_space(runtime, production_engine, source):
|
||||
from app.errors import ApiError
|
||||
runtime.source = source
|
||||
|
||||
async def scenario():
|
||||
await seed()
|
||||
runtime.model_id = "new-configured-space"
|
||||
with pytest.raises(ApiError) as error:
|
||||
await production_engine.search(request())
|
||||
assert error.value.code == "SEMANTIC_INDEX_UNAVAILABLE"
|
||||
assert "Embedding 已可用" in error.value.message
|
||||
assert error.value.details["source"] == source
|
||||
await index_service.rebuild(IndexRebuildRequest())
|
||||
assert (await production_engine.search(request())).items
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_real_embedding_failure_is_not_reported_as_missing_configuration(runtime, production_engine):
|
||||
from app.errors import ApiError
|
||||
|
||||
async def scenario():
|
||||
await seed()
|
||||
runtime.error = ApiError(503, "LOCAL_MODEL_TIMEOUT", "本地模型推理超时。", {"fallback_reason": "PROVIDER_TIMEOUT"})
|
||||
with pytest.raises(ApiError) as error:
|
||||
await production_engine.search(request())
|
||||
assert error.value.code == "LOCAL_MODEL_TIMEOUT"
|
||||
assert error.value.details["fallback_reason"] == "PROVIDER_TIMEOUT"
|
||||
assert (await production_engine.search(SearchRequest(query="apple", mode=SearchMode.hybrid))).items
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["inference", "storage", "space_change"])
|
||||
def test_real_embedding_rebuild_failure_preserves_index(runtime, production_engine, monkeypatch, failure):
|
||||
from app.errors import ApiError
|
||||
|
||||
async def scenario():
|
||||
await seed()
|
||||
tables = ("notes", "blocks", "blocks_fts", "index_meta", "routed_block_vectors")
|
||||
before = {table: [tuple(r) for r in rows(f"SELECT * FROM {table}")] for table in tables}
|
||||
if failure == "inference":
|
||||
runtime.error = ApiError(503, "LOCAL_MODEL_TIMEOUT", "本地模型推理超时。")
|
||||
elif failure == "storage":
|
||||
monkeypatch.setattr(routed_vectors, "store_remote", lambda *args: None)
|
||||
else:
|
||||
original = runtime.embed
|
||||
async def changing(texts):
|
||||
runtime.model_id += "x"
|
||||
return await original(texts)
|
||||
monkeypatch.setattr(runtime, "embed", changing)
|
||||
with pytest.raises(ApiError):
|
||||
await index_service.rebuild(IndexRebuildRequest())
|
||||
assert index_service.get_status().status == "failed"
|
||||
after = {table: [tuple(r) for r in rows(f"SELECT * FROM {table}")] for table in tables}
|
||||
assert before == after
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_empty_vault_vector_search_returns_empty(runtime, production_engine):
|
||||
assert asyncio.run(production_engine.search(request())).items == []
|
||||
|
||||
Reference in New Issue
Block a user