render.py

import json
import time
import uuid
from typing import Any, AsyncIterator

from config import DEFAULT_MODEL, STOP_REASON_MAP, STREAM_RENDER_CHUNK_SIZE
from response_cache import CompletionArtifact
from tools import parse_function_calls_text, split_text_and_tools
from filters import P5jsLeadingFilter, ToolAwareTextBuffer, strip_p5js_noise
from upstream import fetch_completion_artifact


def extract_artifact_parts(artifact: CompletionArtifact) -> tuple[str, list[dict]]:
    clean_text, tool_uses = split_text_and_tools(artifact.raw_text)
    return strip_p5js_noise(clean_text), tool_uses


def render_anthropic_json_from_artifact(artifact: CompletionArtifact, has_tools: bool) -> dict:
    msg_id = f"msg_{uuid.uuid4().hex[:24]}"
    clean_text, tool_uses = extract_artifact_parts(artifact)

    content: list[dict] = []
    if clean_text:
        content.append({"type": "text", "text": clean_text})
    if has_tools:
        for tu in tool_uses:
            content.append({"type": "tool_use", "id": tu["id"], "name": tu["name"], "input": tu["input"]})

    stop_reason = "tool_use" if has_tools and tool_uses else artifact.stop_reason

    return {
        "id": msg_id,
        "type": "message",
        "role": "assistant",
        "model": artifact.model_id,
        "content": content or [{"type": "text", "text": ""}],
        "stop_reason": stop_reason,
        "stop_sequence": None,
        "usage": {
            "input_tokens": artifact.usage_input_tokens,
            "output_tokens": artifact.usage_output_tokens,
        },
    }


def render_openai_json_from_artifact(artifact: CompletionArtifact, has_tools: bool) -> dict:
    clean_text, tool_uses = extract_artifact_parts(artifact)
    finish_reason = "tool_calls" if has_tools and tool_uses else STOP_REASON_MAP.get(artifact.stop_reason, "stop")

    message: dict[str, Any] = {
        "role": "assistant",
        "content": clean_text or (None if has_tools and tool_uses else ""),
    }
    if has_tools and tool_uses:
        message["tool_calls"] = [
            {
                "id": f"call_{tu['id'].removeprefix('toolu_')}",
                "type": "function",
                "function": {
                    "name": tu["name"],
                    "arguments": json.dumps(tu["input"], ensure_ascii=False),
                },
            }
            for tu in tool_uses
        ]

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": artifact.model_id,
        "choices": [
            {
                "index": 0,
                "message": message,
                "finish_reason": finish_reason,
            }
        ],
        "usage": {
            "prompt_tokens": artifact.usage_input_tokens,
            "completion_tokens": artifact.usage_output_tokens,
            "total_tokens": artifact.usage_input_tokens + artifact.usage_output_tokens,
        },
    }


def _iter_text_chunks(text: str, chunk_size: int = STREAM_RENDER_CHUNK_SIZE):
    if not text:
        return
    for index in range(0, len(text), chunk_size):
        yield text[index:index + chunk_size]


def iter_stream_segments(raw_text: str, has_tools: bool):
    p5_filter = P5jsLeadingFilter()
    tool_buffer = ToolAwareTextBuffer()

    for raw_chunk in _iter_text_chunks(raw_text):
        filtered = p5_filter.feed(raw_chunk)
        if not filtered:
            continue
        if has_tools:
            yield from tool_buffer.feed(filtered)
        else:
            yield ("text", filtered)

    tail = p5_filter.flush()
    if tail:
        if has_tools:
            yield from tool_buffer.feed(tail)
        else:
            yield ("text", tail)

    if has_tools:
        yield from tool_buffer.flush()


async def anthropic_stream_from_artifact(artifact: CompletionArtifact, has_tools: bool) -> AsyncIterator[bytes]:
    msg_id = f"msg_{uuid.uuid4().hex[:24]}"
    usage = {
        "input_tokens": artifact.usage_input_tokens,
        "output_tokens": artifact.usage_output_tokens,
    }
    next_index = 0
    text_index: int | None = None
    text_opened = False
    saw_tool_use = False

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

    yield sse("message_start", {
        "type": "message_start",
        "message": {
            "id": msg_id,
            "type": "message",
            "role": "assistant",
            "model": artifact.model_id,
            "content": [],
            "stop_reason": None,
            "stop_sequence": None,
            "usage": usage,
        },
    })

    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

    for kind, payload_text in iter_stream_segments(artifact.raw_text, has_tools):
        if kind == "text":
            if not payload_text:
                continue
            if not text_opened:
                text_index = next_index
                next_index += 1
                yield sse("content_block_start", {
                    "type": "content_block_start",
                    "index": text_index,
                    "content_block": {"type": "text", "text": ""},
                })
                text_opened = True
            for chunk in _iter_text_chunks(payload_text):
                yield sse("content_block_delta", {
                    "type": "content_block_delta",
                    "index": text_index,
                    "delta": {"type": "text_delta", "text": chunk},
                })
        elif kind == "tool_block":
            tool_uses = parse_function_calls_text(payload_text)
            if not tool_uses:
                continue
            closed = close_text_if_open()
            if closed:
                yield closed
            for tu in tool_uses:
                saw_tool_use = True
                index = next_index
                next_index += 1
                yield sse("content_block_start", {
                    "type": "content_block_start",
                    "index": index,
                    "content_block": {"type": "tool_use", "id": tu["id"], "name": tu["name"], "input": {}},
                })
                yield sse("content_block_delta", {
                    "type": "content_block_delta",
                    "index": index,
                    "delta": {"type": "input_json_delta", "partial_json": json.dumps(tu["input"], ensure_ascii=False)},
                })
                yield sse("content_block_stop", {"type": "content_block_stop", "index": index})

    closed = close_text_if_open()
    if closed:
        yield closed

    stop_reason = "tool_use" if saw_tool_use else artifact.stop_reason
    yield sse("message_delta", {
        "type": "message_delta",
        "delta": {"stop_reason": stop_reason, "stop_sequence": None},
        "usage": usage,
    })
    yield sse("message_stop", {"type": "message_stop"})
    yield b"data: [DONE]\n\n"


async def openai_stream_from_artifact(artifact: CompletionArtifact, has_tools: bool) -> AsyncIterator[bytes]:
    chat_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created = int(time.time())
    first_chunk_sent = False
    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": artifact.model_id,
            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
        }
        return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n".encode()

    for kind, payload_text in iter_stream_segments(artifact.raw_text, has_tools):
        if kind == "text":
            if not payload_text:
                continue
            for text_chunk in _iter_text_chunks(payload_text):
                if not first_chunk_sent:
                    yield chunk({"role": "assistant", "content": ""})
                    first_chunk_sent = True
                yield chunk({"content": text_chunk})
        elif kind == "tool_block":
            tool_uses = parse_function_calls_text(payload_text)
            if not tool_uses:
                continue
            if not first_chunk_sent:
                yield chunk({"role": "assistant", "content": None})
                first_chunk_sent = True
            for tu in tool_uses:
                saw_tool_use = True
                index = next_tool_index
                next_tool_index += 1
                yield chunk({"tool_calls": [{
                    "index": index,
                    "id": f"call_{tu['id'].removeprefix('toolu_')}",
                    "type": "function",
                    "function": {"name": tu["name"], "arguments": ""},
                }]})
                yield chunk({"tool_calls": [{
                    "index": index,
                    "function": {"arguments": json.dumps(tu["input"], ensure_ascii=False)},
                }]})

    finish_reason = "tool_calls" if saw_tool_use else STOP_REASON_MAP.get(artifact.stop_reason, "stop")
    if not first_chunk_sent:
        yield chunk({"role": "assistant", "content": ""})
    yield chunk({}, finish=finish_reason)
    yield b"data: [DONE]\n\n"


async def anthropic_aggregate(payload: dict, has_tools: bool) -> dict:
    artifact = await fetch_completion_artifact(payload)
    return render_anthropic_json_from_artifact(artifact, has_tools)


async def openai_aggregate(payload: dict, requested_model: str, has_tools: bool) -> dict:
    artifact = await fetch_completion_artifact(payload)
    return render_openai_json_from_artifact(artifact, has_tools)