vectorvein 0.1.51__py3-none-any.whl → 0.1.52__py3-none-any.whl

vectorvein/chat_clients/utils.py

@@ -2,6 +2,7 @@
 # @Date:   2024-07-26 14:48:55
 import re
 import json
+import warnings
 from math import ceil
 from typing import Iterable
 
@@ -177,7 +178,7 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         if isinstance(endpoint_id, dict):
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
-        tokenize_url = "https://api.moonshot.cn/v1/tokenizers/estimate-token-count"
+        tokenize_url = f"{endpoint.api_base}/tokenizers/estimate-token-count"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
         request_body = {
             "model": model,
@@ -226,11 +227,29 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         result = response.json()
         return result["totalTokens"]
     elif model.startswith("claude"):
-        return (
-            Anthropic()
-            .beta.messages.count_tokens(messages=[{"role": "user", "content": text}], model=model)
-            .input_tokens
-        )
+        backend_settings = settings.get_backend(BackendType.Anthropic)
+        for endpoint_choice in backend_settings.models[model].endpoints:
+            if isinstance(endpoint_choice, dict):
+                endpoint_id = endpoint_choice["endpoint_id"]
+            else:
+                endpoint_id = endpoint_choice
+            endpoint = settings.get_endpoint(endpoint_id)
+
+            if endpoint.is_vertex:
+                continue
+            elif endpoint.api_schema_type == "default":
+                return (
+                    Anthropic(
+                        api_key=endpoint.api_key,
+                        base_url=endpoint.api_base,
+                    )
+                    .beta.messages.count_tokens(messages=[{"role": "user", "content": text}], model=model)
+                    .input_tokens
+                )
+
+        # TODO: Use anthropic token counting
+        warnings.warn("Anthropic token counting is not implemented yet")
+        return len(get_gpt_4o_encoding().encode(text))
     elif model.startswith("deepseek"):
         from deepseek_tokenizer import deepseek_tokenizer
 
@@ -248,7 +267,7 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         if isinstance(endpoint_id, dict):
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
-        tokenize_url = "https://api.stepfun.com/v1/token/count"
+        tokenize_url = f"{endpoint.api_base}/token/count"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
         request_body = {
             "model": model,

{vectorvein-0.1.51.dist-info → vectorvein-0.1.52.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.51
+Version: 0.1.52
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT

{vectorvein-0.1.51.dist-info → vectorvein-0.1.52.dist-info}/RECORD

@@ -1,6 +1,6 @@
-vectorvein-0.1.51.dist-info/METADATA,sha256=Z6zB6JLQZR-jyMfrG8gldHY6DQSdC3hjEDJKoWnDUOc,644
-vectorvein-0.1.51.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-vectorvein-0.1.51.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.52.dist-info/METADATA,sha256=7Tl7WKTmbcQa7SG_fgPtxaZJbR8pOGzM5mOi07ivvpo,644
+vectorvein-0.1.52.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+vectorvein-0.1.52.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
 vectorvein/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/chat_clients/__init__.py,sha256=dW169oK1n3v8Z0uD8itghzlCP72rxiaS-XYn6fvI2xM,16788
 vectorvein/chat_clients/anthropic_client.py,sha256=i1fMYSkUovd-Lc9B64bWMgzSTUFw4S3fj3AJ_pDokT4,34029
@@ -18,7 +18,7 @@ vectorvein/chat_clients/openai_compatible_client.py,sha256=D2VmhpDVct4w2y58s87An
 vectorvein/chat_clients/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/chat_clients/qwen_client.py,sha256=-ryh-m9PgsO0fc4ulcCmPTy1155J8YUy15uPoJQOHA0,513
 vectorvein/chat_clients/stepfun_client.py,sha256=zsD2W5ahmR4DD9cqQTXmJr3txrGuvxbRWhFlRdwNijI,519
-vectorvein/chat_clients/utils.py,sha256=gL2B_q3FeCqLGml8rpfKfXkTHcBQlxP5U3jH9eUEsa8,25604
+vectorvein/chat_clients/utils.py,sha256=HUPtdn-OSxghJNCF5Q8PsL3Ye-1Nu16O84lntPqKyao,26439
 vectorvein/chat_clients/yi_client.py,sha256=RNf4CRuPJfixrwLZ3-DEc3t25QDe1mvZeb9sku2f8Bc,484
 vectorvein/chat_clients/zhipuai_client.py,sha256=Ys5DSeLCuedaDXr3PfG1EW2zKXopt-awO2IylWSwY0s,519
 vectorvein/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -32,4 +32,4 @@ vectorvein/types/llm_parameters.py,sha256=5o-C_yXxxQWZy_e8OWowB2107GTS-Eawx4Mvb1
 vectorvein/types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/utilities/media_processing.py,sha256=cnzLrU1OaJvSv87IOnc36FrDXtmGMDStPbxtIJ33YN4,5880
 vectorvein/utilities/retry.py,sha256=6KFS9R2HdhqM3_9jkjD4F36ZSpEx2YNFGOVlpOsUetM,2208
-vectorvein-0.1.51.dist-info/RECORD,,
+vectorvein-0.1.52.dist-info/RECORD,,