tokenator 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

tokenator/__init__.py

@@ -6,15 +6,13 @@ from . import usage
 from .utils import get_default_db_path, is_colab
 
 __version__ = "0.1.0"
-__all__ = ["OpenAIWrapper", "usage", "get_default_db_path"]
+__all__ = ["OpenAIWrapper", "usage", "get_default_db_path", "is_colab"]
 
 logger = logging.getLogger(__name__)
 
 try:
-    if not is_colab():
-        from .migrations import check_and_run_migrations
-        check_and_run_migrations()
-    else:
-        logger.info("Running in Colab environment - skipping migrations")
+    # Always run migrations, even in Colab
+    from .migrations import check_and_run_migrations
+    check_and_run_migrations()
 except Exception as e:
     logger.warning(f"Failed to run migrations, but continuing anyway: {e}") 

tokenator/create_migrations.py

@@ -1,11 +1,12 @@
 """Development utilities for tokenator."""
 
 import os
+import sys
 from pathlib import Path
 from alembic import command
 from tokenator.migrations import get_alembic_config
 
-def create_migration():
+def create_migration(message: str):
     """Create a new migration based on model changes."""
     config = get_alembic_config()
     
@@ -13,9 +14,12 @@ def create_migration():
     migrations_dir = Path(__file__).parent / "migrations" / "versions"
     migrations_dir.mkdir(parents=True, exist_ok=True)
     
-    # Generate migration
-    command.revision(config, autogenerate=True, message="auto generated migration") 
-
+    # Generate migration with custom message
+    command.revision(config, autogenerate=True, message=message) 
 
 if __name__ == "__main__":
-    create_migration()
+    if len(sys.argv) > 1:
+        msg = " ".join(sys.argv[1:])
+    else:
+        msg = "auto generated migration"
+    create_migration(msg)

tokenator/migrations/env.py

@@ -5,7 +5,7 @@ from sqlalchemy import pool
 
 from alembic import context
 
-from tokenator.models import Base
+from tokenator.schemas import Base
 
 # this is the Alembic Config object
 config = context.config

tokenator/migrations/versions/f6f1f2437513_initial_migration.py

@@ -0,0 +1,49 @@
+"""Initial migration
+
+Revision ID: f6f1f2437513
+Revises: 
+Create Date: 2024-12-21 17:33:27.187221
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'f6f1f2437513'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('token_usage',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('execution_id', sa.String(), nullable=False),
+    sa.Column('provider', sa.String(), nullable=False),
+    sa.Column('model', sa.String(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=False),
+    sa.Column('updated_at', sa.DateTime(), nullable=False),
+    sa.Column('prompt_tokens', sa.Integer(), nullable=False),
+    sa.Column('completion_tokens', sa.Integer(), nullable=False),
+    sa.Column('total_tokens', sa.Integer(), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index('idx_created_at', 'token_usage', ['created_at'], unique=False)
+    op.create_index('idx_execution_id', 'token_usage', ['execution_id'], unique=False)
+    op.create_index('idx_model', 'token_usage', ['model'], unique=False)
+    op.create_index('idx_provider', 'token_usage', ['provider'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index('idx_provider', table_name='token_usage')
+    op.drop_index('idx_model', table_name='token_usage')
+    op.drop_index('idx_execution_id', table_name='token_usage')
+    op.drop_index('idx_created_at', table_name='token_usage')
+    op.drop_table('token_usage')
+    # ### end Alembic commands ### 

tokenator/migrations.py

@@ -1,38 +1,39 @@
-"""Automatic database migrations manager."""
+"""Database migration utilities for tokenator."""
 
 import os
-import logging
 from pathlib import Path
-from alembic import command
 from alembic.config import Config
-from alembic.runtime.migration import MigrationContext
-from sqlalchemy import create_engine
-
+from alembic import command
 from .utils import get_default_db_path
 
-logger = logging.getLogger(__name__)
-
-def check_and_run_migrations():
-    """Check and run any pending database migrations."""
-    try:
+def get_alembic_config(db_path: str = None) -> Config:
+    """Get Alembic config for migrations."""
+    if db_path is None:
         db_path = get_default_db_path()
-        engine = create_engine(f"sqlite:///{db_path}")
 
-        # Create migrations table if it doesn't exist
-        with engine.connect() as conn:
-            context = MigrationContext.configure(conn)
-            current_rev = context.get_current_revision()
+    # Get the directory containing this file
+    migrations_dir = Path(__file__).parent / "migrations"
+    
+    # Create Config object
+    config = Config()
+    config.set_main_option("script_location", str(migrations_dir))
+    config.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
+    
+    return config
 
-        if current_rev is None:
-            # Run migrations
-            config = Config()
-            migrations_dir = os.path.join(os.path.dirname(__file__), "migrations")
-            config.set_main_option("script_location", migrations_dir)
-            config.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
-            
-            command.upgrade(config, "head")
-            logger.info("Database migrations completed successfully")
-    except Exception as e:
-        logger.error(f"Failed to run migrations: {e}")
-        # Don't raise the exception - allow the application to continue
-        # The user can manually initialize the DB later if needed
+def check_and_run_migrations(db_path: str = None):
+    """Check and run any pending database migrations."""
+    if db_path is None:
+        db_path = get_default_db_path()
+        
+    dirname = os.path.dirname(db_path)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+    
+    # Initialize database
+    import sqlite3
+    conn = sqlite3.connect(db_path)
+    conn.close()
+    
+    config = get_alembic_config(db_path)
+    command.upgrade(config, "head")

tokenator/schemas.py

@@ -24,7 +24,7 @@ def get_engine(db_path: str = None):
 def get_session(db_path: str = None):
     """Create a thread-safe session factory."""
     engine = get_engine(db_path)
-    Base.metadata.create_all(engine)
+    # Base.metadata.create_all(engine)
     session_factory = sessionmaker(bind=engine)
     return scoped_session(session_factory)
 

tokenator/utils.py

@@ -18,13 +18,14 @@ def is_colab() -> bool:
 
 def get_default_db_path() -> str:
     """Get the platform-specific default database path."""
-    system = platform.system().lower()
-    
     try:
         if is_colab():
-            # We're in Colab - use current directory
-            db_path = os.path.join(os.getcwd(), "tokenator_usage.db")
-        elif system == "linux" or system == "darwin":
+            # Use in-memory database for Colab
+            return "usage.db"
+            
+        system = platform.system().lower()
+        
+        if system == "linux" or system == "darwin":
             # Follow XDG Base Directory Specification
             xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
             if not xdg_data_home:

{tokenator-0.1.3.dist-info → tokenator-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tokenator
-Version: 0.1.3
+Version: 0.1.5
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari

tokenator-0.1.5.dist-info/RECORD

@@ -0,0 +1,17 @@
+tokenator/__init__.py,sha256=YsjMbOFztjrOKWNyguFcadzPKwhj3uaw0gRJQkCoRaM,566
+tokenator/base_wrapper.py,sha256=vSu_pStKYulho7_5g0jMCNf84KRxC4kTKep0v8YE61M,2377
+tokenator/client_anthropic.py,sha256=1ejWIZBxtk-mWTVaKWeMUvS2hZ_Dn-vNKYa3yopdjAU,6714
+tokenator/client_openai.py,sha256=1xZuRA90kwlflTwEuFkXJHHN584XTeNh1CfEBMLELbQ,6308
+tokenator/create_migrations.py,sha256=n1OVbWrdwvBdaN-Aqqt1gLCPQidfoQfeJtGsab_epGk,746
+tokenator/migrations/env.py,sha256=LR_hONDa8Saiq9CyNUpH8kZCi5PtXLaDlfABs_CePkk,1415
+tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
+tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=DvHcjnREmUHZVX9q1e6PS4wNK_d4qGw-8pz0eS4_3mE,1860
+tokenator/migrations.py,sha256=BFgZRsdIx-Qs_WwDaH6cyi2124mLf5hA8VrIlW7f7Mg,1134
+tokenator/models.py,sha256=EprE_MMJxDS-YXlcIQLZzfekH7xTYbeOC3bx3B2osVw,1171
+tokenator/schemas.py,sha256=V7NYfY9eZvH3J6uOwXJz4dSAU6WYzINRnfFi1wWsTcc,2280
+tokenator/usage.py,sha256=aHjGwzDzaiVznahNk5HqVyk3IxDo5FtFVfOUCeE7DZ4,7833
+tokenator/utils.py,sha256=BzfyWZkKt-2Jw_DCS7tY3iicBynklgvmf_1cPGe-OEI,1883
+tokenator-0.1.5.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.5.dist-info/METADATA,sha256=jVSteW5iVYoqLT727GZUIdmuPNtG3rK7cvPqEBLMmHo,2444
+tokenator-0.1.5.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+tokenator-0.1.5.dist-info/RECORD,,

tokenator-0.1.3.dist-info/RECORD

@@ -1,16 +0,0 @@
-tokenator/__init__.py,sha256=SB81-PEkyU-JHuckDNqeoopfpzbfL1jDweLaETPa4J0,626
-tokenator/base_wrapper.py,sha256=vSu_pStKYulho7_5g0jMCNf84KRxC4kTKep0v8YE61M,2377
-tokenator/client_anthropic.py,sha256=1ejWIZBxtk-mWTVaKWeMUvS2hZ_Dn-vNKYa3yopdjAU,6714
-tokenator/client_openai.py,sha256=1xZuRA90kwlflTwEuFkXJHHN584XTeNh1CfEBMLELbQ,6308
-tokenator/create_migrations.py,sha256=C_3WqB0tOGKXOA4JmvWuLpcyGEysWyRSiSttxX-Kie4,606
-tokenator/migrations/env.py,sha256=eFTw66gG464JV53740RKU32wqEL8uZFReS_INrvkFrU,1414
-tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
-tokenator/migrations.py,sha256=RHm5XI5qh6W-Ib06vz4bXmE9XL211n1lZLzQNHPoSzg,1396
-tokenator/models.py,sha256=EprE_MMJxDS-YXlcIQLZzfekH7xTYbeOC3bx3B2osVw,1171
-tokenator/schemas.py,sha256=eVdBWi6_hTETnPw50glq0OvSh3PbP2pLl_aHdf3fi-M,2278
-tokenator/usage.py,sha256=aHjGwzDzaiVznahNk5HqVyk3IxDo5FtFVfOUCeE7DZ4,7833
-tokenator/utils.py,sha256=UHV6tKLd6zoz7Fml1LokkbGmN1hvQMfXDY4Aulkhar8,1910
-tokenator-0.1.3.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
-tokenator-0.1.3.dist-info/METADATA,sha256=kOesX0EPrxsqvrowcayXbA8phU7Ix9xPmp6Jqb_fYHM,2444
-tokenator-0.1.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-tokenator-0.1.3.dist-info/RECORD,,