feat: new refresh_prices method for synchronous data fetching

.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

pyproject.toml

@@ -10,14 +10,15 @@ tokencost = ["model_prices.json"]
 
 [project]
 name = "tokencost"
-version = "0.1.20"
+version = "0.1.21"
 authors = [
   { name = "Trisha Pan", email = "[email protected]" },
   { name = "Alex Reibman", email = "[email protected]" },
+  { name = "Pratyush Shukla", email = "[email protected]" },
 ]
 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",
@@ -34,13 +35,10 @@ dev = [
     "pytest>=7.4.4",
     "flake8>=3.1.0",
     "coverage[toml]>=7.4.0",
-    "tach==0.6.9",
+    "tach>=0.6.9",
     "tabulate>=0.9.0",
     "pandas>=2.1.0",
 ]
-llama-index = [
-    "llama-index>=0.10.23"
-]
 
 [project.urls]
 Homepage = "https://github.com/AgentOps-AI/tokencost"

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

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

tokencost/constants.py

@@ -52,11 +52,32 @@ 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."""
+    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:
+            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)
 
@@ -67,7 +88,6 @@ async def update_token_costs():
 # 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:

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 =

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,10 @@ 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### )"
+# 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
 
 updated_readme = re.sub(table_pattern, table_replacement, readme_content, flags=re.DOTALL)