Skip to content

Commit

Permalink
Add support for default headers in LLM configurations
Browse files Browse the repository at this point in the history
- Introduced `default_headers` parameter in the `OpenAILLM` class to allow custom headers for OpenAI API requests.
- Updated unit tests for Anyscale, Azure, and Together LLMs to include `default_headers` in their configurations.

This change enhances flexibility in API interactions across multiple LLM implementations.
  • Loading branch information
davidberenstein1957 committed Jan 10, 2025
1 parent 1c81533 commit d823b9a
Show file tree
Hide file tree
Showing 4 changed files with 5 additions and 0 deletions.
1 change: 1 addition & 0 deletions src/distilabel/models/llms/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class OpenAILLM(AsyncLLM):
api_key: the API key to authenticate the requests to the OpenAI API. Defaults to
`None` which means that the value set for the environment variable `OPENAI_API_KEY`
will be used, or `None` if not set.
default_headers: the default headers to use for the OpenAI API requests.
max_retries: the maximum number of times to retry the request to the API before
failing. Defaults to `6`.
timeout: the maximum time in seconds to wait for a response from the API. Defaults
Expand Down
1 change: 1 addition & 0 deletions tests/unit/models/llms/test_anyscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def test_serialization(self) -> None:
"model": self.model_id,
"generation_kwargs": {},
"max_retries": 6,
"default_headers": None,
"base_url": "https://api.endpoints.anyscale.com/v1",
"timeout": 120,
"structured_output": None,
Expand Down
2 changes: 2 additions & 0 deletions tests/unit/models/llms/test_azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ def test_azure_openai_llm_env_vars(self) -> None:
"api_version": "preview",
"generation_kwargs": {},
"max_retries": 6,
"default_headers": None,
"base_url": "https://example-resource.azure.openai.com/",
"timeout": 120,
"structured_output": None,
Expand All @@ -95,6 +96,7 @@ def test_azure_openai_llm_env_vars(self) -> None:
"generation_kwargs": {},
"max_retries": 6,
"base_url": "https://example-resource.azure.openai.com/",
"default_headers": None,
"timeout": 120,
"structured_output": {
"schema": DummyUserDetail.model_json_schema(),
Expand Down
1 change: 1 addition & 0 deletions tests/unit/models/llms/test_together.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def test_serialization(self) -> None:
"model": self.model_id,
"generation_kwargs": {},
"max_retries": 6,
"default_headers": None,
"base_url": "https://api.together.xyz/v1",
"timeout": 120,
"structured_output": None,
Expand Down

0 comments on commit d823b9a

Please sign in to comment.