tokenator 0.1.3__tar.gz → 0.1.5__tar.gz

{tokenator-0.1.3 → tokenator-0.1.5}/PKG-INFO

@@ -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.3 → tokenator-0.1.5}/pyproject.toml

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

{tokenator-0.1.3 → tokenator-0.1.5}/src/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-0.1.3 → tokenator-0.1.5}/src/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-0.1.3 → tokenator-0.1.5}/src/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-0.1.5/src/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-0.1.5/src/tokenator/migrations.py

@@ -0,0 +1,39 @@
+"""Database migration utilities for tokenator."""
+
+import os
+from pathlib import Path
+from alembic.config import Config
+from alembic import command
+from .utils import get_default_db_path
+
+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()
+
+    # 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
+
+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-0.1.3 → tokenator-0.1.5}/src/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-0.1.3 → tokenator-0.1.5}/src/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/src/tokenator/migrations.py

@@ -1,38 +0,0 @@
-"""Automatic database migrations manager."""
-
-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 .utils import get_default_db_path
-
-logger = logging.getLogger(__name__)
-
-def check_and_run_migrations():
-    """Check and run any pending database migrations."""
-    try:
-        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()
-
-        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