diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 8ef77c8c461..e41a700c1d2 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -187,6 +187,7 @@ def __init__( compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, additional_properties: dict[str, Any] | None = None, + timeout: float | None = None, ) -> None: """Initialize a raw Microsoft Foundry chat client. @@ -207,6 +208,8 @@ def __init__( compaction_strategy: Optional per-client compaction override. tokenizer: Optional tokenizer for compaction strategies. additional_properties: Additional properties stored on the client instance. + timeout: HTTP timeout in seconds for requests. When not provided, the + OpenAI SDK default is used (connect: 5s, total: 600s). """ foundry_settings = load_settings( FoundrySettings, @@ -255,9 +258,13 @@ def __init__( if owns_project_client: openai_kwargs["http_client"] = create_foundry_feature_usage_http_client() + openai_client = project_client.get_openai_client(**openai_kwargs) + if timeout is not None: + openai_client = openai_client.with_options(timeout=timeout) + super().__init__( model=resolved_model, - async_client=project_client.get_openai_client(**openai_kwargs), + async_client=openai_client, default_headers=default_headers, instruction_role=instruction_role, compaction_strategy=compaction_strategy, @@ -1129,6 +1136,7 @@ def __init__( additional_properties: dict[str, Any] | None = None, middleware: (Sequence[ChatAndFunctionMiddlewareTypes] | None) = None, function_invocation_configuration: FunctionInvocationConfiguration | None = None, + timeout: float | None = None, ) -> None: """Initialize a Foundry chat client. @@ -1149,6 +1157,8 @@ def __init__( additional_properties: Additional properties stored on the client instance. middleware: Optional sequence of middleware. function_invocation_configuration: Optional function invocation configuration. + timeout: HTTP timeout in seconds for requests. When not provided, the + OpenAI SDK default is used (connect: 5s, total: 600s). """ super().__init__( project_endpoint=project_endpoint, @@ -1165,4 +1175,5 @@ def __init__( additional_properties=additional_properties, middleware=middleware, function_invocation_configuration=function_invocation_configuration, + timeout=timeout, ) diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 490fb2cb3ce..a686057ff04 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -290,6 +290,30 @@ def test_foundry_chat_client_init_uses_explicit_parameters() -> None: assert all(parameter.kind != inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()) +def test_init_applies_timeout_to_openai_client() -> None: + mock_openai_client = _make_mock_openai_client() + timeout_client = _make_mock_openai_client() + mock_openai_client.with_options.return_value = timeout_client + mock_project_client = MagicMock() + mock_project_client.get_openai_client.return_value = mock_openai_client + + client = FoundryChatClient(project_client=mock_project_client, model=_TEST_FOUNDRY_MODEL, timeout=12.5) + + mock_openai_client.with_options.assert_called_once_with(timeout=12.5) + assert client.client is timeout_client + + +def test_init_without_timeout_keeps_openai_client() -> None: + mock_openai_client = _make_mock_openai_client() + mock_project_client = MagicMock() + mock_project_client.get_openai_client.return_value = mock_openai_client + + client = FoundryChatClient(project_client=mock_project_client, model=_TEST_FOUNDRY_MODEL) + + mock_openai_client.with_options.assert_not_called() + assert client.client is mock_openai_client + + def test_init_with_default_header() -> None: default_headers = {"X-Unit-Test": "test-guid"} mock_openai_client = _make_mock_openai_client()