import json

import uuid

import os

import asyncio

import weakref

from pathlib import Path

from typing import Dict, Optional, Tuple, Iterator, List, AsyncGenerator, Any

import struct

import httpx

import importlib.util

def _load_claude_parser():

"""Dynamically load claude_parser module."""

base_dir = Path(__file__).resolve().parent

spec = importlib.util.spec_from_file_location("v2_claude_parser", str(base_dir / "claude_parser.py"))

module = importlib.util.module_from_spec(spec)

spec.loader.exec_module(module)

return module

try:

_parser = _load_claude_parser()

EventStreamParser = _parser.EventStreamParser

extract_event_info = _parser.extract_event_info

except Exception as e:

print(f"Warning: Failed to load claude_parser: {e}")

EventStreamParser = None

extract_event_info = None

class StreamTracker:

def __init__(self):

self.has_content = False

async def track(self, gen: AsyncGenerator[str, None]) -> AsyncGenerator[str, None]:

async for item in gen:

if item:

self.has_content = True

yield item

def _get_proxies() -> Optional[Dict[str, str]]:

proxy = os.getenv("HTTP_PROXY", "").strip()

if proxy:

return {"http": proxy, "https": proxy}

return None

BASE_DIR = Path(__file__).resolve().parent

TEMPLATE_PATH = BASE_DIR / "templates" / "streaming_request.json"

def load_template() -> Tuple[str, Dict[str, str], Dict[str, Any]]:

data = json.loads(TEMPLATE_PATH.read_text(encoding="utf-8"))

url, headers, body = data

assert isinstance(url, str) and isinstance(headers, dict) and isinstance(body, dict)

return url, headers, body

def _merge_headers(as_log: Dict[str, str], bearer_token: str) -> Dict[str, str]:

headers = dict(as_log)

for k in list(headers.keys()):

kl = k.lower()

if kl in ("content-length","host","connection","transfer-encoding"):

headers.pop(k, None)

def set_header(name: str, value: str):

for key in list(headers.keys()):

if key.lower() == name.lower():

del headers[key]

headers[name] = value

set_header("Authorization", f"Bearer {bearer_token}")

set_header("amz-sdk-invocation-id", str(uuid.uuid4()))

return headers

def _parse_event_headers(raw: bytes) -> Dict[str, object]:

headers: Dict[str, object] = {}

i = 0

n = len(raw)

while i < n:

if i + 1 > n:

break

name_len = raw[i]

i += 1

if i + name_len + 1 > n:

break

name = raw[i : i + name_len].decode("utf-8", errors="ignore")

i += name_len

htype = raw[i]

i += 1

if htype == 0:

val = True

elif htype == 1:

val = False

elif htype == 2:

if i + 1 > n: break

val = raw[i]; i += 1

elif htype == 3:

if i + 2 > n: break

val = int.from_bytes(raw[i:i+2],"big",signed=True); i += 2

elif htype == 4:

if i + 4 > n: break

val = int.from_bytes(raw[i:i+4],"big",signed=True); i += 4

elif htype == 5:

if i + 8 > n: break

val = int.from_bytes(raw[i:i+8],"big",signed=True); i += 8

elif htype == 6:

if i + 2 > n: break

l = int.from_bytes(raw[i:i+2],"big"); i += 2

if i + l > n: break

val = raw[i:i+l]; i += l

elif htype == 7:

if i + 2 > n: break

l = int.from_bytes(raw[i:i+2],"big"); i += 2

if i + l > n: break

val = raw[i:i+l].decode("utf-8", errors="ignore"); i += l

elif htype == 8:

if i + 8 > n: break

val = int.from_bytes(raw[i:i+8],"big",signed=False); i += 8

elif htype == 9:

if i + 16 > n: break

import uuid as _uuid

val = str(_uuid.UUID(bytes=bytes(raw[i:i+16]))); i += 16

else:

break

headers[name] = val

return headers

class AwsEventStreamParser:

def __init__(self):

self._buf = bytearray()

def feed(self, data: bytes) -> List[Tuple[Dict[str, object], bytes]]:

if not data:

return []

self._buf.extend(data)

out: List[Tuple[Dict[str, object], bytes]] = []

while True:

if len(self._buf) < 12:

break

total_len, headers_len, _prelude_crc = struct.unpack(">I I I", self._buf[:12])

if total_len < 16 or headers_len > total_len:

self._buf.pop(0)

continue

if len(self._buf) < total_len:

break

msg = bytes(self._buf[:total_len])

del self._buf[:total_len]

headers_raw = msg[12:12+headers_len]

payload = msg[12+headers_len: total_len-4]

headers = _parse_event_headers(headers_raw)

out.append((headers, payload))

return out

def _try_decode_event_payload(payload: bytes) -> Optional[dict]:

try:

return json.loads(payload.decode("utf-8"))

except Exception:

return None

def _extract_text_from_event(ev: dict) -> Optional[str]:

for key in ("assistantResponseEvent","assistantMessage","message","delta","data"):

if key in ev and isinstance(ev[key], dict):

inner = ev[key]

if isinstance(inner.get("content"), str) and inner.get("content"):

return inner["content"]

if isinstance(ev.get("content"), str) and ev.get("content"):

return ev["content"]

for list_key in ("chunks","content"):

if isinstance(ev.get(list_key), list):

buf = []

for item in ev[list_key]:

if isinstance(item, dict):

if isinstance(item.get("content"), str):

buf.append(item["content"])

elif isinstance(item.get("text"), str):

buf.append(item["text"])

elif isinstance(item, str):

buf.append(item)

if buf:

return "".join(buf)

for k in ("text","delta","payload"):

v = ev.get(k)

if isinstance(v, str) and v:

return v

return None

def openai_messages_to_text(messages: List[Dict[str, Any]]) -> str:

lines: List[str] = []

for m in messages:

role = m.get("role","user")

content = m.get("content","")

if isinstance(content, list):

parts = []

for seg in content:

if isinstance(seg, dict) and isinstance(seg.get("text"), str):

parts.append(seg["text"])

elif isinstance(seg, str):

parts.append(seg)

content = "\n".join(parts)

elif not isinstance(content, str):

content = str(content)

lines.append(f"{role}:\n{content}")

return "\n\n".join(lines)

def inject_history(body_json: Dict[str, Any], history_text: str) -> None:

try:

cur = body_json["conversationState"]["currentMessage"]["userInputMessage"]

content = cur.get("content","")

if isinstance(content, str):

cur["content"] = content.replace("你好，你必须讲个故事", history_text)

except Exception:

pass

def inject_model(body_json: Dict[str, Any], model: Optional[str]) -> None:

if not model:

return

try:

body_json["conversationState"]["currentMessage"]["userInputMessage"]["modelId"] = model

except Exception:

pass

async def send_chat_request(

access_token: str,

messages: List[Dict[str, Any]],

model: Optional[str] = None,

stream: bool = False,

timeout: Tuple[int,int] = (15,300),

client: Optional[httpx.AsyncClient] = None,

raw_payload: Optional[Dict[str, Any]] = None

) -> Tuple[Optional[str], Optional[AsyncGenerator[str, None]], StreamTracker, Optional[AsyncGenerator[Any, None]]]:

url, headers_from_log, body_json = load_template()

headers_from_log["amz-sdk-invocation-id"] = str(uuid.uuid4())

if raw_payload:

# Use raw payload if provided (for Claude API)

body_json = raw_payload

# Ensure conversationId is set if missing

if "conversationState" in body_json and "conversationId" not in body_json["conversationState"]:

body_json["conversationState"]["conversationId"] = str(uuid.uuid4())

else:

# Standard OpenAI-compatible logic

try:

body_json["conversationState"]["conversationId"] = str(uuid.uuid4())

except Exception:

pass

history_text = openai_messages_to_text(messages)

inject_history(body_json, history_text)

inject_model(body_json, model)

payload_str = json.dumps(body_json, ensure_ascii=False)

headers = _merge_headers(headers_from_log, access_token)

local_client = False

if client is None:

local_client = True

proxies = _get_proxies()

mounts = None

if proxies:

proxy_url = proxies.get("https") or proxies.get("http")

if proxy_url:

mounts = {

"https://": httpx.AsyncHTTPTransport(proxy=proxy_url),

"http://": httpx.AsyncHTTPTransport(proxy=proxy_url),

}

client = httpx.AsyncClient(mounts=mounts, timeout=httpx.Timeout(timeout[0], read=timeout[1]))

# Use manual request sending to control stream lifetime

req = client.build_request("POST", url, headers=headers, content=payload_str)

resp = None

try:

resp = await client.send(req, stream=True)

if resp.status_code >= 400:

try:

await resp.aread()

err = resp.text

except Exception:

err = f"HTTP {resp.status_code}"

await resp.aclose()

if local_client:

await client.aclose()

raise httpx.HTTPError(f"Upstream error {resp.status_code}: {err}")

parser = AwsEventStreamParser()

tracker = StreamTracker()

# Track if the response has been consumed to avoid double-close

response_consumed = False

async def _iter_events() -> AsyncGenerator[Any, None]:

nonlocal response_consumed

try:

if EventStreamParser and extract_event_info:

# Use proper EventStreamParser

async def byte_gen():

async for chunk in resp.aiter_bytes():

if chunk:

yield chunk

async for message in EventStreamParser.parse_stream(byte_gen()):

event_info = extract_event_info(message)

if event_info:

event_type = event_info.get('event_type')

payload = event_info.get('payload')

if event_type and payload:

yield (event_type, payload)

else:

# Fallback to old parser

async for chunk in resp.aiter_bytes():

if not chunk:

continue

events = parser.feed(chunk)

for ev_headers, payload in events:

parsed = _try_decode_event_payload(payload)

if parsed is not None:

event_type = None

if ":event-type" in ev_headers:

event_type = ev_headers[":event-type"]

yield (event_type, parsed)

except GeneratorExit:

# Client disconnected - ensure cleanup without re-raising

pass

except Exception:

if not tracker.has_content:

raise

finally:

response_consumed = True

if resp and not resp.is_closed:

await resp.aclose()

if local_client and client:

await client.aclose()

async def _iter_text() -> AsyncGenerator[str, None]:

async for event_type, parsed in _iter_events():

text = _extract_text_from_event(parsed)

if isinstance(text, str) and text:

yield text

def _schedule_cleanup():

"""Schedule cleanup when generator is GC'd without being consumed.

- If there's a running loop: spawn tasks for aclose()

- Else: try a synchronous close fallback (best-effort)

"""

try:

if not resp:

return

if not getattr(resp, "is_closed", True):

try:

loop = asyncio.get_event_loop()

if loop.is_running():

asyncio.create_task(resp.aclose())

if local_client and client:

asyncio.create_task(client.aclose())

else:

# No running loop: best-effort close

try:

# Prefer async close via a temporary loop

asyncio.run(resp.aclose())

if local_client and client:

asyncio.run(client.aclose())

except Exception:

# Fallback to sync close if available

if hasattr(resp, "close"):

try:

resp.close() # type: ignore[attr-defined]

except Exception:

pass

if local_client and client and hasattr(client, "close"):

try:

client.close() # type: ignore[attr-defined]

except Exception:

pass

except RuntimeError:

# No event loop; best-effort sync close

try:

asyncio.run(resp.aclose())

if local_client and client:

asyncio.run(client.aclose())

except Exception:

if hasattr(resp, "close"):

try:

resp.close() # type: ignore[attr-defined]

except Exception:

pass

if local_client and client and hasattr(client, "close"):

try:

client.close() # type: ignore[attr-defined]

except Exception:

pass

except Exception:

# Final safety: swallow exceptions in finalizer

pass

if stream:

# If raw_payload is used, we might want the raw event stream

if raw_payload:

# Return event stream with finalizer as safety net

event_gen = _iter_events()

weakref.finalize(event_gen, _schedule_cleanup)

return None, None, tracker, event_gen

# Return text stream with finalizer as safety net

text_gen = tracker.track(_iter_text())

weakref.finalize(text_gen, _schedule_cleanup)

return None, text_gen, tracker, None

else:

buf = []

try:

async for t in tracker.track(_iter_text()):

buf.append(t)

finally:

# Ensure response is closed even if iteration is incomplete

if not response_consumed and resp:

await resp.aclose()

if local_client:

await client.aclose()

return "".join(buf), None, tracker, None

except Exception:

# Critical: close response on any exception before generators are created

if resp and not resp.is_closed:

await resp.aclose()

if local_client and client:

await client.aclose()

raise