216 lines
13 KiB
Python
216 lines
13 KiB
Python
"""Failure recovery, pinned models, deletion and permission regressions."""
|
|
import copy
|
|
import json
|
|
from types import SimpleNamespace
|
|
from unittest import TestCase, mock
|
|
|
|
from archive_store import ArchiveStore
|
|
from knowledge_models import model_choices, safe_model_error
|
|
from knowledge_store import KnowledgeStore
|
|
from knowledge_worker import KnowledgeWorker, enhance
|
|
import test_knowledge as fixtures
|
|
|
|
|
|
def fake_catalog():
|
|
outlets = {name: SimpleNamespace(id=name, name=name, kind=kind,
|
|
config={'model': name, 'api_key': 'private-value', 'timeout_ms': 30000})
|
|
for name, kind in [('alpha', 'openai'), ('beta', 'dify'), ('image', 'comfyui')]}
|
|
return SimpleNamespace(outlets=outlets, plan=lambda: ([outlets['alpha']], None, '', 1))
|
|
|
|
|
|
class KnowledgeTaskTest(TestCase):
|
|
payload = fixtures.KnowledgeTest.payload
|
|
start_job = fixtures.KnowledgeTest.start_job
|
|
process = fixtures.KnowledgeTest.process
|
|
first = fixtures.KnowledgeTest.first
|
|
publish = fixtures.KnowledgeTest.publish
|
|
|
|
def setUp(self):
|
|
fixtures.KnowledgeTest.setUp(self)
|
|
self.catalog = fake_catalog()
|
|
self.patch = mock.patch('knowledge_models.catalog_for', return_value=self.catalog)
|
|
self.patch.start()
|
|
|
|
def tearDown(self):
|
|
self.patch.stop()
|
|
fixtures.KnowledgeTest.tearDown(self)
|
|
|
|
def test_delete_requires_rejection_revision_and_tenant_and_keeps_archive(self):
|
|
self.process()
|
|
item = self.first()
|
|
args = ('default', item['id'], item['revision'])
|
|
with self.assertRaises(ValueError):
|
|
self.store.delete_rejected(*args, self.actor, '')
|
|
self.store.transition(*args, 'reject', self.actor, '')
|
|
with self.assertRaises(KeyError):
|
|
self.store.delete_rejected('other', item['id'], item['revision'], self.actor, '')
|
|
with self.assertRaises(ValueError):
|
|
self.store.delete_rejected('default', item['id'], 99, self.actor, '')
|
|
self.store.delete_rejected(*args, self.actor, '')
|
|
self.assertEqual(self.store.list_items('default')['total'], 0)
|
|
with self.assertRaises(KeyError):
|
|
self.store.detail('default', item['id'])
|
|
with self.db.connect() as db:
|
|
self.assertEqual(db.execute('SELECT COUNT(*) FROM archive_message').fetchone()[0], 2)
|
|
for table in ('knowledge_source', 'knowledge_revision', 'knowledge_fts'):
|
|
self.assertEqual(db.execute(f'SELECT COUNT(*) FROM {table}').fetchone()[0], 0)
|
|
self.assertIsNotNone(db.execute("SELECT 1 FROM audit_log WHERE action='knowledge.delete'").fetchone())
|
|
self.start_job()
|
|
KnowledgeWorker(self.store).step()
|
|
self.assertEqual(self.store.list_items('default')['total'], 0, 'Identical discarded drafts must stay discarded')
|
|
|
|
def test_model_choice_is_pinned_and_image_disabled_unknown_rejected(self):
|
|
self.archive.import_messages(self.payload(), self.actor, '')
|
|
job = self.start_job(engine='model', model_provider_id='beta', model_timeout_seconds=120)
|
|
self.assertEqual(job['options']['model_provider_id'], 'beta')
|
|
self.assertEqual(job['options']['model_name'], 'beta')
|
|
default = self.start_job(engine='model')
|
|
self.assertEqual(default['options']['model_provider_id'], 'alpha')
|
|
for provider in ('missing', 'image'):
|
|
with self.assertRaises(ValueError):
|
|
self.start_job(engine='model', model_provider_id=provider)
|
|
self.assertNotIn('private-value', json.dumps(model_choices(self.db)))
|
|
self.assertEqual(len(model_choices(self.db)), 2)
|
|
|
|
def test_resume_preserves_successful_candidates_in_failed_batch(self):
|
|
self.archive.import_messages(self.payload(messages=[
|
|
('inbound', '预约挂号需要什么材料？'), ('outbound', '请先提供就诊日期和科室。'),
|
|
('inbound', '测试门店周日营业时间是几点？'), ('outbound', '测试门店周日上午九点营业。')]), self.actor, '')
|
|
job = self.start_job(engine='model', model_provider_id='alpha')
|
|
seen = []
|
|
def succeed_then_fail(candidate, database, options, **usage_context):
|
|
seen.append(candidate['question'])
|
|
if len(seen) == 2:
|
|
raise ValueError('模拟超时')
|
|
return candidate, 50, 80
|
|
with mock.patch('knowledge_worker.enhance', side_effect=succeed_then_fail):
|
|
KnowledgeWorker(self.store).step()
|
|
failed = self.store.job('default', job['id'])
|
|
self.assertEqual((failed['status'], failed['processed'], failed['created_items'], failed['pending_items']), ('failed', 4, 1, 1))
|
|
self.assertEqual(failed['input_chars'], 50)
|
|
self.assertNotIn('state', failed)
|
|
self.store.configure_job('default', job['id'], {'engine': 'model', 'model_provider_id': 'beta',
|
|
'model_timeout_seconds': 120, 'max_model_calls': 200}, self.actor, '')
|
|
self.store.job_action('default', job['id'], 'resume', self.actor, '')
|
|
def remaining(candidate, database, options, **usage_context):
|
|
self.assertEqual(candidate['question'], seen[1])
|
|
self.assertEqual(options['model_provider_id'], 'beta')
|
|
return candidate, 40, 60
|
|
with mock.patch('knowledge_worker.enhance', side_effect=remaining) as call:
|
|
KnowledgeWorker(self.store).step()
|
|
self.assertEqual(call.call_count, 1)
|
|
final = self.store.job('default', job['id'])
|
|
self.assertEqual((final['status'], final['created_items'], final['processed'], final['model_calls']), ('completed', 2, 4, 3))
|
|
|
|
def test_budget_can_be_increased_without_losing_pending_candidate(self):
|
|
self.archive.import_messages(self.payload(), self.actor, '')
|
|
job = self.start_job(engine='model', max_model_calls=1)
|
|
with mock.patch('knowledge_worker.enhance', side_effect=ValueError('失败')):
|
|
KnowledgeWorker(self.store).step()
|
|
with self.assertRaises(ValueError):
|
|
self.store.configure_job('default', job['id'], {'engine': 'model', 'max_model_calls': 1}, self.actor, '')
|
|
self.store.configure_job('default', job['id'], {'engine': 'rules'}, self.actor, '')
|
|
self.store.job_action('default', job['id'], 'resume', self.actor, '')
|
|
with mock.patch('knowledge_worker.enhance') as call:
|
|
KnowledgeWorker(self.store).step()
|
|
call.assert_not_called()
|
|
self.assertEqual(self.store.job('default', job['id'])['created_items'], 1)
|
|
|
|
def test_cancel_during_call_fences_result_and_pause_can_continue(self):
|
|
self.archive.import_messages(self.payload(), self.actor, '')
|
|
job = self.start_job(engine='model')
|
|
def pause(candidate, database, options, **usage_context):
|
|
self.store.job_action('default', job['id'], 'pause', self.actor, '')
|
|
return candidate, 10, 20
|
|
with mock.patch('knowledge_worker.enhance', side_effect=pause):
|
|
KnowledgeWorker(self.store).step()
|
|
self.assertEqual(self.store.list_items('default')['total'], 0)
|
|
self.assertEqual(self.store.job('default', job['id'])['pending_items'], 1)
|
|
self.store.job_action('default', job['id'], 'resume', self.actor, '')
|
|
def cancel(candidate, database, options, **usage_context):
|
|
self.store.job_action('default', job['id'], 'cancel', self.actor, '')
|
|
return candidate, 10, 20
|
|
with mock.patch('knowledge_worker.enhance', side_effect=cancel):
|
|
KnowledgeWorker(self.store).step()
|
|
self.assertEqual(self.store.job('default', job['id'])['status'], 'cancelled')
|
|
self.assertEqual(self.store.job('default', job['id'])['pending_items'], 0)
|
|
self.assertEqual(self.store.list_items('default')['total'], 0)
|
|
|
|
def test_config_rejects_active_foreign_invalid_timeout_and_pins_legacy_job(self):
|
|
self.archive.import_messages(self.payload(), self.actor, '')
|
|
job = self.start_job(engine='model')
|
|
with self.assertRaises(ValueError):
|
|
self.store.configure_job('default', job['id'], {'engine': 'rules'}, self.actor, '')
|
|
with self.assertRaises(KeyError):
|
|
self.store.configure_job('other', job['id'], {'engine': 'rules'}, self.actor, '')
|
|
self.store.job_action('default', job['id'], 'pause', self.actor, '')
|
|
for timeout in (12, 181):
|
|
with self.assertRaises(ValueError):
|
|
self.store.configure_job('default', job['id'], {'model_timeout_seconds': timeout}, self.actor, '')
|
|
with self.db.connect() as db:
|
|
options = json.loads(db.execute('SELECT options_json FROM knowledge_job WHERE id=?', (job['id'],)).fetchone()[0])
|
|
options.pop('model_provider_id')
|
|
db.execute('UPDATE knowledge_job SET options_json=? WHERE id=?', (json.dumps(options), job['id']))
|
|
self.store.job_action('default', job['id'], 'resume', self.actor, '')
|
|
with mock.patch('knowledge_worker.enhance', side_effect=lambda c, d, o: (c, 10, 20)):
|
|
KnowledgeWorker(self.store).step()
|
|
self.assertEqual(self.store.job('default', job['id'])['options']['model_provider_id'], 'alpha')
|
|
|
|
def test_selected_model_timeout_output_budget_and_safe_errors(self):
|
|
self.process()
|
|
candidate = self.first()
|
|
fields = {k: candidate[k] for k in ('title', 'question', 'answer', 'category', 'kind')}
|
|
fields['conditions'] = '仅限虚构测试门店'
|
|
options = {'model_provider_id': 'beta', 'model_timeout_seconds': 120}
|
|
with mock.patch('model_gateway.call_outlet', new_callable=mock.AsyncMock) as call:
|
|
call.return_value = {'text': '```json\n' + json.dumps(fields) + '\n```', 'error': ''}
|
|
result, input_chars, output_chars = enhance(copy.deepcopy(candidate), self.db, options)
|
|
self.assertEqual(result['conditions'], fields['conditions'])
|
|
self.assertGreater(input_chars, 0)
|
|
self.assertGreater(output_chars, 0)
|
|
self.assertEqual(call.call_args.args[1].id, 'beta')
|
|
self.assertEqual(call.call_args.args[1].config['timeout_ms'], 120000)
|
|
self.assertEqual(call.call_args.kwargs['deadline'], 120)
|
|
self.assertEqual(call.call_args.kwargs['max_tokens'], 4096)
|
|
call.return_value = {'text': '', 'error': '上游返回 401: private-value https://secret/path'}
|
|
with self.assertRaisesRegex(ValueError, 'HTTP 401') as error:
|
|
enhance(candidate, self.db, options)
|
|
self.assertNotIn('private-value', str(error.exception))
|
|
for text in ('{"title":', '[]', '{"kind":"qa"}'):
|
|
call.return_value = {'text': text, 'error': ''}
|
|
with self.assertRaises(ValueError):
|
|
enhance(copy.deepcopy(candidate), self.db, options)
|
|
self.assertIn('超时', safe_model_error('ReadTimeout private-value'))
|
|
self.assertIn('限流', safe_model_error('上游返回 429 secret'))
|
|
|
|
|
|
class KnowledgeTaskApiTest(TestCase):
|
|
setUp = fixtures.KnowledgeApiTest.setUp
|
|
tearDown = fixtures.KnowledgeApiTest.tearDown
|
|
|
|
def test_delete_confirmation_and_model_metadata_auth(self):
|
|
store = self.client.app.state.knowledge_store
|
|
actor = self.db.authenticate('admin', 'ChangedAdmin@2026')['id']
|
|
ArchiveStore(self.db).import_messages(fixtures.KnowledgeTest.payload(self), actor, '')
|
|
with self.db.connect() as db:
|
|
source = db.execute('SELECT id FROM archive_source_account').fetchone()[0]
|
|
job = store.create_job('default', {'source_account_id': source, 'staff_confirmed': True}, actor, '')
|
|
KnowledgeWorker(store).step()
|
|
item = store.list_items('default')['items'][0]
|
|
store.transition('default', item['id'], 1, 'reject', actor, '')
|
|
url = '/api/v2/knowledge/items/' + item['id'] + '/delete'
|
|
self.assertEqual(self.client.post(url, json={'revision': 1, 'confirmed': True}).status_code, 401)
|
|
self.assertEqual(self.client.post(url, headers=self.headers, json={'revision': 1}).status_code, 400)
|
|
self.assertEqual(self.client.post(url + '?account_id=-1', headers=self.headers, json={'revision': 1, 'confirmed': True}).status_code, 400)
|
|
self.assertEqual(self.client.post(url, headers=self.headers, json={'revision': 1, 'confirmed': True}).status_code, 200)
|
|
self.assertEqual(self.client.get(url.rsplit('/', 1)[0], headers=self.headers).status_code, 404)
|
|
self.assertEqual(self.client.get('/api/v2/knowledge/models').status_code, 401)
|
|
with mock.patch('knowledge_models.catalog_for', return_value=fake_catalog()):
|
|
response = self.client.get('/api/v2/knowledge/models', headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertNotIn('private-value', response.text)
|
|
self.assertEqual(len(response.json()['items']), 2)
|
|
response = self.client.put('/api/v2/knowledge/jobs/' + job['id'] + '/config', headers=self.headers,
|
|
json={'engine': 'model', 'model_timeout_seconds': 12})
|
|
self.assertEqual(response.status_code, 422)
|