stream.py

import asyncio
import json
import logging
import time
import uuid
from typing import AsyncIterator

from config import DEFAULT_MODEL, STOP_REASON_MAP
from response_cache import CompletionArtifact, get_cache_service
from tools import parse_function_calls_text
from filters import P5jsLeadingFilter, ToolAwareTextBuffer
from upstream import iter_upstream_events, LiveArtifactCapture

logger = logging.getLogger(__name__)


async def anthropic_stream_plain(payload: dict, capture: LiveArtifactCapture | None = None) -> AsyncIterator[bytes]:
    p5_filter = P5jsLeadingFilter()
    async for event_name, obj in iter_upstream_events(payload):
        if capture is not None:
            capture.observe(event_name, obj)
        if event_name == "error":
            err = {"type": "error", "error": {"type": "upstream_error", "message": obj.get("body", "")}}
            yield f"event: error\ndata: {json.dumps(err)}\n\n".encode()
            return
        if event_name == "done":
            yield b"data: [DONE]\n\n"
            continue
        if event_name == "content_block_delta":
            delta = obj.get("delta", {})
            if delta.get("type") == "text_delta":
                filtered = p5_filter.feed(delta.get("text", ""))
                if not filtered:
                    continue
                obj = {**obj, "delta": {**delta, "text": filtered}}
        elif event_name == "message_stop":
            tail = p5_filter.flush()
            if tail:
                tail_obj = {
                    "type": "content_block_delta",
                    "index": 0,
                    "delta": {"type": "text_delta", "text": tail},
                }
                yield f"event: content_block_delta\ndata: {json.dumps(tail_obj, ensure_ascii=False)}\n\n".encode()
        yield f"event: {event_name}\ndata: {json.dumps(obj, ensure_ascii=False)}\n\n".encode()


async def openai_stream_plain(payload: dict, requested_model: str, capture: LiveArtifactCapture | None = None) -> AsyncIterator[bytes]:
    chat_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created = int(time.time())
    model_id = requested_model
    first_chunk_sent = False
    finish_reason: str | None = None
    p5_filter = P5jsLeadingFilter()

    def chunk(delta: dict, finish: str | None = None) -> bytes:
        payload_obj = {
            "id": chat_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model_id,
            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
        }
        return f"data: {json.dumps(payload_obj, ensure_ascii=False)}\n\n".encode()

    async for event_name, obj in iter_upstream_events(payload):
        if capture is not None:
            capture.observe(event_name, obj)
        if event_name == "error":
            err = {"error": {"message": obj.get("body", "upstream error"), "type": "upstream_error"}}
            yield f"data: {json.dumps(err)}\n\n".encode()
            yield b"data: [DONE]\n\n"
            return
        if event_name == "message_start":
            m = obj.get("message", {})
            model_id = m.get("model", model_id)
            if not first_chunk_sent:
                yield chunk({"role": "assistant", "content": ""})
                first_chunk_sent = True
        elif event_name == "content_block_delta":
            delta = obj.get("delta", {})
            if delta.get("type") == "text_delta":
                text = p5_filter.feed(delta.get("text", ""))
                if text:
                    if not first_chunk_sent:
                        yield chunk({"role": "assistant", "content": ""})
                        first_chunk_sent = True
                    yield chunk({"content": text})
        elif event_name == "message_delta":
            d = obj.get("delta", {})
            if d.get("stop_reason"):
                finish_reason = STOP_REASON_MAP.get(d["stop_reason"], "stop")
        elif event_name == "message_stop":
            tail = p5_filter.flush()
            if tail:
                yield chunk({"content": tail})
            yield chunk({}, finish=finish_reason or "stop")
        elif event_name == "done":
            yield b"data: [DONE]\n\n"


async def anthropic_stream_with_tools(payload: dict, capture: LiveArtifactCapture | None = None) -> AsyncIterator[bytes]:
    msg_id = f"msg_{uuid.uuid4().hex[:24]}"
    model_id = payload.get("model", DEFAULT_MODEL)
    next_index = 0
    text_index: int | None = None
    text_opened = False
    buf = ToolAwareTextBuffer()
    p5_filter = P5jsLeadingFilter()
    stop_reason = "end_turn"
    saw_tool_use = False
    usage_seed = {"input_tokens": 0, "output_tokens": 0}

    def sse(event: str, obj: dict) -> bytes:
        return f"event: {event}\ndata: {json.dumps(obj, ensure_ascii=False)}\n\n".encode()

    def emit_text(t: str) -> bytes | None:
        nonlocal text_opened, text_index, next_index
        if not t:
            return None
        parts = []
        if not text_opened:
            text_index = next_index
            next_index += 1
            parts.append(sse("content_block_start", {
                "type": "content_block_start",
                "index": text_index,
                "content_block": {"type": "text", "text": ""},
            }))
            text_opened = True
        parts.append(sse("content_block_delta", {
            "type": "content_block_delta",
            "index": text_index,
            "delta": {"type": "text_delta", "text": t},
        }))
        return b"".join(parts)

    def close_text_if_open() -> bytes | None:
        nonlocal text_opened
        if text_opened and text_index is not None:
            text_opened = False
            return sse("content_block_stop", {"type": "content_block_stop", "index": text_index})
        return None

    def emit_tool_block(block: str) -> bytes | None:
        nonlocal next_index, saw_tool_use
        tool_uses = parse_function_calls_text(block)
        if not tool_uses:
            return None
        chunks: list[bytes] = []
        closed = close_text_if_open()
        if closed:
            chunks.append(closed)
        for tu in tool_uses:
            saw_tool_use = True
            idx = next_index
            next_index += 1
            chunks.append(sse("content_block_start", {
                "type": "content_block_start",
                "index": idx,
                "content_block": {"type": "tool_use", "id": tu["id"], "name": tu["name"], "input": {}},
            }))
            chunks.append(sse("content_block_delta", {
                "type": "content_block_delta",
                "index": idx,
                "delta": {"type": "input_json_delta", "partial_json": json.dumps(tu["input"], ensure_ascii=False)},
            }))
            chunks.append(sse("content_block_stop", {"type": "content_block_stop", "index": idx}))
        return b"".join(chunks)

    started = False
    async for event_name, obj in iter_upstream_events(payload):
        if capture is not None:
            capture.observe(event_name, obj)
        if event_name == "error":
            err = {"type": "error", "error": {"type": "upstream_error", "message": obj.get("body", "")}}
            yield sse("error", err)
            return
        if event_name == "message_start":
            m = obj.get("message", {})
            msg_id = m.get("id", msg_id)
            model_id = m.get("model", model_id)
            if "usage" in m:
                usage_seed["input_tokens"] = m["usage"].get("input_tokens", 0)
            if not started:
                started = True
                yield sse("message_start", {
                    "type": "message_start",
                    "message": {
                        "id": msg_id,
                        "type": "message",
                        "role": "assistant",
                        "model": model_id,
                        "content": [],
                        "stop_reason": None,
                        "stop_sequence": None,
                        "usage": usage_seed,
                    },
                })
        elif event_name == "content_block_delta":
            delta = obj.get("delta", {})
            if delta.get("type") == "text_delta":
                t = p5_filter.feed(delta.get("text", ""))
                if not t:
                    continue
                for kind, payload_text in buf.feed(t):
                    if kind == "text":
                        out = emit_text(payload_text)
                        if out:
                            yield out
                    elif kind == "tool_block":
                        out = emit_tool_block(payload_text)
                        if out:
                            yield out
        elif event_name == "message_delta":
            d = obj.get("delta", {})
            if d.get("stop_reason"):
                stop_reason = d["stop_reason"]
            u = obj.get("usage")
            if u and "output_tokens" in u:
                usage_seed["output_tokens"] = u["output_tokens"]
        elif event_name == "message_stop":
            tail = p5_filter.flush()
            if tail:
                for kind, payload_text in buf.feed(tail):
                    if kind == "text":
                        out = emit_text(payload_text)
                        if out:
                            yield out
                    elif kind == "tool_block":
                        out = emit_tool_block(payload_text)
                        if out:
                            yield out
            for kind, payload_text in buf.flush():
                if kind == "text":
                    out = emit_text(payload_text)
                    if out:
                        yield out
                elif kind == "tool_block":
                    out = emit_tool_block(payload_text)
                    if out:
                        yield out
            closed = close_text_if_open()
            if closed:
                yield closed
            if saw_tool_use:
                stop_reason = "tool_use"
            yield sse("message_delta", {
                "type": "message_delta",
                "delta": {"stop_reason": stop_reason, "stop_sequence": None},
                "usage": {"input_tokens": usage_seed["input_tokens"], "output_tokens": usage_seed["output_tokens"]},
            })
            yield sse("message_stop", {"type": "message_stop"})
        elif event_name == "done":
            yield b"data: [DONE]\n\n"


async def openai_stream_with_tools(payload: dict, requested_model: str, capture: LiveArtifactCapture | None = None) -> AsyncIterator[bytes]:
    chat_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created = int(time.time())
    model_id = requested_model
    buf = ToolAwareTextBuffer()
    p5_filter = P5jsLeadingFilter()
    first_chunk_sent = False
    finish_reason: str | None = None
    saw_tool_use = False
    next_tool_index = 0

    def chunk(delta: dict, finish: str | None = None) -> bytes:
        obj = {
            "id": chat_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model_id,
            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
        }
        return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n".encode()

    def emit_text(t: str) -> bytes | None:
        nonlocal first_chunk_sent
        if not t:
            return None
        if not first_chunk_sent:
            first_chunk_sent = True
            return chunk({"role": "assistant", "content": ""}) + chunk({"content": t})
        return chunk({"content": t})

    def emit_tool_block(block: str) -> bytes | None:
        nonlocal next_tool_index, saw_tool_use, first_chunk_sent
        tool_uses = parse_function_calls_text(block)
        if not tool_uses:
            return None
        chunks: list[bytes] = []
        if not first_chunk_sent:
            chunks.append(chunk({"role": "assistant", "content": None}))
            first_chunk_sent = True
        for tu in tool_uses:
            saw_tool_use = True
            idx = next_tool_index
            next_tool_index += 1
            chunks.append(chunk({"tool_calls": [{
                "index": idx,
                "id": f"call_{tu['id'].removeprefix('toolu_')}",
                "type": "function",
                "function": {"name": tu["name"], "arguments": ""},
            }]}))
            chunks.append(chunk({"tool_calls": [{
                "index": idx,
                "function": {"arguments": json.dumps(tu["input"], ensure_ascii=False)},
            }]}))
        return b"".join(chunks)

    async for event_name, obj in iter_upstream_events(payload):
        if capture is not None:
            capture.observe(event_name, obj)
        if event_name == "error":
            err = {"error": {"message": obj.get("body", "upstream error"), "type": "upstream_error"}}
            yield f"data: {json.dumps(err)}\n\n".encode()
            yield b"data: [DONE]\n\n"
            return
        if event_name == "message_start":
            m = obj.get("message", {})
            model_id = m.get("model", model_id)
        elif event_name == "content_block_delta":
            delta = obj.get("delta", {})
            if delta.get("type") == "text_delta":
                t = p5_filter.feed(delta.get("text", ""))
                if not t:
                    continue
                for kind, payload_text in buf.feed(t):
                    if kind == "text":
                        out = emit_text(payload_text)
                        if out:
                            yield out
                    elif kind == "tool_block":
                        out = emit_tool_block(payload_text)
                        if out:
                            yield out
        elif event_name == "message_delta":
            d = obj.get("delta", {})
            if d.get("stop_reason"):
                finish_reason = STOP_REASON_MAP.get(d["stop_reason"], "stop")
        elif event_name == "message_stop":
            tail = p5_filter.flush()
            if tail:
                for kind, payload_text in buf.feed(tail):
                    if kind == "text":
                        out = emit_text(payload_text)
                        if out:
                            yield out
                    elif kind == "tool_block":
                        out = emit_tool_block(payload_text)
                        if out:
                            yield out
            for kind, payload_text in buf.flush():
                if kind == "text":
                    out = emit_text(payload_text)
                    if out:
                        yield out
                elif kind == "tool_block":
                    out = emit_tool_block(payload_text)
                    if out:
                        yield out
            if saw_tool_use:
                finish_reason = "tool_calls"
            yield chunk({}, finish=finish_reason or "stop")
        elif event_name == "done":
            yield b"data: [DONE]\n\n"


def build_cache_headers(status: str, source: str | None = None) -> dict[str, str]:
    headers = {"X-Proxy-Cache": status}
    if source:
        headers["X-Proxy-Cache-Source"] = source
    return headers


def build_stream_headers(status: str, source: str | None = None) -> dict[str, str]:
    headers = {
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no",
    }
    headers.update(build_cache_headers(status, source))
    return headers


async def wait_for_inflight_artifact(future: asyncio.Future[CompletionArtifact]) -> CompletionArtifact | None:
    try:
        return await future
    except Exception:
        return None


async def finalize_stream_cache(
    stream: AsyncIterator[bytes],
    capture: LiveArtifactCapture,
    cache_key: str,
    ttl_secs: int,
) -> AsyncIterator[bytes]:
    cache_service = get_cache_service()
    try:
        async for chunk in stream:
            yield chunk
        if capture.is_cacheable():
            artifact = capture.build()
            await cache_service.set(cache_key, artifact, ttl_secs)
            await cache_service.inflight.resolve(cache_key, artifact)
        else:
            await cache_service.inflight.reject(cache_key, RuntimeError("stream did not produce a cacheable artifact"))
    except asyncio.CancelledError as exc:
        await cache_service.inflight.reject(cache_key, exc)
        raise
    except Exception as exc:
        await cache_service.inflight.reject(cache_key, exc)
        raise