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)