tokenator 0.1.0__tar.gz → 0.1.2__tar.gz

{tokenator-0.1.0 → tokenator-0.1.2}/PKG-INFO

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

{tokenator-0.1.0 → tokenator-0.1.2}/pyproject.toml

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

tokenator-0.1.2/src/tokenator/migrations.py

@@ -0,0 +1,38 @@
+"""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

{tokenator-0.1.0 → tokenator-0.1.2}/src/tokenator/schemas.py

@@ -2,6 +2,7 @@
 
 import uuid
 from datetime import datetime
+import os
 
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, Float, Index
 from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
@@ -12,7 +13,12 @@ Base = declarative_base()
 
 def get_engine(db_path: str = None):
     """Create SQLAlchemy engine with the given database path."""
-    db_path = db_path or get_default_db_path()
+    if db_path is None:
+        try:
+            import google.colab  # type: ignore
+            db_path = '/content/tokenator.db'
+        except ImportError:
+            db_path = get_default_db_path()
     return create_engine(f"sqlite:///{db_path}", echo=False)
 
 def get_session(db_path: str = None):

tokenator-0.1.2/src/tokenator/utils.py

@@ -0,0 +1,49 @@
+"""Shared utility functions for tokenator."""
+
+import os
+import platform
+import logging
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+def is_colab() -> bool:
+    """Check if running in Google Colab."""
+    try:
+        import google.colab  # type: ignore
+        return True
+    except ImportError:
+        return False
+
+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":
+            # Follow XDG Base Directory Specification
+            xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
+            if not xdg_data_home:
+                xdg_data_home = os.path.join(str(Path.home()), ".local", "share")
+            db_path = os.path.join(xdg_data_home, "tokenator", "usage.db")
+        elif system == "windows":
+            # Use %LOCALAPPDATA% on Windows
+            local_app_data = os.environ.get("LOCALAPPDATA", "")
+            if not local_app_data:
+                local_app_data = os.path.join(str(Path.home()), "AppData", "Local")
+            db_path = os.path.join(local_app_data, "tokenator", "usage.db")
+        else:
+            db_path = os.path.join(str(Path.home()), ".tokenator", "usage.db")
+        
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(db_path), exist_ok=True)
+        return db_path
+    except (OSError, IOError) as e:
+        # Fallback to current directory if we can't create the default path
+        fallback_path = os.path.join(os.getcwd(), "tokenator_usage.db")
+        logger.warning(f"Could not create default db path, falling back to {fallback_path}. Error: {e}")
+        return fallback_path 

tokenator-0.1.0/src/tokenator/migrations.py

@@ -1,40 +0,0 @@
-"""Automatic database migrations manager."""
-
-import os
-from pathlib import Path
-from alembic import command
-from alembic.config import Config
-from alembic.runtime.migration import MigrationContext
-from alembic.script import ScriptDirectory
-from sqlalchemy import create_engine
-
-from .utils import get_default_db_path
-
-def get_alembic_config():
-    """Get Alembic config pointing to the package's migrations."""
-    package_dir = Path(__file__).parent
-    migrations_dir = package_dir / "migrations"
-    
-    alembic_cfg = Config()
-    alembic_cfg.set_main_option("script_location", str(migrations_dir))
-    alembic_cfg.set_main_option("sqlalchemy.url", f"sqlite:///{get_default_db_path()}")
-    
-    return alembic_cfg
-
-def check_and_run_migrations():
-    """Check if migrations are needed and run them automatically."""
-    engine = create_engine(f"sqlite:///{get_default_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 latest available revision
-    config = get_alembic_config()
-    script = ScriptDirectory.from_config(config)
-    head_rev = script.get_current_head()
-    
-    # Run migrations if needed
-    if current_rev != head_rev:
-        command.upgrade(config, "head") 

tokenator-0.1.0/src/tokenator/utils.py

@@ -1,35 +0,0 @@
-"""Shared utility functions for tokenator."""
-
-import os
-import platform
-from pathlib import Path
-from typing import Optional
-
-def get_default_db_path() -> str:
-    """Get the platform-specific default database path.
-    
-    Returns:
-        str: Path to the SQLite database file
-        
-    The path follows platform conventions:
-        - Linux/macOS: ~/.local/share/tokenator/usage.db (XDG spec)
-        - Windows: %LOCALAPPDATA%\\tokenator\\usage.db
-        - Others: ~/.tokenator/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:
-            xdg_data_home = os.path.join(str(Path.home()), ".local", "share")
-        return os.path.join(xdg_data_home, "tokenator", "usage.db")
-    elif system == "windows":
-        # Use %LOCALAPPDATA% on Windows
-        local_app_data = os.environ.get("LOCALAPPDATA", "")
-        if not local_app_data:
-            local_app_data = os.path.join(str(Path.home()), "AppData", "Local")
-        return os.path.join(local_app_data, "tokenator", "usage.db")
-    else:
-        # Fallback for other systems
-        return os.path.join(str(Path.home()), ".tokenator", "usage.db")