Spaces:
Paused
Paused
| 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 | |