"""Utilities function"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['n_tokens', 'load_prompt']

# %% ../nbs/00_core.ipynb 2
import os
from pathlib import Path
import tiktoken

# %% ../nbs/00_core.ipynb 3
def _load_dotenv():
    """Load .env file if it exists and env vars aren't already set"""
    # Check if key env vars already exist
    if os.getenv('ANTHROPIC_API_KEY') or os.getenv('MISTRAL_API_KEY'):
        return  # Already set via export, don't override
    
    try:
        from dotenv import load_dotenv
        # For dev: look in project root (one level up from nbs/)
        env_path = Path(__file__).parent.parent / '.env'
        if env_path.exists():
            load_dotenv(env_path)
    except (ImportError, NameError):
        pass  # dotenv not installed or __file__ not available

# %% ../nbs/00_core.ipynb 4
# Auto-load on import
_load_dotenv()

# %% ../nbs/00_core.ipynb 5
def n_tokens(
    text:str, # Text to count tokens in
    model:str='gpt-4' # Model name for tokenizer
    ) -> int:
    "Count the number of tokens in a string"
    enc = tiktoken.encoding_for_model(model)
    return len(enc.encode(text))

# %% ../nbs/00_core.ipynb 7
def load_prompt(
    name:str, # Name of prompt file: 'gcm', 'select_sections', 'srf_ccps', 'srf_enablers', 'srf_outputs'
    path:Path|str|None=None # Directory containing prompt files, defaults to 'files/prompts'
    ) -> str:
    "Load a prompt template from a markdown file"
    if path is None:
        try: path = Path(__file__).parent / 'files' / 'prompts'
        except NameError: path = Path('files/prompts')  # notebook fallback
    return (Path(path)/f'{name}.md').read_text()
