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

Sign up to get free protection for your applications and to get access to all the features.
tokenator/__init__.py CHANGED
@@ -1,12 +1,20 @@
1
1
  """Tokenator - Track and analyze your OpenAI API token usage and costs."""
2
2
 
3
+ import logging
3
4
  from .client_openai import OpenAIWrapper
4
5
  from . import usage
5
- from .utils import get_default_db_path
6
- from .migrations import check_and_run_migrations
6
+ from .utils import get_default_db_path, is_colab
7
7
 
8
8
  __version__ = "0.1.0"
9
9
  __all__ = ["OpenAIWrapper", "usage", "get_default_db_path"]
10
10
 
11
- # Run migrations on import
12
- check_and_run_migrations()
11
+ logger = logging.getLogger(__name__)
12
+
13
+ try:
14
+ if not is_colab():
15
+ from .migrations import check_and_run_migrations
16
+ check_and_run_migrations()
17
+ else:
18
+ logger.info("Running in Colab environment - skipping migrations")
19
+ except Exception as e:
20
+ logger.warning(f"Failed to run migrations, but continuing anyway: {e}")
tokenator/migrations.py CHANGED
@@ -1,40 +1,38 @@
1
1
  """Automatic database migrations manager."""
2
2
 
3
3
  import os
4
+ import logging
4
5
  from pathlib import Path
5
6
  from alembic import command
6
7
  from alembic.config import Config
7
8
  from alembic.runtime.migration import MigrationContext
8
- from alembic.script import ScriptDirectory
9
9
  from sqlalchemy import create_engine
10
10
 
11
11
  from .utils import get_default_db_path
12
12
 
13
- def get_alembic_config():
14
- """Get Alembic config pointing to the package's migrations."""
15
- package_dir = Path(__file__).parent
16
- migrations_dir = package_dir / "migrations"
17
-
18
- alembic_cfg = Config()
19
- alembic_cfg.set_main_option("script_location", str(migrations_dir))
20
- alembic_cfg.set_main_option("sqlalchemy.url", f"sqlite:///{get_default_db_path()}")
21
-
22
- return alembic_cfg
13
+ logger = logging.getLogger(__name__)
23
14
 
24
15
  def check_and_run_migrations():
25
- """Check if migrations are needed and run them automatically."""
26
- engine = create_engine(f"sqlite:///{get_default_db_path()}")
27
-
28
- # Create migrations table if it doesn't exist
29
- with engine.connect() as conn:
30
- context = MigrationContext.configure(conn)
31
- current_rev = context.get_current_revision()
32
-
33
- # Get latest available revision
34
- config = get_alembic_config()
35
- script = ScriptDirectory.from_config(config)
36
- head_rev = script.get_current_head()
37
-
38
- # Run migrations if needed
39
- if current_rev != head_rev:
40
- command.upgrade(config, "head")
16
+ """Check and run any pending database migrations."""
17
+ try:
18
+ db_path = get_default_db_path()
19
+ engine = create_engine(f"sqlite:///{db_path}")
20
+
21
+ # Create migrations table if it doesn't exist
22
+ with engine.connect() as conn:
23
+ context = MigrationContext.configure(conn)
24
+ current_rev = context.get_current_revision()
25
+
26
+ if current_rev is None:
27
+ # Run migrations
28
+ config = Config()
29
+ migrations_dir = os.path.join(os.path.dirname(__file__), "migrations")
30
+ config.set_main_option("script_location", migrations_dir)
31
+ config.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
32
+
33
+ command.upgrade(config, "head")
34
+ logger.info("Database migrations completed successfully")
35
+ except Exception as e:
36
+ logger.error(f"Failed to run migrations: {e}")
37
+ # Don't raise the exception - allow the application to continue
38
+ # The user can manually initialize the DB later if needed
tokenator/utils.py CHANGED
@@ -2,37 +2,48 @@
2
2
 
3
3
  import os
4
4
  import platform
5
+ import logging
5
6
  from pathlib import Path
6
7
  from typing import Optional
7
8
 
9
+ logger = logging.getLogger(__name__)
10
+
11
+ def is_colab() -> bool:
12
+ """Check if running in Google Colab."""
13
+ try:
14
+ import google.colab # type: ignore
15
+ return True
16
+ except ImportError:
17
+ return False
18
+
8
19
  def get_default_db_path() -> str:
9
- """Get the platform-specific default database path.
10
-
11
- Returns:
12
- str: Path to the SQLite database file
13
-
14
- The path follows platform conventions:
15
- - Linux/macOS: ~/.local/share/tokenator/usage.db (XDG spec)
16
- - Windows: %LOCALAPPDATA%\\tokenator\\usage.db
17
- - Others: ~/.tokenator/usage.db
18
- """
20
+ """Get the platform-specific default database path."""
19
21
  system = platform.system().lower()
20
22
 
21
- if system == "linux" or system == "darwin":
22
- # Follow XDG Base Directory Specification
23
- xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
24
- if not xdg_data_home:
25
- xdg_data_home = os.path.join(str(Path.home()), ".local", "share")
26
- db_path = os.path.join(xdg_data_home, "tokenator", "usage.db")
27
- elif system == "windows":
28
- # Use %LOCALAPPDATA% on Windows
29
- local_app_data = os.environ.get("LOCALAPPDATA", "")
30
- if not local_app_data:
31
- local_app_data = os.path.join(str(Path.home()), "AppData", "Local")
32
- db_path = os.path.join(local_app_data, "tokenator", "usage.db")
33
- else:
34
- db_path = os.path.join(str(Path.home()), ".tokenator", "usage.db")
35
-
36
- # Create directory if it doesn't exist
37
- os.makedirs(os.path.dirname(db_path), exist_ok=True)
38
- return db_path
23
+ try:
24
+ if is_colab():
25
+ # We're in Colab - use current directory
26
+ db_path = os.path.join(os.getcwd(), "tokenator_usage.db")
27
+ elif system == "linux" or system == "darwin":
28
+ # Follow XDG Base Directory Specification
29
+ xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
30
+ if not xdg_data_home:
31
+ xdg_data_home = os.path.join(str(Path.home()), ".local", "share")
32
+ db_path = os.path.join(xdg_data_home, "tokenator", "usage.db")
33
+ elif system == "windows":
34
+ # Use %LOCALAPPDATA% on Windows
35
+ local_app_data = os.environ.get("LOCALAPPDATA", "")
36
+ if not local_app_data:
37
+ local_app_data = os.path.join(str(Path.home()), "AppData", "Local")
38
+ db_path = os.path.join(local_app_data, "tokenator", "usage.db")
39
+ else:
40
+ db_path = os.path.join(str(Path.home()), ".tokenator", "usage.db")
41
+
42
+ # Create directory if it doesn't exist
43
+ os.makedirs(os.path.dirname(db_path), exist_ok=True)
44
+ return db_path
45
+ except (OSError, IOError) as e:
46
+ # Fallback to current directory if we can't create the default path
47
+ fallback_path = os.path.join(os.getcwd(), "tokenator_usage.db")
48
+ logger.warning(f"Could not create default db path, falling back to {fallback_path}. Error: {e}")
49
+ return fallback_path
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: tokenator
3
- Version: 0.1.1
3
+ Version: 0.1.3
4
4
  Summary: Token usage tracking wrapper for LLMs
5
5
  License: MIT
6
6
  Author: Ujjwal Maheshwari
@@ -1,16 +1,16 @@
1
- tokenator/__init__.py,sha256=yBw4D1VQWe9fJfDvGavy1OFEEW0BIM2714dVk70HenE,363
1
+ tokenator/__init__.py,sha256=SB81-PEkyU-JHuckDNqeoopfpzbfL1jDweLaETPa4J0,626
2
2
  tokenator/base_wrapper.py,sha256=vSu_pStKYulho7_5g0jMCNf84KRxC4kTKep0v8YE61M,2377
3
3
  tokenator/client_anthropic.py,sha256=1ejWIZBxtk-mWTVaKWeMUvS2hZ_Dn-vNKYa3yopdjAU,6714
4
4
  tokenator/client_openai.py,sha256=1xZuRA90kwlflTwEuFkXJHHN584XTeNh1CfEBMLELbQ,6308
5
5
  tokenator/create_migrations.py,sha256=C_3WqB0tOGKXOA4JmvWuLpcyGEysWyRSiSttxX-Kie4,606
6
6
  tokenator/migrations/env.py,sha256=eFTw66gG464JV53740RKU32wqEL8uZFReS_INrvkFrU,1414
7
7
  tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
8
- tokenator/migrations.py,sha256=GH84M6AzidTcscMP0nQspBaQ1v6bjx00V4FeN-v5XAo,1354
8
+ tokenator/migrations.py,sha256=RHm5XI5qh6W-Ib06vz4bXmE9XL211n1lZLzQNHPoSzg,1396
9
9
  tokenator/models.py,sha256=EprE_MMJxDS-YXlcIQLZzfekH7xTYbeOC3bx3B2osVw,1171
10
10
  tokenator/schemas.py,sha256=eVdBWi6_hTETnPw50glq0OvSh3PbP2pLl_aHdf3fi-M,2278
11
11
  tokenator/usage.py,sha256=aHjGwzDzaiVznahNk5HqVyk3IxDo5FtFVfOUCeE7DZ4,7833
12
- tokenator/utils.py,sha256=WAMRY78d2g0PGSQZzPGoEuhJvQO5Iabgx4Wb6ykYp0Q,1407
13
- tokenator-0.1.1.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
14
- tokenator-0.1.1.dist-info/METADATA,sha256=ot6LXHhCJ6JrPOU6o18D-rgoarh2ZT0mn4uhPgJbcXU,2444
15
- tokenator-0.1.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
16
- tokenator-0.1.1.dist-info/RECORD,,
12
+ tokenator/utils.py,sha256=UHV6tKLd6zoz7Fml1LokkbGmN1hvQMfXDY4Aulkhar8,1910
13
+ tokenator-0.1.3.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
14
+ tokenator-0.1.3.dist-info/METADATA,sha256=kOesX0EPrxsqvrowcayXbA8phU7Ix9xPmp6Jqb_fYHM,2444
15
+ tokenator-0.1.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
16
+ tokenator-0.1.3.dist-info/RECORD,,