tokenstretcher 1.0.0__py3-none-any.whl

tokensaver/__init__.py

@@ -0,0 +1,76 @@
+"""
+TokenSaverAI — Hierarchical AI Task Manager for Token & Cost Efficiency
+
+Public API:
+    from tokensaver import Manager, TokenSaverConfig
+    from tokensaver.manager import Manager
+    from tokensaver.config import load_config
+
+    result = await Manager().run("Build a complete FastAPI auth system")
+"""
+
+from __future__ import annotations
+
+__version__ = "1.0.0"
+__all__ = [
+    "Manager",
+    "Agent",
+    "TokenSaverConfig",
+    "ManagerResult",
+    "load_config",
+    "run",
+    "run_json",
+]
+
+from tokensaver.agent import Agent
+from tokensaver.config import load_config
+from tokensaver.manager import Manager
+from tokensaver.manager import ManagerResult
+from tokensaver.models import TokenSaverConfig
+
+
+async def run(prompt: str, **config_overrides) -> "ManagerResult":
+    """
+    Primary one-liner API for AI agents (Grok Build, OpenCode, etc.).
+
+    Recommended usage from another AI system:
+
+        from tokensaver import run
+        result = await run(
+            "Build a production FastAPI auth system with JWT + refresh tokens",
+            proxy_mode=False,           # or True if using prepaid billing
+            verbose=False               # cleaner output for agents
+        )
+
+        print(result.final_output)
+        print(result.savings.summary())
+    """
+    from tokensaver.manager import ManagerResult
+
+    cfg = load_config()
+    for k, v in config_overrides.items():
+        if hasattr(cfg, k):
+            setattr(cfg, k, v)
+
+    mgr = Manager(config=cfg)
+    return await mgr.run(prompt, return_report=True)
+
+
+async def run_json(prompt: str, **config_overrides) -> dict:
+    """
+    Convenience wrapper that returns a clean dictionary suitable for JSON serialization.
+
+    Especially useful when calling from Grok Build or other agent frameworks.
+
+    Example:
+        from tokensaver import run_json
+        data = await run_json("Refactor the auth module", verbose=False)
+        print(data["savings"]["percent_saved"])
+    """
+    result = await run(prompt, **config_overrides)
+    return {
+        "final_output": result.final_output,
+        "savings": result.savings.model_dump(),
+        "plan_reasoning": result.plan.reasoning,
+        "total_duration_seconds": result.total_duration_seconds,
+    }

tokensaver/agent.py

@@ -0,0 +1,144 @@
+"""
+Specialized Agent implementation for TokenSaverAI.
+
+Each Agent is intentionally narrow in scope and context.
+This is the key to massive token savings.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from litellm import acompletion
+from rich.console import Console
+
+from tokensaver.models import AgentResult, ModelTier, Subtask, TaskStatus, TokenSaverConfig
+from tokensaver.prompts import build_agent_system_prompt
+from tokensaver.utils import (
+    call_llm_with_retry,
+    estimate_cost,
+    filter_context_for_task,
+    get_llm_call_kwargs,
+)
+
+console = Console()
+
+
+class Agent:
+    """
+    A highly specialized, context-constrained AI worker.
+
+    Agents receive:
+    - A precise role + system prompt
+    - Only the minimal context they actually need
+    - A single focused task
+
+    They are cheap and fast by design.
+    """
+
+    def __init__(
+        self,
+        subtask: Subtask,
+        config: TokenSaverConfig,
+        full_project_context: str | None = None,
+    ):
+        self.subtask = subtask
+        self.config = config
+        self.full_project_context = full_project_context or ""
+
+        llm_kwargs = get_llm_call_kwargs(config)
+        force_xai = llm_kwargs.get("force_xai_models", False)
+        self.model = config.get_model_for_tier(subtask.model_preference, force_xai_models=force_xai)
+
+        # Build the tight system prompt for this specialist
+        filtered_context = filter_context_for_task(
+            self.full_project_context,
+            subtask.context_keywords,
+            max_chars=4200 if subtask.model_preference != ModelTier.POWERFUL else 8500,
+        )
+        self.system_prompt = build_agent_system_prompt(subtask.role, filtered_context)
+
+    async def run(self, user_instruction_override: str | None = None) -> AgentResult:
+        """
+        Execute the agent's assigned subtask.
+        Returns structured result with real token and cost data when available.
+        """
+        start = time.perf_counter()
+        instruction = user_instruction_override or self.subtask.description
+
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": instruction},
+        ]
+
+        try:
+            llm_kwargs = get_llm_call_kwargs(self.config)
+            resp = await call_llm_with_retry(
+                acompletion,
+                model=self.model,
+                messages=messages,
+                max_tokens=2800 if self.subtask.model_preference != ModelTier.CHEAP_FAST else 1600,
+                temperature=0.2,
+                timeout=self.config.default_timeout_seconds,
+                **llm_kwargs,
+            )
+
+            output = resp.choices[0].message.content or "(empty response)"
+            usage = getattr(resp, "usage", None)
+
+            if usage:
+                tokens = int(getattr(usage, "total_tokens", 0) or 0)
+                # LiteLLM sometimes puts cost on the response
+                cost = float(getattr(resp, "_response_cost", 0.0) or 0.0)
+                if cost == 0.0:
+                    cost = estimate_cost(tokens, self.model)
+            else:
+                # Fallback estimation
+                tokens = int(len(output) / 3.5) + int(len(self.system_prompt) / 3.8) + 180
+                cost = estimate_cost(tokens, self.model)
+
+            duration = time.perf_counter() - start
+
+            if self.config.verbose:
+                console.print(
+                    f"[dim]  ✓ {self.subtask.id} ({self.subtask.role[:30]}) — "
+                    f"{tokens:,} tokens — ${cost:.5f} — {duration:.1f}s[/dim]"
+                )
+
+            return AgentResult(
+                task_id=self.subtask.id,
+                role=self.subtask.role,
+                output=output.strip(),
+                model_used=self.model,
+                tokens_used=tokens,
+                cost_usd=round(cost, 6),
+                duration_seconds=round(duration, 2),
+                status=TaskStatus.COMPLETED,
+            )
+
+        except Exception as exc:
+            duration = time.perf_counter() - start
+            console.print(f"[red]  ✗ Agent {self.subtask.id} failed: {exc}[/red]")
+
+            return AgentResult(
+                task_id=self.subtask.id,
+                role=self.subtask.role,
+                output="",
+                model_used=self.model,
+                tokens_used=0,
+                cost_usd=0.0,
+                duration_seconds=round(duration, 2),
+                status=TaskStatus.FAILED,
+                error=str(exc),
+            )
+
+
+async def run_agent(
+    subtask: Subtask,
+    config: TokenSaverConfig,
+    context: str | None = None,
+) -> AgentResult:
+    """Convenience wrapper."""
+    agent = Agent(subtask, config, context)
+    return await agent.run()

tokensaver/cli.py

@@ -0,0 +1,354 @@
+"""
+TokenSaverAI Command Line Interface.
+
+Usage examples:
+    tokensaver "Build a FastAPI JWT authentication service with user management"
+    tokensaver "Design a scalable multi-tenant SaaS backend" --budget 0.80
+    tokensaver interactive
+    tokensaver --plan-only "Refactor the entire auth module"
+
+DISCLAIMER
+----------
+Commands related to proxy management, virtual keys, and prepaid billing
+are provided for convenience only. The authors accept no liability for
+any financial, operational, or legal issues arising from commercial use
+of these features. See LICENSE for full terms.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import sys
+from pathlib import Path
+
+from rich.console import Console
+from rich.prompt import Prompt
+
+from tokensaver.config import load_config, save_example_config
+from tokensaver.manager import Manager
+from tokensaver.models import TokenSaverConfig
+from tokensaver.wallet import Wallet
+from tokensaver.key_manager import KeyManager
+from tokensaver.email import send_key_email
+
+console = Console()
+
+
+def create_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="tokensaver",
+        description="TokenSaverAI — Hierarchical AI task manager for maximum quality at minimum cost",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    # Common flags that apply to task execution
+    parser.add_argument("-i", "--interactive", action="store_true", help="Launch interactive REPL mode")
+    parser.add_argument("--plan-only", action="store_true", help="Only show the decomposition plan")
+    parser.add_argument("--model", dest="powerful_model", help="Override powerful model (e.g. xai/grok-4)")
+    parser.add_argument("--budget", type=float, default=None, help="Soft max cost in USD")
+    parser.add_argument("--config", type=Path, default=None, help="Path to custom config")
+    parser.add_argument("--quiet", action="store_true", help="Reduce output")
+    parser.add_argument(
+        "--json", action="store_true",
+        help="Output structured JSON (strongly recommended when called by Grok Build, OpenCode, or other AI agents)"
+    )
+    parser.add_argument("--init-config", action="store_true", help="Write example config and exit")
+    from tokensaver import __version__
+    parser.add_argument("--version", action="version", version=f"TokenSaverAI {__version__}")
+
+    # Subcommands (wallet + proxy management)
+    subparsers = parser.add_subparsers(dest="command", help="Management commands (optional)")
+
+    # balance
+    p_balance = subparsers.add_parser("balance", help="Show current prepaid balance and spending")
+    p_balance.add_argument("--detailed", action="store_true", help="Show recent transactions")
+
+    # topup
+    p_topup = subparsers.add_parser("topup", help="Show instructions to add funds (Stripe / Lemon Squeezy)")
+
+    # add-funds
+    p_add = subparsers.add_parser("add-funds", help="Credit your local wallet (for testing or manual topups)")
+    p_add.add_argument("amount", type=float, help="Amount in USD to add to prepaid balance")
+
+    # proxy
+    p_proxy = subparsers.add_parser("proxy", help="Virtual key & proxy management")
+    proxy_sub = p_proxy.add_subparsers(dest="proxy_action")
+    proxy_sub.add_parser("start", help="Instructions to start the budget-enforcing proxy server")
+
+    p_create = proxy_sub.add_parser("create-key", help="Create a virtual key for a user and email it")
+    p_create.add_argument("email", help="User email address")
+    p_create.add_argument("budget", type=float, nargs="?", default=10.0, help="Initial budget in USD (default 10)")
+
+    p_disable = proxy_sub.add_parser("disable-key", help="Revoke a virtual key")
+    p_disable.add_argument("key", help="The virtual key to disable")
+
+    proxy_sub.add_parser("dashboard", help="Show all virtual keys and wallet status")
+
+    # config
+    subparsers.add_parser("config", help="Interactive setup for proxy mode and payment links")
+
+    return parser
+
+
+async def run_headless(prompt: str, args: argparse.Namespace, config: TokenSaverConfig) -> None:
+    """Execute a single prompt and print the beautiful result."""
+    if args.quiet:
+        config.verbose = False
+
+    if args.powerful_model:
+        config.default_powerful_model = args.powerful_model
+
+    if args.budget:
+        config.max_cost_usd = args.budget
+
+    mgr = Manager(config=config)
+    result = await mgr.run(prompt, return_report=True)
+
+    if args.json:
+        # Machine-friendly output for Grok Build, OpenCode, Cursor agents, etc.
+        import json
+        output = {
+            "final_output": result.final_output,
+            "savings": result.savings.model_dump(),
+            "plan": {
+                "reasoning": result.plan.reasoning,
+                "subtasks": [t.model_dump() for t in result.plan.subtasks],
+                "estimated_total_tokens": result.plan.estimated_total_tokens,
+            },
+            "agent_results": [r.model_dump() for r in result.agent_results],
+            "total_duration_seconds": result.total_duration_seconds,
+            "recursion_depth": result.recursion_depth,
+        }
+        print(json.dumps(output, indent=2, default=str))
+    elif not args.quiet:
+        from tokensaver.utils import print_final_result
+        print_final_result(result)
+    else:
+        print(result.final_output)
+
+
+async def run_interactive(config: TokenSaverConfig) -> None:
+    """Simple interactive session (great for exploration and demos)."""
+    console.rule("[bold cyan]TokenSaverAI Interactive Mode[/bold cyan]")
+    console.print("Type your complex task and press Enter. Type 'exit' or 'quit' to leave.\n")
+
+    mgr = Manager(config=config)
+
+    while True:
+        try:
+            prompt = Prompt.ask("[bold green]Task[/bold green]")
+        except (EOFError, KeyboardInterrupt):
+            console.print("\n[yellow]Goodbye.[/yellow]")
+            break
+
+        if prompt.lower().strip() in {"exit", "quit", "q"}:
+            break
+        if not prompt.strip():
+            continue
+
+        try:
+            result = await mgr.run(prompt)
+            from tokensaver.utils import print_final_result
+            print_final_result(result)
+            console.print("\n[dim]--- next task ---\n[/dim]")
+        except Exception as e:
+            console.print(f"[red]Error: {e}[/red]")
+
+
+def handle_wallet_commands(args: argparse.Namespace) -> None:
+    wallet = Wallet()
+
+    if args.command == "balance":
+        bal = wallet.get_balance()
+        console.print(f"\n[bold cyan]TokenSaverAI Prepay Wallet[/bold cyan]")
+        console.print(f"  Prepaid Balance : [bold green]${bal.prepaid_balance:.2f}[/bold green]")
+        console.print(f"  Total Spent     : ${bal.total_spent:.2f}")
+        console.print(f"  Total Credited  : ${bal.total_credited:.2f}")
+        console.print(f"  Last Updated    : {bal.last_updated}")
+
+        if getattr(args, "detailed", False):
+            txs = wallet.get_recent_transactions(8)
+            if txs:
+                console.print("\n[bold]Recent Transactions:[/bold]")
+                for ts, typ, amt, src, desc in txs:
+                    sign = "+" if typ == "credit" else "-"
+                    console.print(f"  {ts[:19]}  {sign}${amt:.4f}  ({src or desc or typ})")
+
+    elif args.command == "topup":
+        console.print("\n[bold cyan]How to Add Funds to Your TokenSaverAI Wallet[/bold cyan]\n")
+        console.print("1. Visit the payment link provided during onboarding (Stripe or Lemon Squeezy).")
+        console.print("2. After payment, run:")
+        console.print("      [bold]tokensaver add-funds 25[/bold]     (local/manual credit)")
+        console.print("   or contact support for automatic crediting.")
+        console.print("\nThis prepay model guarantees you can never receive a surprise bill.")
+
+    elif args.command == "add-funds":
+        amount = args.amount
+        new_bal = wallet.credit(amount, source="cli")
+        console.print(f"[green]✓ Credited ${amount:.2f}[/green]")
+        console.print(f"New prepaid balance: [bold]${new_bal.prepaid_balance:.2f}[/bold]")
+
+    elif args.command == "config":
+        console.print("\n[bold]TokenSaverAI Configuration Wizard[/bold]")
+        console.print("This will help you set up Proxy Mode (recommended for production use).")
+        console.print("\nFor now, edit .tokensaver/config.toml and set proxy_mode = true when ready.")
+
+
+def handle_proxy_commands(args: argparse.Namespace) -> None:
+    key_mgr = KeyManager()
+    cfg = load_config()  # Load early for proxy commands (port, etc.)
+
+    if args.proxy_action == "start":
+        console.print("[bold cyan]Starting TokenSaverAI Proxy Server...[/bold cyan]\n")
+        console.print("Loading configuration from .env (including XAI_API_KEY)...\n")
+
+        try:
+            import uvicorn
+            from tokensaver.proxy.server import app
+
+            port = getattr(cfg, "proxy_port", 8000)
+            console.print(f"[green]✓ Proxy starting on http://0.0.0.0:{port}[/green]")
+            console.print(f"[dim]Webhook URL: http://localhost:{port}/webhooks/lemonsqueezy[/dim]\n")
+
+            uvicorn.run(
+                "tokensaver.proxy.server:app",
+                host="0.0.0.0",
+                port=port,
+                reload=False,
+                log_level="info"
+            )
+        except ImportError:
+            console.print("[red]Proxy dependencies are missing.[/red]")
+            console.print("Install them with:")
+            console.print("  pip install -e \".[proxy]\"")
+            console.print("\nOr manually run:")
+            console.print(f"  uvicorn tokensaver.proxy.server:app --port 8000")
+        except Exception as e:
+            console.print(f"[red]Failed to start proxy: {e}[/red]")
+
+    elif args.proxy_action == "create-key":
+        email = args.email
+        budget = args.budget
+        vk = key_mgr.create_key(email, budget)
+
+        console.print(f"[green]✓ Virtual key created for {email}[/green]")
+        console.print(f"Key: [bold]{vk.key}[/bold]")
+        console.print(f"Budget: ${budget:.2f}")
+
+        # Try to email it
+        try:
+            config = load_config()
+            email_cfg = getattr(config, "email", {}) or {}
+            sent = send_key_email(
+                to_email=email,
+                virtual_key=vk.key,
+                budget_usd=budget,
+                provider=email_cfg.get("provider", "smtp"),
+                smtp_host=email_cfg.get("smtp_host"),
+                smtp_user=email_cfg.get("smtp_user"),
+                smtp_password=email_cfg.get("smtp_password"),
+                from_email=email_cfg.get("from_email", "noreply@tokensaver.ai"),
+                resend_api_key=email_cfg.get("resend_api_key"),
+            )
+            if sent:
+                console.print("[green]✓ Key emailed to user[/green]")
+            else:
+                console.print("[yellow]Key created but email failed to send. Share it manually.[/yellow]")
+        except Exception as e:
+            console.print(f"[yellow]Key created. Email sending failed: {e}[/yellow]")
+
+    elif args.proxy_action == "disable-key":
+        success = key_mgr.disable_key(args.key)
+        if success:
+            console.print(f"[green]✓ Key disabled: {args.key[:14]}...[/green]")
+        else:
+            console.print("[red]Key not found or already disabled[/red]")
+
+    elif args.proxy_action == "dashboard":
+        keys = key_mgr.list_keys(active_only=False)
+        console.print("\n[bold cyan]TokenSaverAI Proxy Dashboard[/bold cyan]\n")
+        console.print(f"Total Keys: {len(keys)}")
+        console.print(f"Active: {sum(1 for k in keys if k.is_active)}")
+
+        from tokensaver.wallet import Wallet
+        w = Wallet()
+        bal = w.get_balance()
+        console.print(f"Wallet Balance: [bold green]${bal.prepaid_balance:.2f}[/bold green]")
+
+        if keys:
+            console.print("\n[bold]Recent Keys:[/bold]")
+            for k in keys[:8]:
+                status = "ACTIVE" if k.is_active else "DISABLED"
+                rem = k.budget_usd - k.spent_usd
+                console.print(f"  {k.user_email:<25} | ${rem:>6.2f} left | {status}")
+
+        admin_token = os.getenv("ADMIN_TOKEN")
+        if not admin_token:
+            console.print("\n[yellow]Warning: No ADMIN_TOKEN set. Admin routes are unprotected.[/yellow]")
+
+
+def main() -> None:
+    parser = create_parser()
+    raw_args = sys.argv[1:]
+
+    known_commands = {"balance", "topup", "add-funds", "proxy", "config"}
+
+    # Find first non-flag token
+    first_pos = next((a for a in raw_args if not a.startswith("-")), None)
+
+    # Only attempt full subcommand parsing if the first token is a known command
+    if first_pos in known_commands:
+        args = parser.parse_args()
+        if args.init_config:
+            path = save_example_config()
+            console.print(f"[green]Example config written to {path}[/green]")
+            sys.exit(0)
+
+        if args.command in ("balance", "topup", "add-funds", "config"):
+            handle_wallet_commands(args)
+            return
+        if args.command == "proxy":
+            handle_proxy_commands(args)
+            return
+
+    # === Task execution mode ===
+    # Use parse_known_args and be defensive in case it still complains
+    try:
+        args, unknown = parser.parse_known_args()
+    except SystemExit:
+        # Fallback: treat the entire command line as a task prompt
+        prompt_tokens = [a for a in raw_args if not a.startswith("-")]
+        prompt = " ".join(prompt_tokens).strip()
+        args = argparse.Namespace(
+            interactive=False,
+            plan_only="--plan-only" in raw_args,
+            powerful_model=None,
+            budget=None,
+            config=None,
+            quiet="--quiet" in raw_args,
+            json="--json" in raw_args,
+            init_config="--init-config" in raw_args,
+        )
+        unknown = []
+    else:
+        prompt = " ".join(unknown).strip()
+
+    if args.init_config:
+        path = save_example_config()
+        console.print(f"[green]Example config written to {path}[/green]")
+        sys.exit(0)
+
+    config = load_config(args.config)
+    if args.quiet:
+        config.verbose = False
+
+    if args.interactive or not prompt:
+        asyncio.run(run_interactive(config))
+    else:
+        asyncio.run(run_headless(prompt, args, config))
+
+
+if __name__ == "__main__":
+    main()