core/chat.py

from __future__ import annotations

import itertools
import logging
import queue
import threading
import time
from typing import Any, Callable, Iterator, Optional, TypedDict

from config import MODEL, client
from tools import build_tools_list

logger = logging.getLogger(__name__)

# Monotonically-increasing epoch ids stamped on every per-session
# ``ChatState`` and bumped on reset. Streaming handlers capture the
# epoch at entry and check it between yields; if the user clicks "+"
# (``new_chat``) mid-stream, ``reset_state_in_place`` mutates the SAME
# state dict the running generator holds — bumping the epoch — so the
# generator notices on its next iteration and exits without emitting
# any further chat deltas.
_epoch_counter = itertools.count(1)


def _next_epoch() -> int:
    return next(_epoch_counter)

# Floor interval between streaming yields. The per-yield cost on the
# wire is just a small delta payload, so 60ms ≈ 16 yields/sec — close
# to one yield per browser frame, which is the natural ceiling for
# human-perceptible smoothness anyway.
_YIELD_INTERVAL = 0.06

# How often to emit a keep-alive yield when no new chunks arrive. This keeps
# the SSE/WebSocket connection alive through reverse-proxy idle timeouts
# (e.g. HuggingFace Spaces proxy). Without heartbeats, a long pause between
# reasoning and content chunks can cause the proxy to drop the connection,
# silently terminating the generator.
_HEARTBEAT_INTERVAL = 5.0

# Sentinel placed in the chunk queue when the streaming thread finishes.
_STREAM_DONE = object()


def _drain_queue(q: queue.Queue) -> list:
    """Pull every item currently in *q* without blocking. Empty list if none."""
    out: list = []
    while True:
        try:
            out.append(q.get_nowait())
        except queue.Empty:
            return out


class ChatState(TypedDict, total=False):
    messages: list[dict]
    context_start_index: int
    pending_tool_calls: list[dict]
    pending_assistant_msg: Optional[dict]
    submitted_tool_results: list[dict]
    epoch: int


def init_state() -> ChatState:
    """Fresh per-session conversation state.

    Note: there is intentionally NO server-side ``is_streaming`` flag. The
    "model is busy" signal is owned entirely by the UI: a click on Send
    instantly disables the Send button via a ``queue=False`` Gradio chain
    BEFORE the streaming generator is even queued, so a duplicate submission
    is impossible regardless of network latency or queue order.

    ``epoch`` is the cancellation token: bumped by ``reset_state_in_place``
    when the user clicks "+" mid-stream so the running generator can
    detect the reset and abandon further yields.
    """
    return {
        "messages": [],
        "context_start_index": 0,
        "pending_tool_calls": [],
        "pending_assistant_msg": None,
        "submitted_tool_results": [],
        "epoch": _next_epoch(),
    }


def reset_state_in_place(state: ChatState) -> int:
    """Reset *state* in place and bump its epoch. Returns the new epoch.

    Critical: this MUTATES the caller's dict instead of returning a fresh
    one. A streaming generator started before the reset still holds a
    reference to this same dict — the in-place mutation is what lets it
    observe the bumped epoch and stop yielding chat deltas. Returning a
    new dict (and asking Gradio to swap it into the State component)
    would leave the in-flight generator pointed at a stale dict it would
    happily keep streaming into.
    """
    state["messages"] = []
    state["context_start_index"] = 0
    state["pending_tool_calls"] = []
    state["pending_assistant_msg"] = None
    state["submitted_tool_results"] = []
    state["epoch"] = _next_epoch()
    return state["epoch"]


def get_context_messages(state: ChatState) -> list[dict]:
    return state["messages"][state["context_start_index"]:]


def build_messages_for_api(state: ChatState, system_prompt: str) -> list[dict]:
    context = get_context_messages(state)
    if system_prompt and system_prompt.strip():
        return [{"role": "system", "content": system_prompt.strip()}] + context
    return list(context)


def build_api_kwargs(
    state: ChatState,
    system_prompt: str,
    functions_json_str: Optional[str],
    think_level: Optional[str],
    temperature: Optional[float],
    max_tokens: Optional[int],
    top_p: Optional[float],
    preserved_thinking: Optional[bool] = None,
) -> dict:
    """Build the kwargs dict passed to ``client.chat.completions.create``.

    Each knob is omitted from the request entirely when "unset" so the
    server applies its own default — but the meaning of "unset" differs:

    * ``temperature`` is tristate: ``None`` means unset (omit the field),
      while any float — including ``0``, which selects greedy decoding —
      is sent literally. The UI exposes this via a "Use model default"
      checkbox sitting next to the slider; the headless ``api_chat``
      surface uses ``temperature=None`` as its default.
    * ``max_tokens`` and ``top_p`` collapse "unset" and ``0`` into a
      single sentinel: a value of ``0`` (or ``None``) is treated as unset
      and the field is omitted. They have no UI checkbox because explicit
      ``0`` for either knob is not a useful operating point.
    * ``preserved_thinking`` is a tristate boolean like ``temperature``:
      ``None`` means unset (omit the field), while ``True`` / ``False``
      are sent literally. It is a non-standard extension, so it rides in
      ``extra_body`` rather than as a top-level kwarg (the OpenAI SDK
      would reject an unknown top-level argument). The UI exposes it via
      a "Use model default" checkbox next to an on/off toggle.
    """
    api_messages = build_messages_for_api(state, system_prompt)
    tools = build_tools_list(functions_json_str)

    kwargs: dict = dict(
        model=MODEL,
        messages=api_messages,
        stream=True,
        reasoning_effort=think_level or "no_think",
    )
    if max_tokens is not None and int(max_tokens) != 0:
        kwargs["max_tokens"] = int(max_tokens)
    if temperature is not None:
        kwargs["temperature"] = float(temperature)
    if top_p is not None and float(top_p) != 0:
        kwargs["top_p"] = float(top_p)
    if preserved_thinking is not None:
        kwargs["extra_body"] = {"preserved_thinking": bool(preserved_thinking)}
    if tools:
        kwargs["tools"] = tools
    return kwargs


def _accumulate_tool_call(tool_calls_acc: list[dict], delta_tcs: list[Any]) -> None:
    """Merge streamed tool-call deltas into the accumulator."""
    for tc in delta_tcs:
        idx = getattr(tc, "index", 0) or 0
        while len(tool_calls_acc) <= idx:
            tool_calls_acc.append(
                {"id": "", "type": "function", "function": {"name": "", "arguments": ""}}
            )
        if tc.id:
            tool_calls_acc[idx]["id"] = tc.id
        if tc.function:
            if tc.function.name:
                tool_calls_acc[idx]["function"]["name"] += tc.function.name
            if tc.function.arguments:
                tool_calls_acc[idx]["function"]["arguments"] += tc.function.arguments


def _stream_worker(
    kwargs: dict,
    chunk_queue: queue.Queue,
) -> None:
    """Background thread: run the API call and feed chunks into *chunk_queue*."""
    try:
        stream = client.chat.completions.create(**kwargs)
        for chunk in stream:
            chunk_queue.put(chunk)
    except Exception as exc:
        chunk_queue.put(exc)
    finally:
        chunk_queue.put(_STREAM_DONE)


# Hard ceilings: if no chunk has arrived for this long AND the worker thread
# hasn't terminated, we abandon the stream so the UI lock can release. With a
# healthy heartbeat the worker normally posts STREAM_DONE within seconds of
# the model finishing, but reverse proxies / network blips can occasionally
# leave the SSE connection in a half-open state that hangs ``for chunk in
# stream`` indefinitely. Capping the wait guarantees ``send_message`` always
# reaches its final yield (and therefore re-enables the Send button).
#
# Two separate ceilings because the two phases have very different shapes:
#   * Before the first chunk the model may be doing reasoning / queueing /
#     KV-cache warmup, so we allow a generous 30s first-token budget.
#   * Once tokens are flowing we expect them to keep flowing; a 15s gap with
#     nothing arriving (and no STREAM_DONE) almost certainly means the SSE
#     socket is dead.
_FIRST_CHUNK_TIMEOUT = 60.0
_INTER_CHUNK_TIMEOUT = 15.0


# Op type constants — keep in sync with static/chat.js.
OP_REASONING_DELTA = "reasoning_delta"
OP_CONTENT_DELTA = "content_delta"
OP_TOOL_CALLS = "tool_calls"


def stream_response(
    kwargs: dict,
    is_cancelled: Optional[Callable[[], bool]] = None,
) -> Iterator[tuple[list[dict], str, str, list[dict], str]]:
    """Stream chunks from the API and yield delta-op batches.

    The actual HTTP stream runs in a daemon thread so that the generator can
    emit keep-alive yields during API-side pauses (model thinking, network
    hiccups, etc.). Without these heartbeats the SSE connection between the
    browser and a reverse proxy (e.g. HuggingFace Spaces) may be dropped
    for inactivity, silently killing the generator mid-response.

    Drain coalescing
    ----------------
    Each iteration drains EVERY chunk currently buffered into a single batch
    and emits one yield reflecting the merged deltas. Under back-pressure the
    yield rate naturally collapses (more chunks per yield) without losing
    data — the deltas accumulate in ``pending_*`` strings until the next
    successful yield can drain them.

    Two early-exit paths protect the UI from getting stuck:

    * As soon as we see ``finish_reason`` we drain whatever is already in
      the queue without blocking, then break. The model has logically
      finished; waiting on the SSE socket close would only lengthen the
      visible "stuck" window.
    * Two timeout safety nets force a break if the stream stalls while
      the worker is still technically alive.

    Yields ``(ops, assistant_total, reasoning_total, tool_calls, request_id)``
    where ``ops`` is the list of delta dicts since the previous yield.
    Heartbeat yields produce an empty ``ops`` list — callers should treat
    that as "no new content but the stream is still healthy".
    """
    assistant_content = ""
    reasoning_content = ""
    tool_calls_acc: list[dict] = []
    request_id = ""

    # Pending-since-last-yield deltas. Persist across drain iterations so
    # a throttle-suppressed yield doesn't lose the chars; the next yield
    # picks them up.
    pending_reasoning = ""
    pending_content = ""
    tool_calls_dirty = False

    chunk_q: queue.Queue = queue.Queue()
    worker = threading.Thread(
        target=_stream_worker, args=(kwargs, chunk_q), daemon=True,
    )
    worker.start()

    saw_finish_reason = False

    def take_ops() -> list[dict]:
        """Drain pending deltas into an ops list; return [] if nothing pending."""
        nonlocal pending_reasoning, pending_content, tool_calls_dirty
        ops: list[dict] = []
        if pending_reasoning:
            ops.append({"type": OP_REASONING_DELTA, "delta": pending_reasoning})
            pending_reasoning = ""
        if pending_content:
            ops.append({"type": OP_CONTENT_DELTA, "delta": pending_content})
            pending_content = ""
        if tool_calls_dirty:
            ops.append({"type": OP_TOOL_CALLS, "tool_calls": list(tool_calls_acc)})
            tool_calls_dirty = False
        return ops

    def apply_chunk(chunk) -> bool:
        """Fold a single API chunk into the accumulators.

        Returns True when this chunk produced visible-state changes
        (content, reasoning, or tool-call deltas). Sets the outer
        ``saw_finish_reason`` / ``request_id`` as a side effect.
        """
        nonlocal request_id, reasoning_content, assistant_content
        nonlocal pending_reasoning, pending_content, tool_calls_dirty
        nonlocal saw_finish_reason
        if not request_id and getattr(chunk, "id", None):
            request_id = chunk.id
        if not chunk.choices:
            return False
        choice = chunk.choices[0]
        delta = choice.delta
        if getattr(choice, "finish_reason", None):
            saw_finish_reason = True

        changed = False
        rc = getattr(delta, "reasoning_content", None)
        if rc:
            reasoning_content += rc
            pending_reasoning += rc
            changed = True
        if delta.content:
            assistant_content += delta.content
            pending_content += delta.content
            changed = True
        if getattr(delta, "tool_calls", None):
            _accumulate_tool_call(tool_calls_acc, delta.tool_calls)
            tool_calls_dirty = True
            changed = True
        return changed

    last_yield_at = 0.0
    last_chunk_at = time.monotonic()
    got_first_chunk = False
    yielded = False
    done = False

    while not done:
        # Cancellation check — caller (e.g. ``new_chat``) bumped the
        # session epoch, so abandon the stream WITHOUT a final yield.
        # The worker thread keeps running until the upstream API closes
        # the connection, but its chunks pile harmlessly into the
        # garbage-collected queue once we return.
        if is_cancelled is not None and is_cancelled():
            logger.debug("stream cancelled by caller, abandoning")
            return

        # ── block for the next item, with heartbeat / stall guards ──
        try:
            first = chunk_q.get(timeout=_HEARTBEAT_INTERVAL)
        except queue.Empty:
            if not worker.is_alive() and chunk_q.empty():
                break
            stall_budget = (
                _INTER_CHUNK_TIMEOUT if got_first_chunk else _FIRST_CHUNK_TIMEOUT
            )
            if time.monotonic() - last_chunk_at > stall_budget:
                logger.warning(
                    "stream stalled %.1fs with no chunks (%s), abandoning",
                    stall_budget,
                    "inter-chunk" if got_first_chunk else "first-chunk",
                )
                break
            # Heartbeat: re-emit current state with empty ops so Gradio
            # ships an SSE frame and the upstream proxy doesn't consider
            # the channel idle. The empty-ops frame is ~70 bytes and the
            # client treats it as a noop.
            yield [], assistant_content, reasoning_content, tool_calls_acc, request_id
            yielded = True
            last_yield_at = time.monotonic()
            continue

        # ── coalesce: pull every chunk currently buffered ──
        batch = [first] + _drain_queue(chunk_q)

        for item in batch:
            if item is _STREAM_DONE:
                done = True
                continue
            if isinstance(item, Exception):
                raise item
            last_chunk_at = time.monotonic()
            got_first_chunk = True
            apply_chunk(item)

        # ── one throttled yield per drained batch ──
        # Force-emit on done / finish so the final state always ships.
        if pending_reasoning or pending_content or tool_calls_dirty:
            now = time.monotonic()
            if done or saw_finish_reason or now - last_yield_at >= _YIELD_INTERVAL:
                ops = take_ops()
                yield ops, assistant_content, reasoning_content, tool_calls_acc, request_id
                yielded = True
                last_yield_at = now

        # ── finish_reason fast-exit ──
        # Model has logically finished. Drain anything still buffered and
        # exit. Don't wait on the SSE socket close.
        if saw_finish_reason and not done:
            for item in _drain_queue(chunk_q):
                if item is _STREAM_DONE:
                    break
                if isinstance(item, Exception):
                    raise item
                apply_chunk(item)
            ops = take_ops()
            if ops:
                yield ops, assistant_content, reasoning_content, tool_calls_acc, request_id
                yielded = True
            break

    # Final flush — guarantee callers always observe terminal accumulator
    # values, even when every prior content yield was suppressed by the
    # throttle (e.g. a tiny response that finished within the floor).
    ops = take_ops()
    if ops or not yielded:
        yield ops, assistant_content, reasoning_content, tool_calls_acc, request_id


def finalize_response(
    state: ChatState,
    assistant_content: str,
    reasoning_content: str,
    tool_calls_acc: list[dict],
) -> tuple[bool, list[dict]]:
    """Persist the final assistant message into ``state``.

    Returns ``(has_pending_tool_calls, pending_tool_calls)``. The Gradio
    adapter is responsible for turning ``pending_tool_calls`` into UI
    updates (see ``chat.py``).
    """
    assistant_msg: dict = {"role": "assistant", "content": assistant_content or None}
    if reasoning_content:
        assistant_msg["reasoning_content"] = reasoning_content

    if tool_calls_acc:
        assistant_msg["tool_calls"] = tool_calls_acc
        state["messages"].append(assistant_msg)
        state["pending_tool_calls"] = list(tool_calls_acc)
        state["submitted_tool_results"] = []
        state["pending_assistant_msg"] = assistant_msg
        logger.debug("queued %d tool call(s)", len(tool_calls_acc))
        return True, list(tool_calls_acc)

    state["messages"].append(assistant_msg)
    return False, []


def record_tool_result(state: ChatState, tool_call: Any, result_text: str) -> None:
    """Record a single tool-call result in the pending queue."""
    tc_id = tool_call["id"] if isinstance(tool_call, dict) else tool_call.id
    state.setdefault("submitted_tool_results", []).append({
        "role": "tool",
        "tool_call_id": tc_id,
        "content": result_text or "",
    })


def flush_tool_results(state: ChatState) -> None:
    """Move queued tool results into the main message log."""
    for msg in state.get("submitted_tool_results", []):
        state["messages"].append(msg)
    state["submitted_tool_results"] = []
    state["pending_assistant_msg"] = None