business-assistant/src/step3_question_answer/streaming.py

"""Request-scoped public text streaming; internal model fields never leave here."""

from contextlib import contextmanager
from contextvars import ContextVar
import json
from http.client import HTTPException
import time
import urllib.error
import urllib.request
from uuid import uuid4

from src.step1_knowledge_prepare.summarize import RequestFailure, deepseek_response
from .timing import mark


_session = ContextVar('public_reply_stream', default=None)
_field = ContextVar('public_reply_field', default=None)


class ReplyStream:
    def __init__(self, emit, stopped):
        self.emit, self.stopped = emit, stopped
        self.active = {}
        self.prefixes = {}
        self.texts = {}
        self.published = {'source': set(), 'item': set()}

    def recommendation(self, sources, items):
        """Publish finalized evidence/cards once, before public prose starts."""
        for event, rows, key in (('source', sources, 'id'), ('item', items, 'source_id')):
            for row in rows:
                self.check()
                identity = row[key]
                if identity not in self.published[event]:
                    self.emit(event, row)
                    self.published[event].add(identity)

    def check(self):
        if self.stopped.is_set():
            raise RequestFailure('stream_cancelled', False)

    def start(self, target):
        self.check()
        self.abort(target, 'regenerating')
        item = {'stream_id': uuid4().hex, 'target': target, 'sequence': 0}
        self.active[target] = item
        self.texts[target] = ''
        self.emit('answer_start', {**item, 'provisional': True})
        if self.prefixes.get(target):
            self.delta(target, item['stream_id'], self.prefixes[target])
        return item['stream_id']

    def delta(self, target, identity, text):
        self.check()
        item = self.active.get(target)
        if not item or item['stream_id'] != identity:
            raise RequestFailure('stream_cancelled', False)
        item['sequence'] += 1
        self.texts[target] += text
        self.emit('answer_delta', {**item, 'text': text})

    def complete(self, target, text, accepted=True):
        item = self.active.get(target)
        if item is None:
            return {}
        draft = self.texts[target]
        # Program additions are known only after validation. Append them once;
        # replacements (fallbacks, bound links, edits) remain authoritative end.
        if accepted and text.startswith(draft) and len(text) > len(draft):
            self.delta(target, item['stream_id'], text[len(draft):])
        self.active.pop(target)
        self.texts.pop(target, None)
        self.emit('answer_end', {**item, 'text': text, 'operation': 'replace',
                                 'status': 'completed' if accepted else 'fallback', 'provisional': False})
        return {'stream_id': item['stream_id']}

    def abort(self, target, reason='generation_failed'):
        item = self.active.pop(target, None)
        self.texts.pop(target, None)
        if item:
            self.emit('answer_abort', {**item, 'reason': reason, 'operation': 'discard'})


@contextmanager
def reply_stream(emit, stopped):
    session = ReplyStream(emit, stopped)
    token = _session.set(session)
    try:
        yield session
    finally:
        for target in list(session.active):
            session.abort(target, 'not_finalized')
        _session.reset(token)


def current_stream():
    return _session.get()


@contextmanager
def public_text(target='summary', field='text', gate=None):
    token = _field.set((target, field, gate))
    try:
        yield
    except Exception:
        if _session.get() is not None:
            _session.get().abort(target)
        raise
    finally:
        _field.reset(token)


class JsonTextField:
    """Incrementally decode one top-level string, including split surrogate pairs.

    This is a display extractor, not a JSON validator. Full JSON and references
    are still validated by the caller before the final replacement is emitted.
    """
    def __init__(self, field, emit, gate=None):
        self.field, self.emit, self.gate = field, emit, gate
        self.depth = 0
        self.expect_key = False
        self.key = None
        self.string = None
        self.raw = ''
        self.escape = False
        self.pending = ''
        self.fields = {}
        self.finished = False
        self.seen_target = False

    def _decode(self, final=False):
        result, offset = [], 0
        while offset < len(self.pending):
            char = self.pending[offset]
            width = 1
            if char == '\\':
                if offset + 1 >= len(self.pending):
                    break
                width = 6 if self.pending[offset + 1] == 'u' else 2
                if offset + width > len(self.pending):
                    break
                token = self.pending[offset:offset + width]
                if width == 6:
                    try:
                        number = int(token[2:], 16)
                    except ValueError:
                        raise RequestFailure('invalid_response_json', False) from None
                    if 0xD800 <= number <= 0xDBFF:
                        if offset + 12 > len(self.pending):
                            break
                        width = 12
                        token = self.pending[offset:offset + width]
                try:
                    char = json.loads('"' + token + '"')
                except (ValueError, UnicodeError):
                    raise RequestFailure('invalid_response_json', False) from None
            if any(0xD800 <= ord(c) <= 0xDFFF for c in char):
                raise RequestFailure('invalid_response_json', False)
            result.append(char)
            offset += width
        self.pending = self.pending[offset:]
        if final and self.pending:
            raise RequestFailure('invalid_response_json', False)
        text = ''.join(result)
        if text and (self.gate is None or self.gate(self.fields)):
            self.emit(text)

    def feed(self, chunk):
        if self.finished:
            return
        for char in chunk:
            if self.string is not None:
                if char == '"' and not self.escape:
                    kind = self.string
                    self.string = None
                    if kind == 'key':
                        self.key = json.loads('"' + self.raw + '"')
                        self.expect_key = False
                    elif kind == 'target':
                        self._decode(final=True)
                        self.finished = True
                        break
                    elif kind == 'value':
                        self.fields[self.key] = json.loads('"' + self.raw + '"')
                    self.raw = ''
                    continue
                if self.string == 'target':
                    self.pending += char
                elif self.string in ('key', 'value'):
                    self.raw += char
                if self.escape:
                    self.escape = False
                elif char == '\\':
                    self.escape = True
                continue
            if char == '"':
                if self.depth == 1 and self.expect_key:
                    self.string = 'key'
                elif self.depth == 1 and self.key == self.field:
                    if self.seen_target:
                        raise RequestFailure('invalid_response_json', False)
                    self.seen_target = True
                    self.string = 'target'
                else:
                    self.string = 'value' if self.depth == 1 else 'other'
                self.raw = ''
                self.escape = False
            elif char in '{[':
                self.depth += 1
                if self.depth == 1:
                    self.expect_key = True
            elif char in '}]':
                self.depth -= 1
            elif char == ',' and self.depth == 1:
                self.expect_key = True
                self.key = None
        if self.string == 'target':
            self._decode()


def stream_transport(payload, config, on_content, check=lambda: None):
    """Read actual provider SSE deltas and reconstruct its normal response shape."""
    total, event, parts, finish, done = 0, [], [], None, False
    started = time.monotonic()
    try:
        check()
        with deepseek_response({**payload, 'stream': True}, config, stream=True, check=check) as response:
            for line in response_lines(response, check, started + config.timeout):
                check()
                if time.monotonic() - started > config.timeout:
                    raise RequestFailure('request_timeout')
                total += len(line)
                if total > 4_000_000 or len(line) > 1_000_000:
                    raise RequestFailure('response_too_large', False)
                if not line:
                    break
                line = line.decode('utf-8').rstrip('\r\n')
                if line.startswith('data:'):
                    event.append(line[5:].lstrip(' '))
                elif not line and event:
                    data, event = '\n'.join(event), []
                    if data == '[DONE]':
                        if done:
                            raise RequestFailure('invalid_response_json', False)
                        done = True
                        continue
                    if done:
                        raise RequestFailure('invalid_response_json', False)
                    value = json.loads(data)
                    if not isinstance(value, dict) or 'error' in value or not isinstance(value.get('choices'), list):
                        raise RequestFailure('invalid_response_json', False)
                    if not value['choices']:  # optional usage-only frame
                        continue
                    if len(value['choices']) != 1:
                        raise RequestFailure('invalid_response_json', False)
                    choice = value['choices'][0]
                    if not isinstance(choice, dict) or choice.get('index', 0) != 0:
                        raise RequestFailure('invalid_response_json', False)
                    delta = choice.get('delta', {})
                    if not isinstance(delta, dict):
                        raise RequestFailure('invalid_response_json', False)
                    text = delta.get('content')
                    if text is not None and not isinstance(text, str):
                        raise RequestFailure('invalid_response_json', False)
                    if text:
                        if finish is not None:
                            raise RequestFailure('invalid_response_json', False)
                        parts.append(text)
                        on_content(text)
                    reason = choice.get('finish_reason')
                    if reason is not None:
                        if not isinstance(reason, str) or finish is not None:
                            raise RequestFailure('invalid_response_json', False)
                        finish = reason
        if not done or finish is None or event:
            raise RequestFailure('stream_incomplete', False)
        return {'choices': [{'finish_reason': finish, 'message': {'content': ''.join(parts)}}]}
    except urllib.error.HTTPError as exc:
        code = exc.code
        exc.close()
        raise RequestFailure(f'http_{code}', False) from None
    except TimeoutError:
        raise RequestFailure('request_timeout') from None
    except urllib.error.URLError as exc:
        raise RequestFailure('request_timeout' if isinstance(exc.reason, TimeoutError) else 'network_failure') from None
    except OSError:
        raise RequestFailure('network_failure') from None
    except HTTPException:
        raise RequestFailure('response_incomplete', False) from None
    except (ValueError, UnicodeError):
        raise RequestFailure('invalid_response_json', False) from None


def response_lines(response, check=lambda: None, deadline=None):
    """Incremental UTF-8-safe line framing; read1 also verifies HTTP chunks."""
    pending = b''
    total = 0
    while True:
        check()
        if deadline is not None and time.monotonic() > deadline:
            raise RequestFailure('request_timeout')
        chunk = response.read1(4096)
        if not chunk:
            if getattr(response, 'length', None) not in (None, 0):
                raise RequestFailure('response_incomplete', False)
            if pending:
                raise RequestFailure('stream_incomplete', False)
            return
        total += len(chunk)
        if total > 4_000_000:
            raise RequestFailure('response_too_large', False)
        pending += chunk
        while b'\n' in pending:
            line, pending = pending.split(b'\n', 1)
            if len(line) > 1_000_000:
                raise RequestFailure('response_too_large', False)
            yield line + b'\n'
        if len(pending) > 1_000_000:
            raise RequestFailure('response_too_large', False)


def model_response(payload, config, fallback):
    session, spec = _session.get(), _field.get()
    if session is None or spec is None:
        return fallback(payload, config)
    target, field, gate = spec
    identity = None
    pending = ''
    def emit(text):
        nonlocal identity, pending
        # Finance URLs are bound to validated source IDs only after generation.
        # Hold split internal markers and omit them from provisional text.
        markers = {'[[finance_detail:': ']]', '[[service_detail:': ']]',
                   '[[knowledge_image:': ']]', '<image-scope>': '</image-scope>'}
        pending += text
        visible = ''
        while pending:
            opening = next((marker for marker in markers if pending.startswith(marker)), None)
            if opening is not None:
                closing = markers[opening]
                end = pending.find(closing)
                if end < 0:
                    break
                pending = pending[end + len(closing):]
            elif any(marker.startswith(pending) for marker in markers):
                break
            else:
                visible += pending[0]
                pending = pending[1:]
        if not visible:
            return
        if identity is None:
            identity = session.start(target)
        session.delta(target, identity, visible)
        mark('first_model_text')
    extractor = JsonTextField(field, emit, gate)
    try:
        return stream_transport(payload, config, extractor.feed, session.check)
    except Exception:
        if identity is not None:
            session.abort(target)
        raise