|
8 | 8 |
|
9 | 9 | import logging
|
10 | 10 | from contextlib import asynccontextmanager
|
| 11 | +from datetime import timedelta |
11 | 12 | from typing import Any
|
12 | 13 |
|
13 | 14 | import anyio
|
|
37 | 38 | async def streamablehttp_client(
|
38 | 39 | url: str,
|
39 | 40 | headers: dict[str, Any] | None = None,
|
40 |
| - timeout: float = 30, |
41 |
| - sse_read_timeout: float = 60 * 5, |
| 41 | + timeout: timedelta = timedelta(seconds=30), |
| 42 | + sse_read_timeout: timedelta = timedelta(seconds=60 * 5), |
42 | 43 | ):
|
43 | 44 | """
|
44 | 45 | Client transport for StreamableHTTP.
|
@@ -71,7 +72,9 @@ async def streamablehttp_client(
|
71 | 72 | session_id: str | None = None
|
72 | 73 |
|
73 | 74 | async with httpx.AsyncClient(
|
74 |
| - headers=request_headers, timeout=timeout, follow_redirects=True |
| 75 | + headers=request_headers, |
| 76 | + timeout=httpx.Timeout(timeout.seconds, read=sse_read_timeout.seconds), |
| 77 | + follow_redirects=True, |
75 | 78 | ) as client:
|
76 | 79 |
|
77 | 80 | async def post_writer():
|
@@ -225,7 +228,9 @@ async def get_stream():
|
225 | 228 | "GET",
|
226 | 229 | url,
|
227 | 230 | headers=get_headers,
|
228 |
| - timeout=httpx.Timeout(timeout, read=sse_read_timeout), |
| 231 | + timeout=httpx.Timeout( |
| 232 | + timeout.seconds, read=sse_read_timeout.seconds |
| 233 | + ), |
229 | 234 | ) as event_source:
|
230 | 235 | event_source.response.raise_for_status()
|
231 | 236 | logger.debug("GET SSE connection established")
|
|
0 commit comments