Add llama-index user agent (#14377)

implement
run-llama · Jun 25, 2024 · c090f32 · c090f32
1 parent ab8eae6
commit c090f32
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 13 deletions.
diff --git a/...ions/embeddings/llama-index-embeddings-deepinfra/llama_index/embeddings/deepinfra/base.py b/...ions/embeddings/llama-index-embeddings-deepinfra/llama_index/embeddings/deepinfra/base.py
@@ -19,6 +19,8 @@
 DEFAULT_MODEL_ID = "sentence-transformers/clip-ViT-B-32"
 """Maximum batch size for embedding requests."""
 MAX_BATCH_SIZE = 1024
+"""User agent"""
+USER_AGENT = "llama-index-embeddings-deepinfra"
 
 
 class DeepInfraEmbeddingModel(BaseEmbedding):
@@ -93,10 +95,7 @@ def _post(self, data: List[str]) -> List[List[float]]:
                 json={
                     "inputs": chunk,
                 },
-                headers={
-                    "Authorization": f"Bearer {self._api_token}",
-                    "Content-Type": "application/json",
-                },
+                headers=self._get_headers(),
             )
             response.raise_for_status()
             embeddings.extend(response.json()["embeddings"])
@@ -129,10 +128,7 @@ async def _apost(self, data: List[str]) -> List[List[float]]:
                     json={
                         "inputs": chunk,
                     },
-                    headers={
-                        "Authorization": f"Bearer {self._api_token}",
-                        "Content-Type": "application/json",
-                    },
+                    headers=self._get_headers(),
                 ) as resp:
                     response = await resp.json()
                     embeddings.extend(response["embeddings"])
@@ -206,6 +202,16 @@ def _add_text_prefix(self, texts: List[str]) -> List[str]:
             [self._text_prefix + text for text in texts] if self._text_prefix else texts
         )
 
+    def _get_headers(self) -> dict:
+        """
+        Get headers.
+        """
+        return {
+            "Authorization": f"Bearer {self._api_token}",
+            "Content-Type": "application/json",
+            "User-Agent": USER_AGENT,
+        }
+
 
 def _chunk(items: List[str], batch_size: int = MAX_BATCH_SIZE) -> List[List[str]]:
     """

diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-deepinfra/pyproject.toml b/llama-index-integrations/embeddings/llama-index-embeddings-deepinfra/pyproject.toml
@@ -27,7 +27,7 @@ license = "MIT"
 name = "llama-index-embeddings-deepinfra"
 packages = [{include = "llama_index/"}]
 readme = "README.md"
-version = "0.1.1"
+version = "0.1.2"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"

diff --git a/...a-index-integrations/llms/llama-index-llms-deepinfra/llama_index/llms/deepinfra/client.py b/...a-index-integrations/llms/llama-index-llms-deepinfra/llama_index/llms/deepinfra/client.py
@@ -7,9 +7,7 @@
     maybe_decode_sse_data,
     aretry_request,
 )
-from llama_index.llms.deepinfra.constants import (
-    API_BASE,
-)
+from llama_index.llms.deepinfra.constants import API_BASE, USER_AGENT
 
 
 class DeepInfraClient:
@@ -29,6 +27,7 @@ def _get_headers(self) -> Dict[str, str]:
         return {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
+            "User-Agent": USER_AGENT,
         }
 
     def request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:

diff --git a/...ndex-integrations/llms/llama-index-llms-deepinfra/llama_index/llms/deepinfra/constants.py b/...ndex-integrations/llms/llama-index-llms-deepinfra/llama_index/llms/deepinfra/constants.py
@@ -12,3 +12,5 @@
 DEFAULT_MAX_TOKENS = 512
 """Tool choice is auto."""
 TOOL_CHOICE = "auto"
+"""User agent"""
+USER_AGENT = "llama-index-llms-deepinfra"
diff --git a/llama-index-integrations/llms/llama-index-llms-deepinfra/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-deepinfra/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-deepinfra"
 readme = "README.md"
-version = "0.1.3"
+version = "0.1.4"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"