mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-19 04:14:46 +08:00
fix: fix proxy
This commit is contained in:
@@ -51,29 +51,22 @@ EXCLUDED_HEADERS = {
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
async def stream_response(client: httpx.AsyncClient, method: str, url: str, headers: dict, body: bytes | None, timeout: float) -> AsyncGenerator[bytes, None]:
|
async def stream_sse_response(stream_ctx, response: httpx.Response) -> AsyncGenerator[bytes, None]:
|
||||||
"""Stream response from the upstream server.
|
"""Stream SSE response from the upstream server.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
client: The httpx async client.
|
stream_ctx: The httpx stream context manager.
|
||||||
method: HTTP method.
|
response: The httpx streaming response.
|
||||||
url: Target URL.
|
|
||||||
headers: Request headers.
|
|
||||||
body: Request body.
|
|
||||||
timeout: Request timeout.
|
|
||||||
|
|
||||||
Yields:
|
Yields:
|
||||||
Response chunks.
|
Response chunks.
|
||||||
"""
|
"""
|
||||||
async with client.stream(
|
try:
|
||||||
method=method,
|
|
||||||
url=url,
|
|
||||||
headers=headers,
|
|
||||||
content=body,
|
|
||||||
timeout=timeout,
|
|
||||||
) as response:
|
|
||||||
async for chunk in response.aiter_bytes():
|
async for chunk in response.aiter_bytes():
|
||||||
yield chunk
|
yield chunk
|
||||||
|
finally:
|
||||||
|
# Ensure stream is properly closed when done
|
||||||
|
await stream_ctx.__aexit__(None, None, None)
|
||||||
|
|
||||||
|
|
||||||
async def proxy_request(request: Request, path: str) -> Response | StreamingResponse:
|
async def proxy_request(request: Request, path: str) -> Response | StreamingResponse:
|
||||||
@@ -104,22 +97,26 @@ async def proxy_request(request: Request, path: str) -> Response | StreamingResp
|
|||||||
client = get_http_client()
|
client = get_http_client()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# First, make a non-streaming request to check content type
|
# Use streaming request to avoid waiting for full response
|
||||||
response = await client.request(
|
# This allows us to check headers immediately and stream SSE without delay
|
||||||
|
stream_ctx = client.stream(
|
||||||
method=request.method,
|
method=request.method,
|
||||||
url=target_url,
|
url=target_url,
|
||||||
headers=headers,
|
headers=headers,
|
||||||
content=body,
|
content=body,
|
||||||
timeout=config.proxy_timeout,
|
timeout=config.stream_timeout,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
response = await stream_ctx.__aenter__()
|
||||||
|
|
||||||
content_type = response.headers.get("content-type", "")
|
content_type = response.headers.get("content-type", "")
|
||||||
|
|
||||||
# Check if response is SSE (Server-Sent Events)
|
# Check if response is SSE (Server-Sent Events)
|
||||||
if "text/event-stream" in content_type:
|
if "text/event-stream" in content_type:
|
||||||
# For SSE, we need to re-request with streaming
|
# For SSE, stream the response immediately
|
||||||
return StreamingResponse(
|
return StreamingResponse(
|
||||||
stream_response(client, request.method, target_url, headers, body, config.stream_timeout),
|
stream_sse_response(stream_ctx, response),
|
||||||
|
status_code=response.status_code,
|
||||||
media_type="text/event-stream",
|
media_type="text/event-stream",
|
||||||
headers={
|
headers={
|
||||||
"Cache-Control": "no-cache",
|
"Cache-Control": "no-cache",
|
||||||
@@ -128,13 +125,17 @@ async def proxy_request(request: Request, path: str) -> Response | StreamingResp
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# For non-SSE responses, read full content and close the stream
|
||||||
|
content = await response.aread()
|
||||||
|
await stream_ctx.__aexit__(None, None, None)
|
||||||
|
|
||||||
# Prepare response headers
|
# Prepare response headers
|
||||||
response_headers = dict(response.headers)
|
response_headers = dict(response.headers)
|
||||||
for header in ["transfer-encoding", "connection", "keep-alive"]:
|
for header in ["transfer-encoding", "connection", "keep-alive"]:
|
||||||
response_headers.pop(header, None)
|
response_headers.pop(header, None)
|
||||||
|
|
||||||
return Response(
|
return Response(
|
||||||
content=response.content,
|
content=content,
|
||||||
status_code=response.status_code,
|
status_code=response.status_code,
|
||||||
headers=response_headers,
|
headers=response_headers,
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user