tokenator 0.1.11__tar.gz → 0.1.13__tar.gz

{tokenator-0.1.11 → tokenator-0.1.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tokenator
-Version: 0.1.11
+Version: 0.1.13
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -171,6 +171,66 @@ print(usage.last_execution().model_dump_json(indent=4))
 }
 """
 ```
+
+### xAI
+
+You can use xAI models through the `openai` SDK and track usage using `provider` parameter in `tokenator`.
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+xai_client = OpenAI(
+            api_key=os.getenv("XAI_API_KEY"),
+            base_url="https://api.x.ai/v1"
+        )
+
+# Wrap it with Tokenator
+client = tokenator_openai(xai_client, db_path=temp_db, provider="xai")
+
+# Use it exactly like the OpenAI client but with xAI models
+response = client.chat.completions.create(
+    model="grok-2-latest",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+
+print(response)
+
+print(usage.last_execution())
+```
+
+### Other AI model providers through openai SDKs
+
+Today, a variety of AI companies have made their APIs compatible to the `openai` SDK.
+You can track usage of any such AI models using `tokenator`'s `provider` parameter.
+
+For example, let's see how we can track usage of `perplexity` tokens.
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+perplexity_client = OpenAI(
+            api_key=os.getenv("PERPLEXITY_API_KEY"),
+            base_url="https://api.perplexity.ai"
+        )
+
+# Wrap it with Tokenator
+client = tokenator_openai(perplexity_client, db_path=temp_db, provider="perplexity")
+
+# Use it exactly like the OpenAI client but with perplexity models
+response = client.chat.completions.create(
+    model="llama-3.1-sonar-small-128k-online",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+
+print(response)
+
+print(usage.last_execution())
+
+print(usage.provider("perplexity"))
+```
+
 ---
 
 Most importantly, none of your data is ever sent to any server.

{tokenator-0.1.11 → tokenator-0.1.13}/README.md

@@ -149,6 +149,66 @@ print(usage.last_execution().model_dump_json(indent=4))
 }
 """
 ```
+
+### xAI
+
+You can use xAI models through the `openai` SDK and track usage using `provider` parameter in `tokenator`.
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+xai_client = OpenAI(
+            api_key=os.getenv("XAI_API_KEY"),
+            base_url="https://api.x.ai/v1"
+        )
+
+# Wrap it with Tokenator
+client = tokenator_openai(xai_client, db_path=temp_db, provider="xai")
+
+# Use it exactly like the OpenAI client but with xAI models
+response = client.chat.completions.create(
+    model="grok-2-latest",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+
+print(response)
+
+print(usage.last_execution())
+```
+
+### Other AI model providers through openai SDKs
+
+Today, a variety of AI companies have made their APIs compatible to the `openai` SDK.
+You can track usage of any such AI models using `tokenator`'s `provider` parameter.
+
+For example, let's see how we can track usage of `perplexity` tokens.
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+perplexity_client = OpenAI(
+            api_key=os.getenv("PERPLEXITY_API_KEY"),
+            base_url="https://api.perplexity.ai"
+        )
+
+# Wrap it with Tokenator
+client = tokenator_openai(perplexity_client, db_path=temp_db, provider="perplexity")
+
+# Use it exactly like the OpenAI client but with perplexity models
+response = client.chat.completions.create(
+    model="llama-3.1-sonar-small-128k-online",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+
+print(response)
+
+print(usage.last_execution())
+
+print(usage.provider("perplexity"))
+```
+
 ---
 
 Most importantly, none of your data is ever sent to any server.

{tokenator-0.1.11 → tokenator-0.1.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tokenator"
-version = "0.1.11"
+version = "0.1.13"
 description = "Token usage tracking wrapper for LLMs"
 authors = ["Ujjwal Maheshwari <your.email@example.com>"]
 readme = "README.md"

{tokenator-0.1.11 → tokenator-0.1.13}/src/tokenator/openai/client_openai.py

@@ -14,7 +14,9 @@ logger = logging.getLogger(__name__)
 
 
 class BaseOpenAIWrapper(BaseWrapper):
-    provider = "openai"
+    def __init__(self, client, db_path=None, provider: str = "openai"):
+        super().__init__(client, db_path)
+        self.provider = provider
 
     def _process_response_usage(
         self, response: ResponseType
@@ -134,6 +136,7 @@ class AsyncOpenAIWrapper(BaseOpenAIWrapper):
 def tokenator_openai(
     client: OpenAI,
     db_path: Optional[str] = None,
+    provider: str = "openai",
 ) -> OpenAIWrapper: ...
 
 
@@ -141,23 +144,26 @@ def tokenator_openai(
 def tokenator_openai(
     client: AsyncOpenAI,
     db_path: Optional[str] = None,
+    provider: str = "openai",
 ) -> AsyncOpenAIWrapper: ...
 
 
 def tokenator_openai(
     client: Union[OpenAI, AsyncOpenAI],
     db_path: Optional[str] = None,
+    provider: str = "openai",
 ) -> Union[OpenAIWrapper, AsyncOpenAIWrapper]:
     """Create a token-tracking wrapper for an OpenAI client.
 
     Args:
         client: OpenAI or AsyncOpenAI client instance
         db_path: Optional path to SQLite database for token tracking
+        provider: Provider name, defaults to "openai"
     """
     if isinstance(client, OpenAI):
-        return OpenAIWrapper(client=client, db_path=db_path)
+        return OpenAIWrapper(client=client, db_path=db_path, provider=provider)
 
     if isinstance(client, AsyncOpenAI):
-        return AsyncOpenAIWrapper(client=client, db_path=db_path)
+        return AsyncOpenAIWrapper(client=client, db_path=db_path, provider=provider)
 
     raise ValueError("Client must be an instance of OpenAI or AsyncOpenAI")