business-assistant/src/step4_api/app.py

"""SSE over POST with bounded workers, buffering and per-session exclusion.

Run ONE process per Redis session namespace. A disconnected request retains its
slot and session exclusion until the synchronous graph worker actually exits.
"""

import asyncio
import json
import logging
import re
import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from queue import Empty, Full, Queue
from threading import Event, Lock
from uuid import uuid4

from starlette.applications import Starlette
from starlette.requests import ClientDisconnect
from starlette.responses import JSONResponse, StreamingResponse
from starlette.routing import Route

from .config import load_settings
from .runtime import Runtime
from .company_classification import ClassificationService, classify_company
from .clock import shanghai_now
from .logging import configure_logging, log, safe_code, session_ref, summary_diagnostics, candidate_counts, validation_counts
from .progress import STAGES, progress
from .attachments import AttachmentError, resolve_arrays
from .memory import Admission
from src.step3_question_answer.timing import Timings, bind


def sse(event, data):
    return f'event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n'


class Job:
    def __init__(self, thread_id, timeout, request_id=None, received_at=None, received_monotonic=None):
        self.thread_id = thread_id
        self.request_id = request_id or uuid4().hex
        self.received_at = received_at or shanghai_now()
        self.session_ref = session_ref(thread_id)
        self.outcome = 'running'
        self.events = Queue(maxsize=8)
        self.stopped = Event()
        self.finished = Event()
        self.started_at = time.monotonic()
        self.timings = Timings(self.started_at if received_monotonic is None else received_monotonic)
        self.deadline = self.started_at + timeout

    def emit(self, event, data):
        if self.stopped.is_set():
            return
        if event == 'progress' and data.get('stage') in STAGES and data.get('status') in ('running', 'completed', 'waiting_input', 'failed'):
            info = progress(data['stage'], data['status'])
            self.log('stage', info.pop('message'), **info)
        elif event == 'error':
            self.outcome = 'failed'
            self.log('request_failed', '本轮处理失败', level=logging.ERROR, code=safe_code(data.get('code')))
        elif event == 'result':
            recommendation = data.get('recommendation') or {}
            errors = data.get('errors') or []
            analysis_errors = recommendation.get('analysis_errors') or []
            summary_status = (recommendation.get('summary') or {}).get('status')
            highlight_errors = recommendation.get('support_highlight_errors') or []
            degraded = bool(errors or analysis_errors or highlight_errors or recommendation.get('complete') is False or summary_status == 'fallback')
            if degraded:
                self.outcome = 'degraded'
            self.log('result_ready', '本轮结果已生成', level=logging.WARNING if degraded else logging.INFO,
                     degraded=degraded, error_count=len(errors), analysis_error_count=len(analysis_errors),
                     candidate_outcome_counts=candidate_counts(recommendation.get('candidate_diagnostics')),
                     support_highlight_error_codes=[safe_code(e.get('code')) for e in highlight_errors[:16]],
                     validation_error_counts=validation_counts(recommendation.get('candidate_diagnostics')),
                     error_codes=[safe_code(c) for c in errors[:16]],
                     analysis_error_codes=[safe_code(e.get('code')) for e in analysis_errors[:16]],
                     summary_error_code=safe_code((recommendation.get('summary') or {}).get('error')) if (recommendation.get('summary') or {}).get('error') else None,
                     summary_diagnostics=summary_diagnostics((recommendation.get('summary') or {}).get('diagnostics')),
                     planning_error_code=safe_code((recommendation.get('search_plan') or {}).get('error')) if (recommendation.get('search_plan') or {}).get('error') else None,
                     summary_status=summary_status if summary_status in ('generated', 'fallback', 'no_recommendations') else None)
        elif event == 'done':
            if data.get('status') == 'needs_input':
                self.outcome = 'needs_input'
            elif self.outcome == 'running':
                self.outcome = 'completed'
        frame = sse(event, {'thread_id': self.thread_id, 'request_id': self.request_id, 'data': data})
        enqueue_started = self.timings.clock()
        while not self.stopped.is_set():
            if time.monotonic() >= self.deadline:
                self.stopped.set()
                return
            try:
                self.events.put(frame, timeout=0.1)
                self.timings.add('event_queue', self.timings.clock() - enqueue_started)
                if event == 'item':
                    self.timings.mark('first_card')
                elif event in ('answer_delta', 'answer', 'summary') and isinstance(data.get('text'), str) and data['text'].strip():
                    self.timings.mark('first_text')
                elif event == 'result':
                    self.timings.mark('result_ready')
                elif event in ('done', 'error'):
                    self.timings.mark('terminal_ready')
                return
            except Full:
                continue

    def log(self, event, message, **fields):
        log(event, message, request_id=self.request_id, session_ref=self.session_ref,
            elapsed_ms=round((time.monotonic() - self.started_at) * 1000), **fields)


class Service:
    def __init__(self, runtime, settings):
        self.runtime, self.settings = runtime, settings
        self.admission = getattr(runtime, 'admission', None) or Admission(settings)
        self.pool = ThreadPoolExecutor(max_workers=settings.max_concurrent, thread_name_prefix='chat')
        self.lock = Lock()
        self.active = {}
        self.closing = False

    def submit(self, thread_id, question, request_id=None, received_at=None, pictures=None, files=None, received_monotonic=None):
        with self.lock:
            if thread_id in self.active:
                return None, 409, 'thread_busy'
            if self.closing or len(self.active) >= self.settings.max_concurrent:
                return None, 429, 'server_busy'
            lease, error = self.admission.acquire()
            if error:
                return None, 503 if error == 'memory_status_unavailable' else 429, error
            try:
                job = Job(thread_id, self.settings.stream_timeout_seconds, request_id, received_at, received_monotonic)
                job.memory_lease = lease
                self.active[thread_id] = job
                job.log('request_accepted', '请求已接纳', active_requests=len(self.active))
                self.pool.submit(self._run, job, question, pictures or [], files or [])
            except Exception:
                self.active.pop(thread_id, None)
                lease.release()
                return None, 503, 'service_unavailable'
        return job, 200, None

    def _run(self, job, question, pictures=None, files=None):
        job.timings.add('worker_queue', job.timings.clock() - job.started_at)
        job.timings.mark('worker_started')
        job.log('worker_started', '开始处理问题')
        try:
            extra = {'pictures': pictures, 'files': files} if pictures or files else {}
            with bind(job.timings):
                self.runtime.run(job.thread_id, question, job.emit, job.stopped,
                                 request_id=job.request_id, received_at=job.received_at, **extra)
        except Exception:
            job.emit('error', {'code': 'processing_failed', 'message': '处理失败，请稍后重试或新建会话。'})
        finally:
            # Never free the lock merely because the HTTP caller disconnected.
            with self.lock:
                self.active.pop(job.thread_id, None)
            job.memory_lease.release()
            job.log('request_performance', '本轮耗时统计', metrics=job.timings.snapshot(),
                    thread_id=job.thread_id,
                    outcome=job.outcome if job.outcome != 'running' else 'stopped' if job.stopped.is_set() else 'incomplete')
            job.log('worker_finished', '后台处理结束，会话锁已释放',
                    outcome=job.outcome if job.outcome != 'running' else 'stopped' if job.stopped.is_set() else 'incomplete')
            job.finished.set()

    async def stream(self, job):
        heartbeat = time.monotonic()
        reason = 'client_disconnected'
        try:
            yield sse('accepted', {'thread_id': job.thread_id, 'request_id': job.request_id,
                                   'data': {'message': '收到，我先了解一下您的问题。'}})
            while True:
                if time.monotonic() >= job.deadline:
                    reason = 'timeout'
                    job.log('stream_timeout', '流响应超时', level=logging.WARNING, code='stream_timeout')
                    yield sse('error', {'thread_id': job.thread_id, 'request_id': job.request_id,
                                        'data': {'code': 'stream_timeout'}})
                    return
                try:
                    frame = job.events.get_nowait()
                except Empty:
                    if job.finished.is_set():
                        if reason != 'terminal_event':
                            reason = 'worker_ended_without_terminal'
                        return
                    if time.monotonic() - heartbeat >= self.settings.heartbeat_seconds:
                        yield sse('heartbeat', {'thread_id': job.thread_id, 'request_id': job.request_id,
                                               'data': {'status': 'processing',
                                                        'message': '还在为您准备答复，请再稍等一会儿。',
                                                        'elapsed_seconds': int(time.monotonic() - job.started_at)}})
                        heartbeat = time.monotonic()
                    await asyncio.sleep(0.02)
                else:
                    if frame.startswith(('event: done\n', 'event: error\n')):
                        reason = 'terminal_event'
                    yield frame
        finally:
            job.stopped.set()
            job.log('stream_closed', '流连接结束', reason=reason, worker_finished=job.finished.is_set())

    def close(self):
        with self.lock:
            self.closing = True
            for job in self.active.values():
                job.stopped.set()
        self.pool.shutdown(wait=True)
        self.runtime.close()


def create_app(settings=None, runtime_factory=Runtime):
    settings = settings or load_settings()

    @asynccontextmanager
    async def lifespan(app):
        configure_logging()
        log('service_starting', '正在初始化问答服务')
        try:
            runtime = await asyncio.to_thread(runtime_factory, settings)
        except Exception:
            log('service_startup_failed', '问答服务初始化失败', level=logging.ERROR)
            raise
        app.state.service = Service(runtime, settings)
        classify = getattr(runtime, 'classify_company', None)
        app.state.classification_service = ClassificationService(
            classify, settings.max_concurrent, app.state.service.admission) if classify else None
        app.state.classification_timeout = settings.stream_timeout_seconds
        log('service_ready', '问答服务已就绪', max_concurrent_requests=settings.max_concurrent)
        try:
            yield
        finally:
            log('service_stopping', '正在停止问答服务')
            if app.state.classification_service is not None:
                await asyncio.to_thread(app.state.classification_service.close)
            await asyncio.to_thread(app.state.service.close)
            log('service_stopped', '问答服务已停止')

    async def chat(request):
        request_id = uuid4().hex
        received_at = time.monotonic()
        received_time = shanghai_now()
        log('request_received', '收到问答请求', request_id=request_id, method='POST', path='/api/chat')

        def reject(code, status, thread_id=None, retry=False, detail=None):
            log('request_rejected', '请求未被接纳', level=logging.WARNING, request_id=request_id,
                session_ref=session_ref(thread_id) if thread_id is not None else None,
                status_code=status, code=code, elapsed_ms=round((time.monotonic() - received_at) * 1000))
            headers = {'X-Request-ID': request_id}
            if retry:
                headers['Retry-After'] = '2'
            return JSONResponse({'error': code, **(detail or {})}, status_code=status, headers=headers)

        if request.headers.get('content-type', '').split(';')[0].strip().lower() != 'application/json':
            return reject('application_json_required', 415)
        try:
            body = bytearray()
            async with asyncio.timeout(15):
                async for chunk in request.stream():
                    if len(body) + len(chunk) > 65536:
                        return reject('request_too_large', 413)
                    body.extend(chunk)
            value = json.loads(body)
            if (not isinstance(value, dict) or not {'thread_id', 'question'} <= set(value)
                    or set(value) - {'thread_id', 'question', 'pictures', 'files'}):
                raise ValueError()
            pictures, files = value.get('pictures', []), value.get('files', [])
            pictures, files = resolve_arrays(pictures, files, settings.attachment_base_url)
            thread_id, question = value['thread_id'], value['question']
            if (not isinstance(thread_id, str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,128}', thread_id)
                    or not isinstance(question, str) or not question.strip() or len(question) > 8000):
                raise ValueError()
        except TimeoutError:
            return reject('request_timeout', 408)
        except AttachmentError as exc:
            kind, index = getattr(exc, 'kind', None), getattr(exc, 'index', None)
            if kind and index:
                title = ('图片' if kind == 'picture' else '附件') + str(index)
                return reject('invalid_request', 422, detail={
                    'message': title + '的链接格式无效，请更换后重新提问。',
                    'attachment': {'kind': kind, 'index': index, 'title': title,
                                   'error': 'invalid_attachment_url'}})
            return reject('invalid_request', 422)
        except (ValueError, UnicodeError, ClientDisconnect):
            return reject('invalid_request', 422)
        job, status, error = request.app.state.service.submit(thread_id, question, request_id, received_time,
                                                            pictures=pictures, files=files, received_monotonic=received_at)
        if error:
            return reject(error, status, thread_id, retry=True)
        return StreamingResponse(request.app.state.service.stream(job), media_type='text/event-stream',
                                 headers={'Cache-Control': 'no-cache, no-transform', 'X-Accel-Buffering': 'no',
                                          'X-Request-ID': request_id})

    async def health(request):
        return JSONResponse({'status': 'ready'})

    return Starlette(routes=[Route('/api/chat', chat, methods=['POST']),
                             Route('/api/company/classify', classify_company, methods=['POST']),
                             Route('/health', health)], lifespan=lifespan)