From aa910e5cd9fd86172f3f00fa69fd8b374ab7c762 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 03:55:11 +0530 Subject: [PATCH 01/10] new `refresh_prices` for sync token update --- tokencost/__init__.py | 2 +- tokencost/constants.py | 24 ++++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/tokencost/__init__.py b/tokencost/__init__.py index c79f30c..3184582 100644 --- a/tokencost/__init__.py +++ b/tokencost/__init__.py @@ -6,4 +6,4 @@ calculate_all_costs_and_tokens, calculate_cost_by_tokens, ) -from .constants import TOKEN_COSTS_STATIC, TOKEN_COSTS, update_token_costs +from .constants import TOKEN_COSTS_STATIC, TOKEN_COSTS, update_token_costs, refresh_prices diff --git a/tokencost/constants.py b/tokencost/constants.py index b4dcea9..5de44ca 100644 --- a/tokencost/constants.py +++ b/tokencost/constants.py @@ -52,11 +52,35 @@ async def update_token_costs(): # Safely remove 'sample_spec' if it exists TOKEN_COSTS.update(fetched_costs) TOKEN_COSTS.pop("sample_spec", None) + return TOKEN_COSTS except Exception as e: logger.error(f"Failed to update TOKEN_COSTS: {e}") raise +def refresh_prices(write_file=True): + """Synchronous wrapper for update_token_costs that optionally writes to model_prices.json.""" + import asyncio + try: + # Run the async function in a new event loop + updated_costs = asyncio.run(update_token_costs()) + + # Write to file if requested + if write_file: + import json + import os + file_path = os.path.join(os.path.dirname(__file__), "model_prices.json") + with open(file_path, "w") as f: + json.dump(TOKEN_COSTS, f, indent=4) + logger.info(f"Updated prices written to {file_path}") + + return updated_costs + except Exception as e: + logger.error(f"Failed to refresh prices: {e}") + # Return the static prices as fallback + return TOKEN_COSTS + + with open(os.path.join(os.path.dirname(__file__), "model_prices.json"), "r") as f: TOKEN_COSTS_STATIC = json.load(f) From 52d195391f5d4b779d38e564adfcb0eb184a7e17 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 04:04:02 +0530 Subject: [PATCH 02/10] update script to use `refresh_prices` method --- update_prices.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/update_prices.py b/update_prices.py index b49b89f..f0eb4ad 100644 --- a/update_prices.py +++ b/update_prices.py @@ -5,6 +5,8 @@ import re # Update model_prices.json with the latest costs from the LiteLLM cost tracker +print("Fetching latest prices...") +tokencost.refresh_prices(write_file=False) def diff_dicts(dict1, dict2): @@ -21,19 +23,21 @@ def diff_dicts(dict1, dict2): else: print("No differences found.") - if differences: - return True - else: - return False + return bool(differences) +# Load the current file for comparison with open("tokencost/model_prices.json", "r") as f: model_prices = json.load(f) +# Compare the refreshed TOKEN_COSTS with the file if diff_dicts(model_prices, tokencost.TOKEN_COSTS): print("Updating model_prices.json") with open("tokencost/model_prices.json", "w") as f: json.dump(tokencost.TOKEN_COSTS, f, indent=4) + print("File updated successfully") +else: + print("File is already up to date") # Load the data df = pd.DataFrame(tokencost.TOKEN_COSTS).T @@ -52,7 +56,7 @@ def format_cost(x): else: price_per_million = Decimal(str(x)) * Decimal(str(1_000_000)) normalized = price_per_million.normalize() - formatted_price = "{:2f}".format(normalized) + formatted_price = "{:.2f}".format(normalized) formatted_price = ( formatted_price.rstrip("0").rstrip(".") @@ -103,7 +107,7 @@ def format_cost(x): readme_content = f.read() # Find and replace just the table in the README, preserving the header text -table_pattern = r"\| Model Name.*?(?=\n\n### )" +table_pattern = r"(?s)\| Model Name.*?\n\n(?=#)" table_replacement = table_md updated_readme = re.sub(table_pattern, table_replacement, readme_content, flags=re.DOTALL) From e0a2766234460d7c9a9cd92906e321a191161e8d Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 04:04:24 +0530 Subject: [PATCH 03/10] remove `llama-index` based code --- tests/test_llama_index_callbacks.py | 57 --------------------- tokencost/callbacks/__init__.py | 0 tokencost/callbacks/llama_index.py | 79 ----------------------------- 3 files changed, 136 deletions(-) delete mode 100644 tests/test_llama_index_callbacks.py delete mode 100644 tokencost/callbacks/__init__.py delete mode 100644 tokencost/callbacks/llama_index.py diff --git a/tests/test_llama_index_callbacks.py b/tests/test_llama_index_callbacks.py deleted file mode 100644 index 1169974..0000000 --- a/tests/test_llama_index_callbacks.py +++ /dev/null @@ -1,57 +0,0 @@ -# test_llama_index.py -import pytest -from tokencost.callbacks import llama_index -from llama_index.core.callbacks.schema import EventPayload - -# Mock the calculate_prompt_cost and calculate_completion_cost functions - -# 4 tokens -STRING = "Hello, world!" - - -# Mock the ChatMessage class in LlamaIndex -@pytest.fixture -def mock_chat_message(monkeypatch): - class MockChatMessage: - def __init__(self, text): - self.text = text - - def __str__(self): - return self.text - - monkeypatch.setattr("llama_index.core.llms.ChatMessage", MockChatMessage) - return MockChatMessage - - -# Test the _calc_llm_event_cost method for prompt and completion - - -def test_calc_llm_event_cost_prompt_completion(capsys): - handler = llama_index.TokenCostHandler(model="gpt-3.5-turbo") - payload = {EventPayload.PROMPT: STRING, EventPayload.COMPLETION: STRING} - handler._calc_llm_event_cost(payload) - captured = capsys.readouterr() - assert "# Prompt cost: 0.0000060" in captured.out - assert "# Completion: 0.000008" in captured.out - - -# Test the _calc_llm_event_cost method for messages and response - - -def test_calc_llm_event_cost_messages_response(mock_chat_message, capsys): - handler = llama_index.TokenCostHandler(model="gpt-3.5-turbo") - payload = { - EventPayload.MESSAGES: [ - mock_chat_message("message 1"), - mock_chat_message("message 2"), - ], - EventPayload.RESPONSE: "test response", - } - handler._calc_llm_event_cost(payload) - captured = capsys.readouterr() - assert "# Prompt cost: 0.0000105" in captured.out - assert "# Completion: 0.000004" in captured.out - - -# Additional tests can be written for start_trace, end_trace, on_event_start, and on_event_end -# depending on the specific logic and requirements of those methods. diff --git a/tokencost/callbacks/__init__.py b/tokencost/callbacks/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tokencost/callbacks/llama_index.py b/tokencost/callbacks/llama_index.py deleted file mode 100644 index d11c1ab..0000000 --- a/tokencost/callbacks/llama_index.py +++ /dev/null @@ -1,79 +0,0 @@ -from typing import Any, Dict, List, Optional, cast -from llama_index.core.callbacks.base_handler import BaseCallbackHandler -from llama_index.core.callbacks.schema import CBEventType, EventPayload -from llama_index.core.llms import ChatMessage -from tokencost import calculate_all_costs_and_tokens - - -class TokenCostHandler(BaseCallbackHandler): - """Callback handler for printing llms inputs/outputs.""" - - def __init__(self, model) -> None: - super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[]) - self.model = model - self.prompt_cost = 0 - self.completion_cost = 0 - self.prompt_tokens = 0 - self.completion_tokens = 0 - - def start_trace(self, trace_id: Optional[str] = None) -> None: - return - - def end_trace( - self, - trace_id: Optional[str] = None, - trace_map: Optional[Dict[str, List[str]]] = None, - ) -> None: - return - - def _calc_llm_event_cost(self, payload: dict) -> None: - if EventPayload.PROMPT in payload: - prompt = str(payload.get(EventPayload.PROMPT)) - completion = str(payload.get(EventPayload.COMPLETION)) - estimates = calculate_all_costs_and_tokens(prompt, completion, self.model) - - elif EventPayload.MESSAGES in payload: - messages = cast(List[ChatMessage], payload.get(EventPayload.MESSAGES, [])) - messages_str = "\n".join([str(x) for x in messages]) - response = str(payload.get(EventPayload.RESPONSE)) - estimates = calculate_all_costs_and_tokens( - messages_str, response, self.model - ) - else: - return - - self.prompt_cost += estimates["prompt_cost"] - self.completion_cost += estimates["completion_cost"] - self.prompt_tokens += estimates["prompt_tokens"] - self.completion_tokens += estimates["completion_tokens"] - - print(f"# Prompt cost: {estimates['prompt_cost']}") - print(f"# Completion: {estimates['completion_cost']}") - print("\n") - - def reset_counts(self) -> None: - self.prompt_cost = 0 - self.completion_cost = 0 - self.prompt_tokens = 0 - self.completion_tokens = 0 - - def on_event_start( - self, - event_type: CBEventType, - payload: Optional[Dict[str, Any]] = None, - event_id: str = "", - parent_id: str = "", - **kwargs: Any, - ) -> str: - return event_id - - def on_event_end( - self, - event_type: CBEventType, - payload: Optional[Dict[str, Any]] = None, - event_id: str = "", - **kwargs: Any, - ) -> None: - """Count the LLM or Embedding tokens as needed.""" - if event_type == CBEventType.LLM and payload is not None: - self._calc_llm_event_cost(payload) From 135af9376ec0ea722c37a549ff13c463f13e22b9 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 04:04:37 +0530 Subject: [PATCH 04/10] remove `llama-index` based code --- pyproject.toml | 8 +++----- tach.yml | 3 --- tox.ini | 13 ++++--------- 3 files changed, 7 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8d34c0c..e12fee9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,14 +10,15 @@ tokencost = ["model_prices.json"] [project] name = "tokencost" -version = "0.1.20" +version = "0.2.0" authors = [ { name = "Trisha Pan", email = "trishaepan@gmail.com" }, { name = "Alex Reibman", email = "areibman@gmail.com" }, + { name = "Pratyush Shukla", email = "ps4534@nyu.edu" }, ] description = "To calculate token and translated USD cost of string and message calls to OpenAI, for example when used by AI agents" readme = "README.md" -requires-python = ">=3.7" +requires-python = ">=3.10" classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: MIT License", @@ -38,9 +39,6 @@ dev = [ "tabulate>=0.9.0", "pandas>=2.1.0", ] -llama-index = [ - "llama-index>=0.10.23" -] [project.urls] Homepage = "https://github.com/AgentOps-AI/tokencost" diff --git a/tach.yml b/tach.yml index bf4cfdb..6a470c4 100644 --- a/tach.yml +++ b/tach.yml @@ -4,9 +4,6 @@ modules: depends_on: - tokencost.constants - tokencost.costs - - path: tokencost.callbacks.llama_index - depends_on: - - tokencost - path: tokencost.constants depends_on: [] - path: tokencost.costs diff --git a/tox.ini b/tox.ini index 48d7815..9a5e086 100644 --- a/tox.ini +++ b/tox.ini @@ -1,24 +1,19 @@ [tox] -envlist = py3, flake8, py3-llama-index +envlist = py3, flake8 isolated_build = true [testenv] deps = pytest coverage +commands = + coverage run --source tokencost -m pytest {posargs} + coverage report -m [testenv:flake8] deps = flake8 commands = flake8 tokencost/ -[testenv:py3-llama-index] -deps = - {[testenv]deps} - .[llama-index] -commands = - coverage run --source tokencost -m pytest {posargs} - coverage report -m - [flake8] max-line-length = 120 per-file-ignores = From 17b3c5ee0ea98ce15ec7e9c7fdb398c30c128c48 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 04:28:07 +0530 Subject: [PATCH 05/10] add comment explaining the regex pattern matching --- update_prices.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/update_prices.py b/update_prices.py index f0eb4ad..a89eabd 100644 --- a/update_prices.py +++ b/update_prices.py @@ -107,6 +107,9 @@ def format_cost(x): readme_content = f.read() # Find and replace just the table in the README, preserving the header text +# The regex pattern matches a markdown table starting with the "Model Name" header +# and ending before the next markdown header. DOTALL mode is enabled to allow +# the `.` character to match newline characters. table_pattern = r"(?s)\| Model Name.*?\n\n(?=#)" table_replacement = table_md From abbe54a8fcc046e08ca3ea123e3419697382796e Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 05:42:32 +0530 Subject: [PATCH 06/10] remove duplicate imports --- tokencost/constants.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tokencost/constants.py b/tokencost/constants.py index 5de44ca..1b56377 100644 --- a/tokencost/constants.py +++ b/tokencost/constants.py @@ -60,7 +60,6 @@ async def update_token_costs(): def refresh_prices(write_file=True): """Synchronous wrapper for update_token_costs that optionally writes to model_prices.json.""" - import asyncio try: # Run the async function in a new event loop updated_costs = asyncio.run(update_token_costs()) @@ -91,7 +90,6 @@ def refresh_prices(write_file=True): # Only run in a non-async context if __name__ == "__main__": try: - import asyncio asyncio.run(update_token_costs()) print("Token costs updated successfully") except Exception: From c62876ff8b3f063478ed9890176111c07e2e63d5 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 05:43:45 +0530 Subject: [PATCH 07/10] remove more duplicate imports --- tokencost/constants.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tokencost/constants.py b/tokencost/constants.py index 1b56377..619ef5a 100644 --- a/tokencost/constants.py +++ b/tokencost/constants.py @@ -66,8 +66,6 @@ def refresh_prices(write_file=True): # Write to file if requested if write_file: - import json - import os file_path = os.path.join(os.path.dirname(__file__), "model_prices.json") with open(file_path, "w") as f: json.dump(TOKEN_COSTS, f, indent=4) From 8de86d2159a01bec6503d4d6fb53430689b4bc56 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 05:49:54 +0530 Subject: [PATCH 08/10] use latest tach version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e12fee9..1eff442 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dev = [ "pytest>=7.4.4", "flake8>=3.1.0", "coverage[toml]>=7.4.0", - "tach==0.6.9", + "tach>=0.29.0", "tabulate>=0.9.0", "pandas>=2.1.0", ] From 90c675a6529ed288ced7e96dd470f115d3919192 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 22:33:03 +0530 Subject: [PATCH 09/10] use `tach==0.6.9` --- .github/workflows/ci.yml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0313dd0..029e54d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install tach + pip install tach==0.6.9 - name: Run Tach run: tach check \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1eff442..a71a764 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dev = [ "pytest>=7.4.4", "flake8>=3.1.0", "coverage[toml]>=7.4.0", - "tach>=0.29.0", + "tach>=0.6.9", "tabulate>=0.9.0", "pandas>=2.1.0", ] From 5590ade825c9ae159d2a8e5def697d7259162700 Mon Sep 17 00:00:00 2001 From: Pratyush Shukla Date: Wed, 7 May 2025 22:38:18 +0530 Subject: [PATCH 10/10] make this version `0.1.21` --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a71a764..057f6db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ tokencost = ["model_prices.json"] [project] name = "tokencost" -version = "0.2.0" +version = "0.1.21" authors = [ { name = "Trisha Pan", email = "trishaepan@gmail.com" }, { name = "Alex Reibman", email = "areibman@gmail.com" },