business-assistant/src/step3_question_answer/graph.py

"""Company workflow: unique exact user identifiers can bypass confirmation."""

from typing import TypedDict
from langchain_core.runnables import RunnableConfig

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import interrupt

from .qcc import QccError, nonempty
from .persona import DIRECT_RESPONSES
from .dialogue import collect_history, model_context
from .timing import measured


MORE_COMPANIES_KEY = 'action:search_more_companies'


class AgentState(TypedDict, total=False):
    question: str
    original_question: str
    conversation_kind: str
    cipa: bool
    needs_company: bool | None
    company_hint: str
    page: int
    searches: int
    search_result: dict
    company_search_provider: str
    selected_company: dict | None
    profile: dict | None
    honors: dict | None
    compang_info: dict
    errors: list[str]
    route: str
    handoff: dict
    response: str
    direct_response: str
    search_plan: dict
    retrieval: dict
    recommendation: dict
    conversation_history: list[dict]
    conversation_context: dict
    previous_company: dict | None
    context_relation: str
    context_query: str
    attachments: list[dict]
    incoming_attachments: list[dict]


def build_graph(client, analyzer=None, checkpointer=None, retriever=None, recommender=None):
    """Inject analyzer(question)->intent and client for offline tests.

    New questions retain same-thread dialogue, but reuse company facts only explicitly.
    Hosting code must bind thread_id to its authenticated user; this is not an API.
    """
    saver = checkpointer if checkpointer is not None else InMemorySaver()

    def prepare(state, config: RunnableConfig):
        question = state.get('question')
        if not nonempty(question):
            raise ValueError('question_required')
        history = collect_history(state, saver, config)
        from src.step4_api.attachments import routing_context
        context = model_context(history)
        attachments = state.get('incoming_attachments', [])
        if attachments:
            context['attachments'] = routing_context(attachments, question)
        previous_company = state.get('previous_company')
        if state.get('selected_company') and state.get('profile'):
            previous_company = {key: state.get(key) for key in ('selected_company', 'profile', 'honors', 'compang_info')}
        elif state.get('company_hint'):
            # A failed/ambiguous switch must not leave an earlier company reusable.
            previous_company = None
        return {'original_question': question, 'conversation_kind': '', 'cipa': False, 'route': '',
                'attachments': attachments, 'incoming_attachments': [],
                'conversation_history': history, 'conversation_context': context,
                'previous_company': previous_company, 'context_relation': 'new_topic', 'context_query': '',
                'needs_company': None, 'company_hint': '',
                'page': 1, 'searches': 0, 'search_result': {}, 'company_search_provider': 'dms', 'selected_company': None,
                'profile': None, 'honors': None, 'compang_info': {}, 'errors': [], 'handoff': {}, 'response': '', 'direct_response': '',
                'search_plan': {}, 'retrieval': {}, 'recommendation': {}}

    def analyze(state):
        if analyzer is None:
            # --manual is an explicit operator override for business workflows.
            return {'conversation_kind': 'specific_need', 'route': 'clarify'}
        try:
            context = dict(state['conversation_context'])
            previous_company = state.get('previous_company') or {}
            identity = previous_company.get('selected_company') or {}
            context['confirmed_company'] = {key: identity[key] for key in ('Name', 'CreditCode') if identity.get(key)}
            intent = (analyzer(state['original_question'], conversation_context=context)
                      if getattr(analyzer, 'supports_context', False) is True else analyzer(state['original_question']))
            relation = intent.get('context_relation', 'new_topic')
            query = intent.get('context_query', '')
            if relation not in ('follow_up', 'new_topic') or not isinstance(query, str) or len(query) > 500:
                raise ValueError('invalid_context_resolution')
            cipa = intent.get('cipa', False)
            if type(cipa) is not bool:
                raise ValueError('invalid_cipa')
            common = {'context_relation': relation, 'cipa': cipa,
                      'context_query': query.strip() if context.get('attachments') or relation == 'follow_up' and context['turns'] else ''}
            interpretation = intent.get('term_interpretation')
            if interpretation is not None:
                if not isinstance(interpretation, dict):
                    raise ValueError('invalid_term_interpretation')
                original, interpreted = interpretation.get('original'), interpretation.get('interpreted')
                confidence = interpretation.get('confidence')
                if (not isinstance(original, str) or not original.strip() or len(original) > 40
                        or state['original_question'].count(original) != 1
                        or not isinstance(interpreted, str) or not interpreted.strip() or len(interpreted) > 40
                        or confidence not in ('high', 'uncertain')
                        or original == interpreted
                        or original in (intent.get('company_hint') or '')
                        or any(ch.isdigit() for ch in original + interpreted)):
                    raise ValueError('invalid_term_interpretation')
                if confidence == 'uncertain':
                    return {**common, 'cipa': False, 'conversation_kind': 'policy_overview',
                            'needs_company': False, 'route': 'direct',
                            'direct_response': f'您说的“{original}”是指“{interpreted}”吗？也可以直接说想解决什么问题，我再帮您查找。'}
                resolved = state['original_question'].replace(original, interpreted, 1)
                # Preserve the raw question and identity; share only a provisional
                # interpretation across planning, rerank and evidence analysis.
                common['context_query'] = resolved
                common['conversation_context'] = {**state['conversation_context'],
                    'term_interpretation': {'original': original, 'interpreted': interpreted,
                                            'query': resolved, 'status': 'hypothesis'}}
            kind = intent['conversation_kind']
            if kind not in ('smalltalk', 'policy_overview', 'specific_need'):
                raise ValueError('invalid_conversation_kind')
            if kind == 'smalltalk' and state.get('attachments'):
                # An attachment-only factual question must reach grounded synthesis.
                kind = 'policy_overview'
            if kind == 'smalltalk':
                reply = intent.get('direct_response')
                valid = isinstance(reply, str) and bool(reply.strip()) and len(reply) <= 4000
                return {**common, 'conversation_kind': kind, 'needs_company': False, 'route': 'direct',
                        'direct_response': reply.strip() if valid else '',
                        'errors': [] if valid else ['direct_response_unavailable']}
            if kind != 'specific_need':
                # Non-chat answers must come from retrieved knowledge, even if the
                # classifier incorrectly supplies an ungrounded direct_response.
                return {**common, 'conversation_kind': kind, 'needs_company': False, 'route': 'handoff'}
            needed, hint = intent['needs_company'], intent.get('company_hint', '')
            if needed is not None and type(needed) is not bool:
                raise ValueError('invalid_intent')
            reuse = intent.get('reuse_company') is True and relation == 'follow_up' and needed is True
            if reuse and identity and (not hint or hint in (identity.get('Name'), identity.get('CreditCode'))):
                return {**common, **previous_company, 'conversation_kind': kind, 'needs_company': True,
                        'company_hint': identity.get('Name', ''), 'route': 'handoff'}
            if not isinstance(hint, str) or (hint and hint not in state['original_question']):
                raise ValueError('invented_company_hint')
            return {**common, 'conversation_kind': kind, 'needs_company': needed, 'company_hint': hint, 'route':
                    'handoff' if needed is False else 'search' if needed is True and hint else 'clarify'}
        except Exception:
            # Never serialize model/network exceptions (may contain request data).
            return {'errors': ['intent_analysis_failed'], 'conversation_kind': 'policy_overview',
                    'needs_company': False, 'route': 'handoff'}

    def direct(state):
        return {'response': state.get('direct_response') or DIRECT_RESPONSES[state['conversation_kind']]}

    def clarify(state):
        answer = interrupt({'kind': 'company_need', 'question': '本次是否需要结合具体公司情况？如需要，请补充公司名称或信用代码。',
                            'schema': {'needs_company': 'boolean', 'company_hint': 'string'}})
        if not isinstance(answer, dict) or type(answer.get('needs_company')) is not bool:
            return {'route': 'clarify'}
        if answer['needs_company'] is False:
            return {'needs_company': False, 'route': 'handoff'}
        hint = answer.get('company_hint')
        if not nonempty(hint):
            return {'route': 'clarify'}
        return {'needs_company': True, 'company_hint': hint.strip(), 'route': 'search'}

    def search(state):
        if state['searches'] >= 5:
            return {'errors': state['errors'] + ['search_limit'], 'route': 'handoff'}
        try:
            result = (client.search_qcc(state['company_hint'], state['page'])
                      if state.get('company_search_provider') == 'qcc' and hasattr(client, 'search_qcc')
                      else client.search(state['company_hint'], state['page']))
            common = {'company_search_provider': result.get('provider', state.get('company_search_provider', 'dms')),
                      'errors': state['errors'] + result.get('warnings', []) +
                      (['search:' + result['error']] if result.get('error') else [])}
            hint = state['company_hint'].strip()
            matches = [row for row in result['candidates']
                       if hint and hint in (row.get('Name'), row.get('CreditCode'))]
            if not result.get('error') and len(matches) == 1 and result.get('exact_unique', True):
                return {**common, 'search_result': result, 'searches': state['searches'] + 1,
                        'selected_company': dict(matches[0]), 'route': 'basic'}
            return {**common, 'search_result': result, 'searches': state['searches'] + 1, 'route': 'confirm'}
        except QccError as exc:
            return {'search_result': {'candidates': [], 'may_have_more': False, 'error': str(exc)},
                    'errors': state['errors'] + ['search:' + str(exc)],
                    'searches': state['searches'] + 1, 'route': 'confirm'}

    def confirm(state):
        result = state['search_result']
        status = 'failed' if result.get('error') else 'found' if result['candidates'] else 'not_found'
        question = '请选择您所指的公司；也可以补充名称重新检索或取消。'
        actions = {'refine': {'keyword': '注册全名或统一社会信用代码'}, 'skip': {}, 'cancel': {}}
        if status == 'not_found':
            question = ('本次未检索到与您提供的名称或代码匹配的企业。请核对并补充工商注册全名或统一社会信用代码；'
                        '也可以输入“跳过公司查询”，先按原问题了解相关政策。')
        elif status == 'failed':
            question = ('企业查询服务暂时未能完成查询，目前无法判断是否有匹配企业。'
                        '您可以输入/retry重试，或输入“跳过公司查询”，先按原问题了解相关政策。')
            actions['retry'] = {}
        else:
            actions['select'] = {'key_no': '当前候选KeyNo'}
            if result.get('can_search_more'):
                actions['more'] = {}
                question = '请选择您所指的公司；如不在候选中，可选择“搜索更多公司”继续查询。'
            if result.get('may_have_more'):
                actions['next'] = {}
        if result.get('warnings'):
            question = '本地企业资料暂未查全，已尝试通过企查查继续查询。' + question
        display_result = result
        if status == 'found' and result.get('can_search_more'):
            # Reuse the existing candidate UI and selection protocol. This row
            # is an action, never a company fact or a stored search candidate.
            more = {key: None for key in ('CreditCode', 'StartDate', 'OperName', 'Status', 'No', 'Address')}
            more.update(KeyNo=MORE_COMPANIES_KEY, Name='搜索更多公司')
            display_result = {**result, 'candidates': [*result['candidates'], more]}
        answer = interrupt({'kind': 'company_selection' if status == 'found' else 'company_need',
                            'status': status, 'question': question,
                            'result': display_result, 'page': state['page'], 'actions': actions})
        if not isinstance(answer, dict):
            return {'route': 'confirm'}
        action = answer.get('action')
        if action == 'skip':
            return {'needs_company': False, 'route': 'handoff'}
        if action == 'cancel':
            return {'route': 'handoff'}
        if action == 'refine' and nonempty(answer.get('keyword')):
            return {'company_hint': answer['keyword'].strip(), 'page': 1,
                    'company_search_provider': 'dms', 'route': 'search'}
        if action == 'more' and result.get('can_search_more'):
            return {'company_search_provider': 'qcc', 'page': 1, 'route': 'search'}
        if action == 'next' and state['search_result'].get('may_have_more'):
            return {'page': state['page'] + 1, 'route': 'search'}
        if action == 'retry' and state['search_result'].get('error'):
            return {'route': 'search'}
        if action == 'select':
            if answer.get('key_no') == MORE_COMPANIES_KEY and result.get('can_search_more'):
                return {'company_search_provider': 'qcc', 'page': 1, 'route': 'search'}
            for row in state['search_result']['candidates']:
                if row['KeyNo'] == answer.get('key_no'):
                    return {'selected_company': dict(row), 'route': 'basic'}
        return {'route': 'confirm'}

    def basic(state):
        try:
            profile = client.basic(state['selected_company'])
            return {'profile': profile, 'route': 'honors',
                    'compang_info': {'company': state['selected_company'],
                                     'profile': profile, 'honors': None}}
        except QccError as exc:
            return {'errors': state['errors'] + ['basic:' + str(exc)], 'route': 'handoff'}

    def honors(state):
        result = client.honors(state['selected_company'])
        return {'honors': result,
                'compang_info': {'company': state['selected_company'],
                                 'profile': state['profile'], 'honors': result},
                'errors': state['errors'] +
                (['honors:' + result['error']] if result['error'] else [])}

    def handoff(state):
        ready = state['needs_company'] is False or state['profile'] is not None
        response = ('需求信息已整理，下一步需检索政策原文并核验适用条件。' if ready else
                    '公司信息尚未完整确认或获取，请补充名称后重试；目前无法据此判断企业政策资格。')
        if state['honors'] and not state['honors']['complete']:
            response += '企业荣誉信息未获取完整，不能据此排除其他资质。'
        if retriever is None or recommender is None:
            response += '本次运行未启用政策检索与推荐。'
        return {'handoff': {'original_question': state['original_question'],
                            'attachments': state.get('attachments', []),
                            'conversation_context': state['conversation_context'],
                            'context_relation': state['context_relation'], 'context_query': state['context_query'],
                            'conversation_kind': state['conversation_kind'],
                            'needs_company': state['needs_company'],
                            'company': state['selected_company'], 'profile': state['profile'],
                            'honors': state['honors'], 'errors': state['errors'],
                            'ready_for_retrieval': ready, 'policy_evidence': [],
                            'qualification_status': 'not_evaluated'}, 'response': response}

    def plan(state):
        return {'search_plan': recommender.plan(state['handoff'])}

    def retrieve(state):
        try:
            queries = list(state['search_plan']['queries'])
            if state.get('context_query') and state['context_query'] not in queries:
                queries.insert(0, state['context_query'])
            return {'retrieval': retriever.search(state['original_question'], queries)}
        except Exception:
            return {'retrieval': {'candidates': [], 'queries': [state['original_question']],
                                  'complete': False, 'failed_databases': ['retrieval_service'],
                                  'content_failures': []}}

    def recommend(state):
        result = recommender.answer({**state['handoff'], 'search_plan': state['search_plan']}, state['retrieval'])
        interpretation = state['conversation_context'].get('term_interpretation')
        if interpretation:
            notice = f'我先按“{interpretation["interpreted"]}”理解您说的“{interpretation["original"]}”；如果不是这个意思，您可以直接纠正。'
            result.setdefault('summary', {}).setdefault('notices', []).insert(0, notice)
        search_plan = state['search_plan']
        result['search_plan'] = search_plan
        if search_plan.get('error'):
            notice = '需求拆解失败，本次仅按原问题检索。'
            result['overview'] += '\n' + notice
            if result.get('summary'):
                result['summary'].setdefault('notices', []).append(notice)
            result['complete'] = result['retrieval_analysis_complete'] = False
        if 'intent_analysis_failed' in state['errors']:
            notice = '需求识别未完成，本次按原问题检索，未推测主体信息。'
            result['overview'] += '\n' + notice
            result.setdefault('summary', {}).setdefault('notices', []).append(notice)
            result['complete'] = False
        return {'recommendation': result, 'response': result['overview']}

    graph = StateGraph(AgentState)
    for name, node in [('prepare', prepare), ('analyze', analyze), ('clarify', clarify),
                       ('search', search), ('confirm', confirm), ('basic', basic),
                       ('honors', honors), ('handoff', handoff), ('direct', direct)]:
        graph.add_node(name, measured(name)(node))
    graph.add_edge(START, 'prepare')
    graph.add_edge('prepare', 'analyze')
    graph.add_edge('direct', END)
    for name, targets in [('analyze', ['clarify', 'search', 'handoff', 'direct']),
                          ('clarify', ['clarify', 'search', 'handoff']),
                          ('search', ['confirm', 'basic', 'handoff']),
                          ('confirm', ['confirm', 'search', 'basic', 'handoff']),
                          ('basic', ['honors', 'handoff'])]:
        graph.add_conditional_edges(name, lambda state: state['route'], {x: x for x in targets})
    graph.add_edge('honors', 'handoff')
    if retriever is not None and recommender is not None:
        graph.add_node('plan', measured('plan')(plan))
        graph.add_node('retrieve', measured('retrieve')(retrieve))
        graph.add_node('recommend', measured('recommend')(recommend))
        graph.add_edge('handoff', 'plan')
        graph.add_edge('plan', 'retrieve')
        graph.add_edge('retrieve', 'recommend')
        graph.add_edge('recommend', END)
    else:
        graph.add_edge('handoff', END)
    return graph.compile(checkpointer=saver)