Files
NotesAgentic/backend/tests/test_chat_context.py
admin 6eb97bf9ab feat: 添加知识库检索功能和改进模型路由错误处理
- 在ChatRequest中添加Citation事件类型,支持引用来源展示
- 实现聊天上下文准备服务,构建带源元数据的受限聊天上下文
- 添加ThreadedProcess类以支持Windows平台的子进程操作
- 改进检索引擎中的错误处理和向量搜索逻辑
- 实现严格的嵌入模型验证和索引重建机制
- 添加前端聊天界面的知识库检索开关
- 实现搜索历史记录功能和错误降级处理
- 更新模型路由设置提示信息以反映索引重建需求
2026-09-04 13:02:08 +08:00

57 lines
2.5 KiB
Python

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())