added check to limit headers that can be used

unifyai · Jul 22, 2024 · 132a9fd · 132a9fd
1 parent 802b345
commit 132a9fd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/unify/clients.py b/unify/clients.py
@@ -5,6 +5,7 @@
 from unify.exceptions import BadRequestError, UnifyError, status_error_map
 from unify.utils import (  # noqa:WPS450
     _available_dynamic_modes,
+    _available_extra_headers,
     _validate_api_key,
     _validate_endpoint,
 )
@@ -178,6 +179,7 @@ def generate(  # noqa: WPS234, WPS211
         else:
             raise UnifyError("You must provider either the user_prompt or messages!")
 
+        extra_headers = {key: extra_headers[key] for key in _available_extra_headers}
         if stream:
             return self._generate_stream(
                 contents,
@@ -480,6 +482,7 @@ async def generate(  # noqa: WPS234, WPS211
         else:
             raise UnifyError("You must provide either the user_prompt or messages!")
 
+        extra_headers = {key: extra_headers[key] for key in _available_extra_headers}
         if stream:
             return self._generate_stream(
                 contents,

diff --git a/unify/utils.py b/unify/utils.py
@@ -18,6 +18,10 @@
     "tks-per-sec",
 ]
 
+_available_extra_headers = [
+    "anthropic-beta"
+]
+
 _base_url = "https://api.unify.ai/v0"