tokenator 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

tokenator/migrations.py

@@ -1,40 +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 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
+logger = logging.getLogger(__name__)
 
 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") 
+    """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/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/utils.py

@@ -2,34 +2,48 @@
 
 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.
-    
-    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
-    """
+    """Get the platform-specific default database path."""
     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") 
+    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.dist-info → tokenator-0.1.2.dist-info}/METADATA

@@ -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.dist-info → tokenator-0.1.2.dist-info}/RECORD

@@ -5,12 +5,12 @@ tokenator/client_openai.py,sha256=1xZuRA90kwlflTwEuFkXJHHN584XTeNh1CfEBMLELbQ,63
 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=GH84M6AzidTcscMP0nQspBaQ1v6bjx00V4FeN-v5XAo,1354
+tokenator/migrations.py,sha256=RHm5XI5qh6W-Ib06vz4bXmE9XL211n1lZLzQNHPoSzg,1396
 tokenator/models.py,sha256=EprE_MMJxDS-YXlcIQLZzfekH7xTYbeOC3bx3B2osVw,1171
-tokenator/schemas.py,sha256=KQrKv_9Euoai3hBZcuiVclU3BlxQkdfEbwYoFV8MKXY,2112
+tokenator/schemas.py,sha256=eVdBWi6_hTETnPw50glq0OvSh3PbP2pLl_aHdf3fi-M,2278
 tokenator/usage.py,sha256=aHjGwzDzaiVznahNk5HqVyk3IxDo5FtFVfOUCeE7DZ4,7833
-tokenator/utils.py,sha256=VF-t6eL16_KokICMfZvO1rIJX9Hm7rexUcaeuLI2iRA,1311
-tokenator-0.1.0.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
-tokenator-0.1.0.dist-info/METADATA,sha256=dlsCQKC_2qVw4OEfj8ISrt5QwIY774PO_H2S3ilj92g,2444
-tokenator-0.1.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-tokenator-0.1.0.dist-info/RECORD,,
+tokenator/utils.py,sha256=UHV6tKLd6zoz7Fml1LokkbGmN1hvQMfXDY4Aulkhar8,1910
+tokenator-0.1.2.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.2.dist-info/METADATA,sha256=VnT8VlRydPmUPqNUesUjrhWWJT1PlliNriCRHAqrRtY,2444
+tokenator-0.1.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+tokenator-0.1.2.dist-info/RECORD,,