from __future__ import annotations import json import uuid from typing import Any, AsyncGenerator, Dict import httpx from .logging import logger from .config import BRIDGE_BASE_URL from .helpers import _get async def stream_openai_sse(packet: Dict[str, Any], completion_id: str, created_ts: int, model_id: str) -> AsyncGenerator[str, None]: try: first = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {"role": "assistant"}}], } # 打印转换后的首个 SSE 事件（OpenAI 格式） try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): %s", json.dumps(first, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(first, ensure_ascii=False)}\n\n" timeout = httpx.Timeout(60.0) async with httpx.AsyncClient(http2=True, timeout=timeout, trust_env=True) as client: def _do_stream(): return client.stream( "POST", f"{BRIDGE_BASE_URL}/api/warp/send_stream_sse", headers={"accept": "text/event-stream"}, json={"json_data": packet, "message_type": "warp.multi_agent.v1.Request"}, ) # 首次请求 response_cm = _do_stream() async with response_cm as response: if response.status_code == 429: try: r = await client.post(f"{BRIDGE_BASE_URL}/api/auth/refresh", timeout=10.0) logger.warning("[OpenAI Compat] Bridge returned 429. Tried JWT refresh -> HTTP %s", r.status_code) except Exception as _e: logger.warning("[OpenAI Compat] JWT refresh attempt failed after 429: %s", _e) # 重试一次 response_cm2 = _do_stream() async with response_cm2 as response2: response = response2 if response.status_code != 200: error_text = await response.aread() error_content = error_text.decode("utf-8") if error_text else "" logger.error(f"[OpenAI Compat] Bridge HTTP error {response.status_code}: {error_content[:300]}") raise RuntimeError(f"bridge error: {error_content}") current = "" tool_calls_emitted = False async for line in response.aiter_lines(): if line.startswith("data:"): payload = line[5:].strip() if not payload: continue # 打印接收到的 Protobuf SSE 原始事件片段 try: logger.info("[OpenAI Compat] 接收到的 Protobuf SSE(data): %s", payload) except Exception: pass if payload == "[DONE]": break current += payload continue if (line.strip() == "") and current: try: ev = json.loads(current) except Exception: current = "" continue current = "" event_data = (ev or {}).get("parsed_data") or {} # 打印接收到的 Protobuf 事件（解析后） try: logger.info("[OpenAI Compat] 接收到的 Protobuf 事件(parsed): %s", json.dumps(event_data, ensure_ascii=False)) except Exception: pass if "init" in event_data: pass client_actions = _get(event_data, "client_actions", "clientActions") if isinstance(client_actions, dict): actions = _get(client_actions, "actions", "Actions") or [] for action in actions: append_data = _get(action, "append_to_message_content", "appendToMessageContent") if isinstance(append_data, dict): message = append_data.get("message", {}) agent_output = _get(message, "agent_output", "agentOutput") or {} text_content = agent_output.get("text", "") if text_content: delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {"content": text_content}}], } # 打印转换后的 OpenAI SSE 事件 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" messages_data = _get(action, "add_messages_to_task", "addMessagesToTask") if isinstance(messages_data, dict): messages = messages_data.get("messages", []) for message in messages: tool_call = _get(message, "tool_call", "toolCall") or {} call_mcp = _get(tool_call, "call_mcp_tool", "callMcpTool") or {} if isinstance(call_mcp, dict) and call_mcp.get("name"): try: args_obj = call_mcp.get("args", {}) or {} args_str = json.dumps(args_obj, ensure_ascii=False) except Exception: args_str = "{}" tool_call_id = tool_call.get("tool_call_id") or str(uuid.uuid4()) delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{ "index": 0, "delta": { "tool_calls": [{ "index": 0, "id": tool_call_id, "type": "function", "function": {"name": call_mcp.get("name"), "arguments": args_str}, }] } }], } # 打印转换后的 OpenAI 工具调用事件 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit tool_calls): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" tool_calls_emitted = True else: agent_output = _get(message, "agent_output", "agentOutput") or {} text_content = agent_output.get("text", "") if text_content: delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {"content": text_content}}], } try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" if "finished" in event_data: done_chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {}, "finish_reason": ("tool_calls" if tool_calls_emitted else "stop")}], } try: logger.info("[OpenAI Compat] 转换后的 SSE(emit done): %s", json.dumps(done_chunk, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(done_chunk, ensure_ascii=False)}\n\n" # 打印完成标记 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): [DONE]") except Exception: pass yield "data: [DONE]\n\n" return if response.status_code != 200: error_text = await response.aread() error_content = error_text.decode("utf-8") if error_text else "" logger.error(f"[OpenAI Compat] Bridge HTTP error {response.status_code}: {error_content[:300]}") raise RuntimeError(f"bridge error: {error_content}") current = "" tool_calls_emitted = False async for line in response.aiter_lines(): if line.startswith("data:"): payload = line[5:].strip() if not payload: continue # 打印接收到的 Protobuf SSE 原始事件片段 try: logger.info("[OpenAI Compat] 接收到的 Protobuf SSE(data): %s", payload) except Exception: pass if payload == "[DONE]": break current += payload continue if (line.strip() == "") and current: try: ev = json.loads(current) except Exception: current = "" continue current = "" event_data = (ev or {}).get("parsed_data") or {} # 打印接收到的 Protobuf 事件（解析后） try: logger.info("[OpenAI Compat] 接收到的 Protobuf 事件(parsed): %s", json.dumps(event_data, ensure_ascii=False)) except Exception: pass if "init" in event_data: pass client_actions = _get(event_data, "client_actions", "clientActions") if isinstance(client_actions, dict): actions = _get(client_actions, "actions", "Actions") or [] for action in actions: append_data = _get(action, "append_to_message_content", "appendToMessageContent") if isinstance(append_data, dict): message = append_data.get("message", {}) agent_output = _get(message, "agent_output", "agentOutput") or {} text_content = agent_output.get("text", "") if text_content: delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {"content": text_content}}], } # 打印转换后的 OpenAI SSE 事件 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" messages_data = _get(action, "add_messages_to_task", "addMessagesToTask") if isinstance(messages_data, dict): messages = messages_data.get("messages", []) for message in messages: tool_call = _get(message, "tool_call", "toolCall") or {} call_mcp = _get(tool_call, "call_mcp_tool", "callMcpTool") or {} if isinstance(call_mcp, dict) and call_mcp.get("name"): try: args_obj = call_mcp.get("args", {}) or {} args_str = json.dumps(args_obj, ensure_ascii=False) except Exception: args_str = "{}" tool_call_id = tool_call.get("tool_call_id") or str(uuid.uuid4()) delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{ "index": 0, "delta": { "tool_calls": [{ "index": 0, "id": tool_call_id, "type": "function", "function": {"name": call_mcp.get("name"), "arguments": args_str}, }] } }], } # 打印转换后的 OpenAI 工具调用事件 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit tool_calls): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" tool_calls_emitted = True else: agent_output = _get(message, "agent_output", "agentOutput") or {} text_content = agent_output.get("text", "") if text_content: delta = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {"content": text_content}}], } try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): %s", json.dumps(delta, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(delta, ensure_ascii=False)}\n\n" if "finished" in event_data: done_chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {}, "finish_reason": ("tool_calls" if tool_calls_emitted else "stop")}], } try: logger.info("[OpenAI Compat] 转换后的 SSE(emit done): %s", json.dumps(done_chunk, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(done_chunk, ensure_ascii=False)}\n\n" # 打印完成标记 try: logger.info("[OpenAI Compat] 转换后的 SSE(emit): [DONE]") except Exception: pass yield "data: [DONE]\n\n" except Exception as e: logger.error(f"[OpenAI Compat] Stream processing failed: {e}") error_chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created_ts, "model": model_id, "choices": [{"index": 0, "delta": {}, "finish_reason": "error"}], "error": {"message": str(e)}, } try: logger.info("[OpenAI Compat] 转换后的 SSE(emit error): %s", json.dumps(error_chunk, ensure_ascii=False)) except Exception: pass yield f"data: {json.dumps(error_chunk, ensure_ascii=False)}\n\n" yield "data: [DONE]\n\n"