"""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