85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
"""Common utilities shared between img2typ and solve scripts."""
|
|
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
from rich.console import Console
|
|
from rich.logging import RichHandler
|
|
from rich.theme import Theme
|
|
|
|
console = Console(
|
|
theme=Theme({"info": "cyan", "warning": "yellow", "error": "bold red"})
|
|
)
|
|
|
|
SCRIPT_DIR = Path(__file__).parent
|
|
DATA_DIR = SCRIPT_DIR.parent / "data"
|
|
GLOBAL_ENV_DIR = Path.home() / ".config" / "phomework"
|
|
GLOBAL_ENV_FILE = GLOBAL_ENV_DIR / ".env"
|
|
LOCAL_ENV_FILE = SCRIPT_DIR / ".env"
|
|
|
|
|
|
def setup_logging(name: str, verbose: bool) -> logging.Logger:
|
|
"""Setup logging with rich handler."""
|
|
level = logging.DEBUG if verbose else logging.INFO
|
|
logging.basicConfig(
|
|
level=level,
|
|
format="%(message)s",
|
|
datefmt="[%X]",
|
|
handlers=[RichHandler(console=console, rich_tracebacks=True)],
|
|
)
|
|
return logging.getLogger(name)
|
|
|
|
|
|
def load_env() -> dict[str, str]:
|
|
"""Load environment variables from .env file.
|
|
|
|
Priority: scripts/.env > ~/.config/phomework/.env
|
|
"""
|
|
env_loaded = False
|
|
|
|
if LOCAL_ENV_FILE.exists():
|
|
load_dotenv(LOCAL_ENV_FILE)
|
|
env_loaded = True
|
|
elif GLOBAL_ENV_FILE.exists():
|
|
load_dotenv(GLOBAL_ENV_FILE)
|
|
env_loaded = True
|
|
console.print(f"[dim]Using config from {GLOBAL_ENV_FILE}[/dim]")
|
|
else:
|
|
console.print(
|
|
f"[yellow]Warning: .env not found at {LOCAL_ENV_FILE} or {GLOBAL_ENV_FILE}[/yellow]"
|
|
)
|
|
|
|
api_endpoint = os.environ.get("IMG2TYP_API_ENDPOINT", "")
|
|
api_key = os.environ.get("IMG2TYP_API_KEY", "")
|
|
api_model = os.environ.get("IMG2TYP_MODEL", "qwen-vl-plus")
|
|
|
|
if not api_endpoint:
|
|
console.print("[yellow]Warning: IMG2TYP_API_ENDPOINT not set[/yellow]")
|
|
if not api_key:
|
|
console.print("[yellow]Warning: IMG2TYP_API_KEY not set[/yellow]")
|
|
|
|
return {"endpoint": api_endpoint, "key": api_key, "model": api_model}
|
|
|
|
|
|
def load_prompt(filename: str) -> str:
|
|
"""Load a prompt template from file."""
|
|
prompt_path = SCRIPT_DIR / filename
|
|
if not prompt_path.exists():
|
|
console.print(
|
|
f"[yellow]Warning: Prompt file not found at {prompt_path}[/yellow]"
|
|
)
|
|
return ""
|
|
return prompt_path.read_text(encoding="utf-8")
|
|
|
|
|
|
def find_attachments(question: str) -> list[str]:
|
|
"""Find all attachment files for a given question."""
|
|
attachments = []
|
|
question_prefix = question + "_"
|
|
for file_path in DATA_DIR.iterdir():
|
|
if file_path.is_file() and file_path.name.startswith(question_prefix):
|
|
attachments.append(file_path.name)
|
|
return sorted(attachments)
|