mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-12 21:00:07 +02:00
80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
from typing import TypedDict
|
|
|
|
from fastapi.datastructures import Headers
|
|
|
|
from onyx.configs.model_configs import LITELLM_EXTRA_HEADERS
|
|
from onyx.configs.model_configs import LITELLM_PASS_THROUGH_HEADERS
|
|
from onyx.configs.tool_configs import CUSTOM_TOOL_PASS_THROUGH_HEADERS
|
|
from onyx.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
class HeaderItemDict(TypedDict):
|
|
key: str
|
|
value: str
|
|
|
|
|
|
def clean_header_list(headers_to_clean: list[HeaderItemDict]) -> dict[str, str]:
|
|
cleaned_headers: dict[str, str] = {}
|
|
for item in headers_to_clean:
|
|
key = item["key"]
|
|
value = item["value"]
|
|
if key in cleaned_headers:
|
|
logger.warning(
|
|
f"Duplicate header {key} found in custom headers, ignoring..."
|
|
)
|
|
continue
|
|
cleaned_headers[key] = value
|
|
return cleaned_headers
|
|
|
|
|
|
def header_dict_to_header_list(header_dict: dict[str, str]) -> list[HeaderItemDict]:
|
|
return [{"key": key, "value": value} for key, value in header_dict.items()]
|
|
|
|
|
|
def header_list_to_header_dict(header_list: list[HeaderItemDict]) -> dict[str, str]:
|
|
return {header["key"]: header["value"] for header in header_list}
|
|
|
|
|
|
def get_relevant_headers(
|
|
headers: dict[str, str] | Headers, desired_headers: list[str] | None
|
|
) -> dict[str, str]:
|
|
if not desired_headers:
|
|
return {}
|
|
|
|
pass_through_headers: dict[str, str] = {}
|
|
for key in desired_headers:
|
|
if key in headers:
|
|
pass_through_headers[key] = headers[key]
|
|
else:
|
|
# fastapi makes all header keys lowercase, handling that here
|
|
lowercase_key = key.lower()
|
|
if lowercase_key in headers:
|
|
pass_through_headers[lowercase_key] = headers[lowercase_key]
|
|
|
|
return pass_through_headers
|
|
|
|
|
|
def get_litellm_additional_request_headers(
|
|
headers: dict[str, str] | Headers
|
|
) -> dict[str, str]:
|
|
return get_relevant_headers(headers, LITELLM_PASS_THROUGH_HEADERS)
|
|
|
|
|
|
def build_llm_extra_headers(
|
|
additional_headers: dict[str, str] | None = None
|
|
) -> dict[str, str]:
|
|
extra_headers: dict[str, str] = {}
|
|
if additional_headers:
|
|
extra_headers.update(additional_headers)
|
|
if LITELLM_EXTRA_HEADERS:
|
|
extra_headers.update(LITELLM_EXTRA_HEADERS)
|
|
return extra_headers
|
|
|
|
|
|
def get_custom_tool_additional_request_headers(
|
|
headers: dict[str, str] | Headers
|
|
) -> dict[str, str]:
|
|
return get_relevant_headers(headers, CUSTOM_TOOL_PASS_THROUGH_HEADERS)
|