business-assistant/src/step4_api/runtime.py

"""Bridge synchronous, checkpointed graph execution to validated public events."""

import json
import re
from threading import BoundedSemaphore
from time import monotonic
from uuid import uuid4

from langgraph.types import Command

from src.step3_question_answer.__main__ import LazyClient
from src.step3_question_answer.context import RedisContextSaver, load_context_config
from src.step3_question_answer.graph import build_graph
from src.step3_question_answer.intent import deepseek_analyzer
from src.step3_question_answer.recommendation import JsonModel, Recommender, analysis_concurrency, synthesis_model
from src.step3_question_answer.retrieval import Retriever
from src.step3_question_answer.rerank import create_reranker
from src.step3_question_answer.session_archive import ArchivedTurn
from src.step3_question_answer.streaming import reply_stream, current_stream
from src.step5_session_analysis.analyzer import SessionAnalyzer
from src.company_classification.classifier import CompanyClassifier
from src.step5_session_analysis.dms_observer import create_dms_observer

from .config import ROOT
from .clock import shanghai_now
from .progress import STAGES, progress
from .attachments import AttachmentProcessor, VisionReader
from src.step3_question_answer.timing import slot, span, measured


def limited(call, semaphore, resource='model'):
    def invoke(*args, **kwargs):
        with slot(semaphore, resource):
            return call(*args, **kwargs)
    # Preserve model budget metadata through the concurrency wrapper.
    for name in ('config', 'MAX_OUTPUT_TOKENS', 'MIN_OUTPUT_TOKENS', 'supports_context'):
        if hasattr(call, name):
            setattr(invoke, name, getattr(call, name))
    if type(getattr(call, 'MIN_OUTPUT_TOKENS', None)) is int:
        invoke.with_output_tokens = lambda limit: limited(call.with_output_tokens(limit), semaphore, resource)
    return invoke


def matches_candidate_display(text, candidate):
    """Recognize whitespace-joined display fields, without accepting prose.

    Every piece must be an exact field from the same current candidate, and
    at least one must identify the company. Status/person alone cannot select.
    """
    identity = {candidate.get(key) for key in ('KeyNo', 'Name', 'CreditCode')
                if isinstance(candidate.get(key), str) and candidate[key]}
    fields = {candidate.get(key) for key in
              ('KeyNo', 'Name', 'CreditCode', 'Status', 'OperName', 'No', 'Address', 'StartDate')
              if isinstance(candidate.get(key), str) and candidate[key]}
    if not fields:
        return False
    atom = '(?:' + '|'.join(re.escape(v) for v in sorted(fields, key=lambda v: (-len(v), v))) + ')'
    if not re.fullmatch(atom + r'(?:\s+' + atom + ')+', text):
        return False
    # A field is an identity only when it occupies a whole displayed piece.
    return any(re.search(r'(?:^|\s)' + re.escape(value) + r'(?=\s|$)', text) for value in identity)


def resume_answer(value, question):
    """Only explicit user choices approve a company; no model guesses identity."""
    text = question.strip()
    # Initial company clarification uses needs_company; search follow-ups use
    # actions, including when no candidates now carry the company_need kind.
    if value['kind'] == 'company_need' and 'result' not in value:
        return {'needs_company': text not in ('不需要', '无需公司信息', '/skip'),
                'company_hint': text}
    candidates = value['result']['candidates']
    if text in ('/skip', '跳过公司查询', '不需要', '无需公司信息'):
        return {'action': 'skip'}
    if text.isascii() and text.isdigit() and 1 <= int(text) <= len(candidates):
        return {'action': 'select', 'key_no': candidates[int(text) - 1]['KeyNo']}
    if text in ('/next', '/retry', '/cancel'):
        return {'action': text[1:]}
    if text in ('/more', '搜索更多公司', '更多的公司信息搜索'):
        return {'action': 'more'} if 'more' in value.get('actions', {}) else {'action': 'invalid'}
    # Match only this interrupt's server-side candidates. Never trust returned
    # profile fields or use fuzzy matching to approve an identity.
    if text.startswith(('{', '[')):
        try:
            selected = json.loads(text)
        except (ValueError, RecursionError):
            return {'action': 'invalid'}
        if not isinstance(selected, dict):
            return {'action': 'invalid'}
        fields = {key: selected[key] for key in ('KeyNo', 'Name', 'CreditCode') if key in selected}
        if 'key_no' in selected:
            if 'KeyNo' in fields and fields['KeyNo'] != selected['key_no']:
                return {'action': 'invalid'}
            fields['KeyNo'] = selected['key_no']
        # The public candidate schema permits a missing credit code.
        if fields.get('CreditCode', '') is None:
            fields.pop('CreditCode', None)
        if (not fields or selected.get('action', 'select') != 'select'
                or any(not isinstance(v, str) or not v.strip() for v in fields.values())):
            return {'action': 'invalid'}
        matches = [row for row in candidates if all(row.get(k) == v.strip() for k, v in fields.items())]
    else:
        matches = [row for row in candidates
                   if any(row.get(key) == text for key in ('KeyNo', 'Name', 'CreditCode'))]
        if not matches:
            matches = [row for row in candidates if matches_candidate_display(text, row)]
        if not matches:
            return {'action': 'refine', 'keyword': text}
    if len(matches) == 1:
        return {'action': 'select', 'key_no': matches[0]['KeyNo']}
    return {'action': 'invalid'}


class GraphRuntime:
    def __init__(self, graph, archive=None, attachments=None):
        self.graph = graph
        self.attachments = attachments
        self.archive = archive
        self.history = graph.checkpointer if isinstance(graph.checkpointer, RedisContextSaver) else None

    def run(self, thread_id, question, emit, stopped, *, request_id=None, received_at=None, pictures=None, files=None):
        started = monotonic()
        record = {'schema_version': 1, 'request_id': request_id or uuid4().hex,
                  'thread_id': thread_id, 'question': question, 'original_question': None,
                  'input_kind': 'unknown', 'received_at': received_at or shanghai_now(),
                  'started_at': shanghai_now(), 'finished_at': None, 'duration_ms': None,
                  'status': 'running', 'degraded': False, 'error_code': None}
        terminal = None
        turn = ArchivedTurn(self.archive, record)

        def tracked_emit(event, data):
            nonlocal terminal
            turn.capture(event, data)
            if event in ('done', 'error'):
                terminal = (event, data)
                record['status'] = data['status'] if event == 'done' else 'failed'
                record['error_code'] = data.get('code') if event == 'error' else None
                return
            if event == 'result':
                recommendation = data.get('recommendation') or {}
                record['degraded'] = bool(data.get('errors') or recommendation.get('analysis_errors') or recommendation.get('support_highlight_errors')
                                          or recommendation.get('complete') is False
                                          or (recommendation.get('summary') or {}).get('status') == 'fallback')
            emit(event, data)

        turn.save()
        try:
            self._save_request(record)
            with reply_stream(tracked_emit, stopped):
                if pictures or files:
                    self._run(thread_id, question, tracked_emit, stopped, record, pictures or [], files or [])
                else:
                    self._run(thread_id, question, tracked_emit, stopped, record)
            if terminal is None:
                record['status'] = 'stopped' if stopped.is_set() else 'incomplete'
        except Exception:
            record.update(status='failed', error_code='processing_failed')
            turn.capture('error', {'code': 'processing_failed', 'message': '处理失败，请稍后重试或新建会话。'})
            raise
        finally:
            record.update(finished_at=shanghai_now(), duration_ms=round((monotonic() - started) * 1000))
            turn.save()
            self._save_request(record)
        # Persistence must succeed before announcing completion to the caller.
        if terminal is not None:
            emit(*terminal)

    @measured('history_save')
    def _save_request(self, record):
        if self.history is not None:
            self.history.put_request(record)

    def _run(self, thread_id, question, emit, stopped, record, pictures=None, files=None):
        config = {'configurable': {'thread_id': thread_id}, 'recursion_limit': 100}
        emit('progress', progress('session'))
        with span('session_load', 'session_load'):
            previous = self.graph.get_state(config)
        emit('progress', progress('session', 'completed'))
        interrupts = [item for task in previous.tasks for item in task.interrupts]
        record['original_question'] = previous.values.get('original_question') if interrupts or previous.next else question
        record['input_kind'] = ('clarification_reply' if interrupts else
                                ('retry' if question.strip() == '/retry' else 'unfinished_turn_input')
                                if previous.next else 'new_question')
        self._save_request(record)
        if stopped.is_set():
            return
        if (interrupts or previous.next) and (pictures or files):
            emit('error', {'code': 'attachments_during_resume',
                           'message': '本轮正在等待公司确认或恢复，请先完成当前步骤，再以新问题提交附件。'})
            return
        if interrupts:
            value = Command(resume=resume_answer(interrupts[0].value, question))
        elif previous.next:
            if question.strip() != '/retry':
                # Do not silently discard a new question when recovering a failed run.
                emit('error', {'code': 'unfinished_turn', 'message': '上一轮尚未完成，请输入/retry恢复，或使用新的thread_id。'})
                return
            value = None
        else:
            records = []
            if pictures or files:
                if self.attachments is None:
                    emit('error', {'code': 'attachments_unavailable', 'message': '附件处理暂不可用。'})
                    return
                emit('progress', {'stage': 'attachments', 'status': 'running', 'message': '正在读取您提供的图片和附件……'})
                with span('attachments', 'attachments'):
                    records = self.attachments.process(pictures, files, question, stopped)
                emit('attachments', {'items': [{k: v for k, v in row.items() if k not in ('sections', 'ok')}
                                                for row in records]})
                failed = next((row for row in records if row['status'] == 'failed'), None)
                if failed:
                    message = failed['title'] + '无法读取：' + failed['message'] + '请更换后重新提问。'
                    emit('progress', {'stage': 'attachments', 'status': 'failed', 'message': message})
                    emit('error', {'code': 'attachment_processing_failed', 'message': message,
                                   'attachment': {'kind': failed['kind'], 'index': failed['index'],
                                                  'title': failed['title'], 'error': failed['error']}})
                    return
                emit('progress', {'stage': 'attachments', 'status': 'completed', 'message': '图片和附件处理已结束。'})
            value = {'question': question, 'incoming_attachments': records}
        for mode, update in self.graph.stream(value, config, stream_mode=['tasks', 'updates']):
            if stopped.is_set():
                return
            if mode == 'tasks':
                # Start notifications arrive before slow synchronous node execution.
                # Whitelist only node name/status: task input/result/error may be private.
                node = update.get('name')
                if node in STAGES:
                    if 'input' in update:
                        emit('progress', progress(node))
                    elif update.get('interrupts'):
                        emit('progress', progress(node, 'waiting_input'))
                    elif update.get('error'):
                        emit('progress', progress(node, 'failed'))
                continue
            for node, fields in update.items():
                if node == '__interrupt__':
                    continue
                if node in STAGES:
                    emit('progress', progress(node, 'completed'))
                if node == 'direct':
                    text = fields.get('response', '')
                    metadata = current_stream().complete('answer', text)
                    emit('answer', {'text': text, **metadata})
                elif node == 'analyze' and fields.get('route') != 'direct':
                    current_stream().abort('answer', 'not_used')
                elif node == 'recommend':
                    answer = fields['recommendation']
                    # Do not stream the summary (or its opening) twice. Keep
                    # independent fallback prose and appended failure notices.
                    summary = answer.get('summary') or {}
                    summary_text = str(summary.get('text') or '').strip()
                    opening = fields.get('response', '').strip()
                    if summary_text and opening in summary_text:
                        opening = ''
                    elif summary_text and opening.startswith(summary_text):
                        opening = opening[len(summary_text):].strip()
                    elif summary_text and opening.startswith(summary_text.split('\n\n', 1)[0]):
                        opening = opening[len(summary_text.split('\n\n', 1)[0]):].strip()
                    current_stream().recommendation(answer.get('sources', []), answer.get('items', []))
                    if summary_text:
                        metadata = current_stream().complete('summary', summary_text, summary.get('status') != 'fallback')
                        emit('summary', {**summary, **metadata})
                    # Independent failure notices follow the finalized summary,
                    # never interrupt its active stream or repeat its opening.
                    if opening:
                        emit('answer', {'text': opening})
        with span('session_result_load', 'session_result_load'):
            state = self.graph.get_state(config)
        pending = [item.value for task in state.tasks for item in task.interrupts]
        if pending:
            prompt = dict(pending[0])
            if prompt.get('status') in ('not_found', 'failed'):
                prompt['input_help'] = ('输入工商注册全名或统一社会信用代码重新查询；'
                                        '输入“跳过公司查询”或/skip继续，/cancel取消公司查询。'
                                        + ('输入/retry重试。' if prompt['status'] == 'failed' else ''))
                emit('answer', {'text': prompt['question']})
            else:
                prompt['input_help'] = ('输入公司名称或信用代码；不需要则输入“不需要”。'
                                    if prompt['kind'] == 'company_need' else
                                    '输入候选序号、准确全名、KeyNo或信用代码确认，也可提交候选显示信息或JSON字符串；'
                                    '重名请使用序号或KeyNo。输入新关键词可重新搜索；支持/next、/retry、/cancel。')
            if 'more' in prompt.get('actions', {}):
                prompt['input_help'] += '选择列表最后一项“搜索更多公司”可继续查询，也支持/more。'
            emit('interrupt', prompt)
        # Preserve the complete existing recommendation contract and full sources.
        response = (prompt['question'] if pending and prompt.get('status') in ('not_found', 'failed')
                    else state.values.get('response', ''))
        emit('result', {'response': response,
                        'attachments': [{k: v for k, v in row.items() if k not in ('sections', 'ok')}
                                        for row in state.values.get('attachments', [])],
                        'cipa': state.values.get('cipa') is True,
                        # Keep legacy checkpoint storage readable; expose corrected API spelling only.
                        'company_info': state.values.get('compang_info', {}),
                        'recommendation': state.values.get('recommendation', {}),
                        'errors': state.values.get('errors', [])})
        emit('done', {'status': 'needs_input' if pending else 'completed'})

    def close(self):
        pass


class Runtime(GraphRuntime):
    def __init__(self, settings):
        from .memory import Admission, require_hard_limit
        require_hard_limit()
        self.admission = Admission(settings)
        env = ROOT / '.env'
        self.saver = RedisContextSaver(load_context_config(env))
        self.retriever = None
        self.observer = None
        try:
            self.saver._call('ping')
            self.retriever = Retriever(env, ROOT / 'knowledge/prepared/ingestion/manifest.json')
            self.retriever.warmup()
            # One shared cap covers intent, planning, per-item analysis and summary.
            llm_slots = BoundedSemaphore(settings.llm_concurrent)
            model = JsonModel(env)
            self.classify_company = CompanyClassifier(limited(model, llm_slots))
            self.retriever.search = limited(self.retriever.search, BoundedSemaphore(settings.retrieval_concurrent), 'retrieval')
            super().__init__(build_graph(
                LazyClient(env), limited(deepseek_analyzer(env), llm_slots),
                checkpointer=self.saver, retriever=self.retriever,
                recommender=Recommender(limited(model, llm_slots), workers=analysis_concurrency(env),
                                        summary_model=limited(synthesis_model(env), llm_slots),
                                        reranker=create_reranker(env, JsonModel, llm_slots))),
                attachments=AttachmentProcessor(VisionReader(model.config, llm_slots)))
            self.observer = create_dms_observer(
                SessionAnalyzer(limited(model, llm_slots), model.config.model), env)
            self.observer.admission = self.admission
            self.observer.start()
        except Exception:
            self.close()
            raise RuntimeError('api_startup_failed') from None

    def close(self):
        if self.observer is not None:
            self.observer.close()
        if self.retriever is not None:
            self.retriever.close()
        self.saver.close()