feat(server): add workbench API, DSH planner worker and RAG pipeline
This commit is contained in:
yi
@@ -0,0 +1,233 @@
|
||||
"""Cheap, advisory estimates and progress tied to completed corpus/model work."""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import copy
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from sqlmodel import select
|
||||
|
||||
from .models import KnowledgeBase, LLMCall, ModelProvider
|
||||
|
||||
|
||||
# Estimation is advisory and the same selection is commonly requested several
|
||||
# times while the user changes a form or keeps the page open. Keep a very
|
||||
# short in-process cache so it never waits behind the long-running automatic
|
||||
# reading writer. The cache is deliberately small and expires quickly so
|
||||
# newly completed reading batches are reflected without a restart.
|
||||
_ESTIMATE_CACHE: dict[str, tuple[float, dict]] = {}
|
||||
_ESTIMATE_CACHE_LOCK = threading.RLock()
|
||||
_ESTIMATE_CACHE_TTL = 12.0
|
||||
|
||||
|
||||
def _cache_key(kb_ids, mode, organization_id, selected_model, requested_days, kb_category_ids, user_id):
|
||||
return json.dumps({
|
||||
'kb_ids': sorted({int(value) for value in kb_ids}),
|
||||
'mode': mode,
|
||||
'organization_id': int(organization_id or 0),
|
||||
'selected_model': list(selected_model) if selected_model else None,
|
||||
'requested_days': requested_days,
|
||||
'kb_category_ids': kb_category_ids or {},
|
||||
'user_id': int(user_id or 0),
|
||||
}, ensure_ascii=False, sort_keys=True, separators=(',', ':'))
|
||||
|
||||
|
||||
def assessment_estimate(session, kb_ids, mode, organization_id, selected_model=None, requested_days=None, kb_category_ids=None, user_id=0):
|
||||
"""Advisory metadata-only estimate; current exact cache is checked during scanning."""
|
||||
cache_key = _cache_key(kb_ids, mode, organization_id, selected_model, requested_days, kb_category_ids, user_id)
|
||||
now = time.monotonic()
|
||||
with _ESTIMATE_CACHE_LOCK:
|
||||
cached = _ESTIMATE_CACHE.get(cache_key)
|
||||
if cached and now - cached[0] < _ESTIMATE_CACHE_TTL:
|
||||
return copy.deepcopy(cached[1])
|
||||
from .training_scope import scope_counts
|
||||
from .training_reading import reading_parallelism, reading_cache_signature
|
||||
from .training_knowledge import reuse_inventory
|
||||
scoped = scope_counts(session, kb_ids, kb_category_ids or {})
|
||||
libraries = len(scoped)
|
||||
total_chunks = sum(item['chunk_count'] for item in scoped)
|
||||
total_files = sum(item['file_count'] for item in scoped)
|
||||
parallelism = reading_parallelism(session, user_id, organization_id)
|
||||
reuse = {'reusable_chunks': 0, 'reusable_files': 0, 'reusable_batches': 0}
|
||||
if mode == 'reuse':
|
||||
reuse = reuse_inventory(session, kb_ids, kb_category_ids or {}, reading_cache_signature(session, selected_model))
|
||||
reuse['reusable_chunks'] = min(total_chunks, reuse['reusable_chunks'])
|
||||
reuse['reusable_files'] = min(total_files, reuse['reusable_files'])
|
||||
sizes = [item['chunk_count'] * (max(120, item['chunk_target'] or 500) + 250) for item in scoped]
|
||||
try:
|
||||
batch_chars = max(16000, min(32000, int(os.environ.get('YAWAN_AUTO_READ_BATCH_CHARS', '24000'))))
|
||||
except (TypeError, ValueError):
|
||||
batch_chars = 24000
|
||||
reading_batches = reducing_batches = reading_waves = reducing_waves = 0
|
||||
if mode != 'quick' and (mode == 'reuse' or sum(sizes) > 16000):
|
||||
quota = max(1, ((32000 if mode == 'reuse' else 40000) - 2) // max(1, libraries) - 2)
|
||||
warm_fraction = reuse['reusable_chunks'] / max(1, total_chunks)
|
||||
for item, size in zip(scoped, sizes):
|
||||
batches = max(item['file_count'] if mode == 'reuse' else 0, math.ceil(size / batch_chars)) if item['chunk_count'] else 0
|
||||
reading_batches += batches
|
||||
reading_waves += math.ceil(batches * (1 - warm_fraction) / parallelism)
|
||||
summary_size = batches * min(4000 if mode == 'reuse' else 2800, quota)
|
||||
while summary_size > quota:
|
||||
batches = math.ceil(summary_size / 14000)
|
||||
reducing_batches += batches
|
||||
reducing_waves += math.ceil(batches / parallelism)
|
||||
next_size = batches * min(2800, quota)
|
||||
if next_size >= summary_size:
|
||||
break
|
||||
summary_size = next_size
|
||||
provider = session.get(ModelProvider, selected_model[0]) if selected_model else None
|
||||
def call_history(agent_name):
|
||||
if not provider:
|
||||
return []
|
||||
cutoff = (datetime.now() - timedelta(days=30)).strftime('%Y-%m-%d %H:%M:%S')
|
||||
calls = session.exec(select(LLMCall.duration_ms).where(
|
||||
LLMCall.organization_id == organization_id, LLMCall.provider == provider.name,
|
||||
LLMCall.model == selected_model[1], LLMCall.agent_name == agent_name,
|
||||
LLMCall.status == 'ok', LLMCall.duration_ms > 0, LLMCall.created_at >= cutoff,
|
||||
).order_by(LLMCall.id.desc()).limit(30)).all()
|
||||
return sorted(value / 1000 for value in calls)
|
||||
def bounds(values):
|
||||
if len(values) < 3:
|
||||
return 25, 120
|
||||
return max(5, values[int((len(values) - 1) * .2)] * .65), max(30, values[math.ceil((len(values) - 1) * .8)] * 1.7)
|
||||
durations, read_durations = call_history('训练营导图'), call_history('训练营资料阅读')
|
||||
call_low, call_high = bounds(durations)
|
||||
read_low, read_high = bounds(read_durations)
|
||||
basis = (f'按所选范围的索引计数和当前模型最近 {len(durations)} 次成功导图调用估算'
|
||||
if len(durations) >= 3 else '按所选范围的索引计数估算；当前模型成功导图记录不足，采用每次 AI 调用 25–120 秒的参考范围')
|
||||
scan_low, scan_high = (2 + libraries, 15 + libraries * 4) if mode == 'quick' else (2 + total_chunks / 1500, 10 + total_chunks / 150)
|
||||
stages = [{'id': 'scanning', 'label': '检索相关资料' if mode == 'quick' else '核对所选资料与阅读记录',
|
||||
'seconds_low': math.ceil(scan_low), 'seconds_high': math.ceil(scan_high)},
|
||||
{'id': 'reading', 'label': '并行阅读与汇总',
|
||||
'seconds_low': math.ceil((reading_waves + reducing_waves) * read_low),
|
||||
'seconds_high': math.ceil((reading_waves + reducing_waves) * read_high)}]
|
||||
if mode == 'reuse':
|
||||
stages.append({'id': 'rechecking', 'label': '按主题回查原文', 'seconds_low': 3, 'seconds_high': 45})
|
||||
basis += f"；可复用约 {reuse['reusable_files']} 个文件、{reuse['reusable_chunks']} 段的已读底稿，开始后核验精确内容"
|
||||
stages.append({'id': 'planning', 'label': '生成推荐天数与课程方案',
|
||||
'seconds_low': math.ceil(call_low), 'seconds_high': math.ceil(call_high)})
|
||||
if not total_chunks:
|
||||
for stage in stages:
|
||||
stage['seconds_low'] = stage['seconds_high'] = 0
|
||||
basis = '所选范围尚无已索引片段，请调整分类选择或先完成资料索引；当前无法估算评估耗时'
|
||||
else:
|
||||
basis += f'；阅读按最多 {parallelism} 批并行估算，模型限流、排队、重试和其他运行任务会延长时间；不含后续逐日导图'
|
||||
new_batches = max(0, reading_batches - reuse['reusable_batches'])
|
||||
result = {'mode': mode, 'total_chunks': total_chunks, 'total_files': total_files,
|
||||
'parallelism': parallelism, **reuse, 'new_chunks': total_chunks - reuse['reusable_chunks'],
|
||||
'new_batches': new_batches, 'estimated_batches': reading_batches + reducing_batches,
|
||||
'estimated_reading_batches': reading_batches, 'estimated_reducing_batches': reducing_batches,
|
||||
'seconds_low': sum(stage['seconds_low'] for stage in stages),
|
||||
'seconds_high': sum(stage['seconds_high'] for stage in stages), 'basis': basis, 'stages': stages,
|
||||
'per_day_seconds_low': math.ceil(call_low), 'per_day_seconds_high': math.ceil(call_high * 2),
|
||||
'history_samples': len(durations), 'reading_history_samples': len(read_durations),
|
||||
'reading_call_seconds_low': read_low, 'reading_call_seconds_high': read_high,
|
||||
'kb_category_ids': kb_category_ids or {}, 'batch_chars': batch_chars}
|
||||
with _ESTIMATE_CACHE_LOCK:
|
||||
_ESTIMATE_CACHE[cache_key] = (time.monotonic(), result)
|
||||
# Avoid retaining old selections indefinitely in a long-lived worker.
|
||||
if len(_ESTIMATE_CACHE) > 128:
|
||||
cutoff = time.monotonic() - _ESTIMATE_CACHE_TTL
|
||||
for key, (created, _) in list(_ESTIMATE_CACHE.items()):
|
||||
if created < cutoff:
|
||||
_ESTIMATE_CACHE.pop(key, None)
|
||||
return copy.deepcopy(result)
|
||||
|
||||
|
||||
def update_timing(coverage, current_epoch=None):
|
||||
"""Monotone work-based progress. Elapsed time affects ETA, never completion."""
|
||||
data = dict(coverage or {})
|
||||
current = time.time() if current_epoch is None else current_epoch
|
||||
estimate = data.get('estimate')
|
||||
started = data.get('started_epoch')
|
||||
if not isinstance(estimate, dict) or not started:
|
||||
return data # Historical drafts have no trustworthy estimate.
|
||||
stage = data.get('stage', 'scanning')
|
||||
stopped = stage in ('complete', 'failed', 'cancelled')
|
||||
if stopped:
|
||||
current = data.setdefault('finished_epoch', current)
|
||||
elapsed = max(0, current - started)
|
||||
data['elapsed_seconds'] = round(elapsed, 1)
|
||||
if stage == 'complete':
|
||||
data.update(progress_percent=100, remaining_seconds_low=0, remaining_seconds_high=0, overdue=False)
|
||||
return data
|
||||
if stopped:
|
||||
data.pop('remaining_seconds_low', None)
|
||||
data.pop('remaining_seconds_high', None)
|
||||
data['overdue'] = False
|
||||
return data
|
||||
stages = {item['id']: item for item in estimate.get('stages', [])}
|
||||
scanning = stages.get('scanning', {})
|
||||
planning = stages.get('planning', {})
|
||||
call_low = max(1, planning.get('seconds_low') or 25)
|
||||
call_high = max(call_low, planning.get('seconds_high') or 120)
|
||||
reading_low = max(1, estimate.get('reading_call_seconds_low', call_low))
|
||||
reading_high = max(reading_low, estimate.get('reading_call_seconds_high', call_high))
|
||||
total = max(estimate.get('total_chunks', 0), data.get('scanned_chunks', 0), 1)
|
||||
scanned = max(0, data.get('scanned_chunks', 0))
|
||||
read = max(0, data.get('read_chunks', 0))
|
||||
done = max(0, data.get('processed_batches', data.get('batches_done', 0)))
|
||||
estimated_batches = estimate.get('estimated_batches', 0)
|
||||
timings = data.get('timing') or {}
|
||||
actual_calls = timings.get('summary_calls_completed', 0)
|
||||
if actual_calls >= 2:
|
||||
average = timings.get('summary_model_seconds', 0) / actual_calls
|
||||
reading_low, reading_high = max(1, average * .65), max(5, average * 1.7)
|
||||
in_flight = max(0, current - timings.get('current_model_started_epoch', current))
|
||||
stage_elapsed = max(0, current - data.get('stage_started_epoch', started))
|
||||
if stage == 'planning':
|
||||
progress = 95 if data.get('mode') != 'quick' else 75
|
||||
low, high = max(1, call_low - stage_elapsed), max(1, call_high - stage_elapsed)
|
||||
overdue = stage_elapsed > call_high
|
||||
elif stage == 'rechecking':
|
||||
progress = 92
|
||||
low, high = max(1, 3 - stage_elapsed) + call_low, max(1, 45 - stage_elapsed) + call_high
|
||||
overdue = stage_elapsed > 45
|
||||
elif data.get('mode') == 'quick':
|
||||
progress = 0 # Vector retrieval has no per-item completion callback.
|
||||
low = max(1, scanning.get('seconds_low', 2) - elapsed) + call_low
|
||||
high = max(1, scanning.get('seconds_high', 20) - elapsed) + call_high
|
||||
overdue = elapsed > scanning.get('seconds_high', 20)
|
||||
else:
|
||||
raw_total = max(1, (scanned if stage != 'scanning' else total) - data.get('duplicate_chunks', 0))
|
||||
batch_fraction = min(1, done / max(done + 1, estimated_batches, 1))
|
||||
progress = 10 * min(1, scanned / total) + 65 * min(1, read / raw_total) + 15 * batch_fraction
|
||||
remaining_batches = max(0, estimated_batches - done)
|
||||
if data.get('mode') == 'reuse':
|
||||
# Expected hits are never billed as fresh model calls in the ETA.
|
||||
# Actual progress still advances only when each cache entry is verified.
|
||||
expected_raw = data.get('total_read_batches') or estimate.get('estimated_reading_batches', 0)
|
||||
raw_remaining = max(0, raw_total - read)
|
||||
reusable_remaining = max(0, estimate.get('reusable_chunks', 0) - data.get('reused_read_chunks', 0))
|
||||
fresh_remaining = max(0, raw_remaining - reusable_remaining)
|
||||
remaining_batches = math.ceil(fresh_remaining * expected_raw / raw_total) + max(
|
||||
0, estimate.get('estimated_reducing_batches', 0) - timings.get('reducing_calls_completed', 0))
|
||||
parallelism = max(1, data.get('effective_parallelism', data.get('parallelism', estimate.get('parallelism', 1))))
|
||||
# After real reading samples, use actual scanned size and throughput; source counters
|
||||
# remain authoritative even when different libraries alternate reading/reducing.
|
||||
if data.get('mode') != 'reuse' and read > 0 and done >= 2 and stage != 'scanning':
|
||||
raw_remaining = max(0, raw_total - read)
|
||||
reading_calls = max(1, data.get('new_read_batches', 0) + data.get('cached_batches', 0), timings.get('reading_calls_completed', 0))
|
||||
reading_remaining = math.ceil(raw_remaining * reading_calls / read)
|
||||
reduction_remaining = max(0, estimate.get('estimated_reducing_batches', 0) - timings.get('reducing_calls_completed', 0))
|
||||
remaining_batches = reading_remaining + reduction_remaining
|
||||
# The current call has not returned, so at least this batch still remains.
|
||||
if timings.get('current_model_started_epoch'):
|
||||
remaining_batches = max(1, remaining_batches)
|
||||
scan_left_low = max(0, scanning.get('seconds_low', 2) - elapsed) if stage == 'scanning' else 0
|
||||
scan_left_high = max(0, scanning.get('seconds_high', 10) - elapsed) if stage == 'scanning' else 0
|
||||
low = scan_left_low + max(0, math.ceil(remaining_batches / parallelism) * reading_low - min(in_flight, reading_low)) + call_low
|
||||
high = scan_left_high + max(0, math.ceil(remaining_batches / parallelism) * reading_high - min(in_flight, reading_high)) + call_high
|
||||
if data.get('mode') == 'reuse':
|
||||
low += stages.get('rechecking', {}).get('seconds_low', 3)
|
||||
high += stages.get('rechecking', {}).get('seconds_high', 45)
|
||||
overdue = (stage == 'scanning' and elapsed > scanning.get('seconds_high', 10)) or in_flight > reading_high
|
||||
data.update(progress_percent=round(min(99, max(data.get('progress_percent', 0), progress)), 1),
|
||||
remaining_seconds_low=math.ceil(low), remaining_seconds_high=math.ceil(max(low, high)),
|
||||
overdue=bool(overdue), estimated_total_chunks=estimate.get('total_chunks', 0),
|
||||
estimated_total_batches=estimated_batches)
|
||||
return data
|
||||
Reference in New Issue
Block a user