[0.5.7] chore

Author: yashbonde

docs/changelog.rst

@@ -7,6 +7,11 @@ minor versions.
 
 All relevant steps to be taken will be mentioned here.
 
+0.5.7
+-----
+
+- Implement ``extra_headers`` via ``__init__`` as well.
+
 0.5.6
 -----
 

docs/conf.py

@@ -13,7 +13,7 @@
 project = "tuneapi"
 copyright = "2024, Frello Technologies"
 author = "Frello Technologies"
-release = "0.5.6"
+release = "0.5.7"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tuneapi"
-version = "0.5.6"
+version = "0.5.7"
 description = "Tune AI APIs."
 authors = ["Frello Technology Private Limited <[email protected]>"]
 license = "MIT"

tuneapi/apis/model_anthropic.py

@@ -18,10 +18,12 @@ def __init__(
         self,
         id: Optional[str] = "claude-3-haiku-20240307",
         base_url: str = "https://api.anthropic.com/v1/messages",
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id
         self.base_url = base_url
         self.api_token = tu.ENV.ANTHROPIC_TOKEN("")
+        self.extra_headers = extra_headers
 
     def set_api_token(self, token: str) -> None:
         self.api_token = token
@@ -159,6 +161,7 @@ def stream_chat(
             for t in tools:
                 t["input_schema"] = t.pop("parameters")
         headers, system, claude_messages = self._process_input(chats=chats, token=token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
 

tuneapi/apis/model_gemini.py

@@ -18,10 +18,12 @@ def __init__(
         self,
         id: Optional[str] = "gemini-1.5-pro-latest",
         base_url: str = "https://generativelanguage.googleapis.com/v1beta/models/{id}:{rpc}",
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id
         self.base_url = base_url
         self.api_token = tu.ENV.GEMINI_TOKEN("")
+        self.extra_headers = extra_headers
 
     def set_api_token(self, token: str) -> None:
         self.api_token = token
@@ -152,6 +154,7 @@ def stream_chat(
         if isinstance(chats, tt.Thread):
             tools = [x.to_dict() for x in chats.tools]
         headers, system, messages, params = self._process_input(chats, token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         data = {

tuneapi/apis/model_groq.py

@@ -17,10 +17,12 @@ def __init__(
         self,
         id: Optional[str] = "llama3-70b-8192",
         base_url: str = "https://api.groq.com/openai/v1/chat/completions",
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id
         self.base_url = base_url
         self.api_token = tu.ENV.GROQ_TOKEN("")
+        self.extra_headers = extra_headers
 
     def set_api_token(self, token: str) -> None:
         self.api_token = token
@@ -133,6 +135,7 @@ def stream_chat(
         if isinstance(chats, tt.Thread):
             tools = [{"type": "function", "function": x.to_dict()} for x in chats.tools]
         headers, messages = self._process_input(chats, token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         data = {

tuneapi/apis/model_mistral.py

@@ -19,10 +19,12 @@ def __init__(
         self,
         id: Optional[str] = "mistral-small-latest",
         base_url: str = "https://api.mistral.ai/v1/chat/completions",
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id
         self.base_url = base_url
         self.api_token = ENV.MISTRAL_TOKEN("")
+        self.extra_headers = extra_headers
 
     def set_api_token(self, token: str) -> None:
         self.api_token = token
@@ -136,6 +138,7 @@ def stream_chat(
         if isinstance(chats, Thread):
             tools = [{"type": "function", "function": x.to_dict()} for x in chats.tools]
         headers, messages = self._process_input(chats, token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         data = {

tuneapi/apis/model_openai.py

@@ -18,10 +18,12 @@ def __init__(
         self,
         id: Optional[str] = "gpt-4o",
         base_url: str = "https://api.openai.com/v1/chat/completions",
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id
         self.base_url = base_url
         self.api_token = tu.ENV.OPENAI_TOKEN("")
+        self.extra_headers = extra_headers
 
     def set_api_token(self, token: str) -> None:
         self.api_token = token
@@ -131,6 +133,7 @@ def stream_chat(
         raw: bool = False,
     ):
         headers, messages = self._process_input(chats, token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         data = {
@@ -187,8 +190,6 @@ def stream_chat(
             yield fn_call
         return
 
-    # def _process_chat_to_string_for_embedding(self, chat: tt.Thread):
-
     def embedding(
         self,
         chats: tt.Thread | List[str] | str,
@@ -203,6 +204,7 @@ def embedding(
         text = []
 
         headers = self._process_header(token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         if isinstance(chats, tt.Thread):

tuneapi/apis/model_tune.py

@@ -21,11 +21,13 @@ def __init__(
         base_url: str = "https://proxy.tune.app/chat/completions",
         api_token: Optional[str] = None,
         org_id: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         self.model_id = id or tu.ENV.TUNEAPI_MODEL("")
         self.base_url = base_url
         self.api_token = api_token or tu.ENV.TUNEAPI_TOKEN("")
         self.org_id = org_id or tu.ENV.TUNEORG_ID("")
+        self.extra_headers = extra_headers
 
     def __repr__(self) -> str:
         out = f"<TuneModel: {self.model_id}"
@@ -156,6 +158,7 @@ def stream_chat(
                 "Tune Model ID not found. Please set TUNEAPI_MODEL environment variable or pass through function"
             )
         headers, messages = self._process_input(chats, token)
+        extra_headers = extra_headers or self.extra_headers
         if extra_headers:
             headers.update(extra_headers)
         data = {

tuneapi/types/chats.py

@@ -321,6 +321,7 @@ def chat(
         temperature: float = 1,
         token: Optional[str] = None,
         timeout=(5, 30),
+        extra_headers: Optional[Dict[str, str]] = None,
         **kwargs,
     ) -> str | Dict[str, Any]:
         """This is the main function to block chat with the model"""
@@ -335,6 +336,7 @@ def stream_chat(
         timeout=(5, 60),
         raw: bool = False,
         debug: bool = False,
+        extra_headers: Optional[Dict[str, str]] = None,
     ):
         """This is the main function to stream chat with the model where each token is iteratively generated"""