tinyprompt.py

#!/usr/bin/env python
"""
Tinyprompt: A simple CLI tool for semantic prompt compression using LLMLingua-2.

Compresses text prompts for LLMs, preserving meaning with 10-20x token reduction.
Pipe-friendly, fast (<1s for 1k tokens on CPU), minimal (~200MB deps).

Usage:
    echo "verbose prompt" | tinyprompt --ratio 0.3
    cat long_prompt.txt | tinyprompt --query "summarize" --verbose
    tinyprompt --input prompt.txt --format json | jq '.compressed_prompt'

Options:
    input              Input text or file (defaults to stdin if piped)
    --ratio FLOAT      Compression ratio (0.1-0.5, default 0.2)
    --query STR        Query for relevance boosting (optional)
    --model STR        Scorer model (default: gpt2-small)
    --verbose          Output metrics (tokens, ratio, BERTScore) to stderr
    --format STR       Output format: text (default) or json
    --cost-savings     Show per-run and total cost savings in USD
    --price-per-1k FLOAT  API price per 1k input tokens (default 0.01)
    --reset-stats      Reset cumulative savings counter

Examples:
    # Basic compression
    echo "This is a very long prompt with redundant info" | tinyprompt

    # Query-aware with cost savings
    cat system_prompt.txt | tinyprompt --query "focus on ethics" --cost-savings

    # JSON output with metrics
    tinyprompt --input prompt.txt --format json --verbose | llm -m gpt-4o

Install:
    uv venv
    source .venv/bin/activate
    uv pip install llmlingua transformers sentence-transformers tiktoken torch evaluate
"""

import sys
import argparse
import json
import os
from types import SimpleNamespace
from http.server import BaseHTTPRequestHandler, HTTPServer
import threading
import urllib.request
import urllib.error
import hashlib
from pathlib import Path
import platform
from http.server import ThreadingHTTPServer

def silence_transformers_warnings():
    try:
        from transformers.utils import logging as hf_logging  # type: ignore
        hf_logging.set_verbosity_error()
    except Exception:
        pass

# Singleton models for speed (lazy)
compressor = None
embedder = None
bertscore = None
encoder = None

def get_torch():
    import torch  # type: ignore
    return torch

def get_prompt_compressor():
    from llmlingua import PromptCompressor  # type: ignore
    return PromptCompressor

def get_sentence_transformer():
    from sentence_transformers import SentenceTransformer  # type: ignore
    return SentenceTransformer

def get_evaluate_load():
    from evaluate import load  # type: ignore
    return load

def get_encoder():
    global encoder
    if encoder is None:
        import tiktoken  # type: ignore
        encoder = tiktoken.get_encoding('cl100k_base')
    return encoder

# Stats file for cumulative savings
STATS_FILE = os.path.expanduser('~/.tinyprompt_stats.json')
DEFAULT_CACHE_DIR = os.path.expanduser('~/.cache/tinyprompt')

# Model presets (conservative choices that are known to work)
PRESETS = {
    'base': 'microsoft/llmlingua-2-xlm-roberta-large-meetingbank',
    'large': 'microsoft/llmlingua-2-xlm-roberta-large-original-multilingual',
}

def load_stats():
    """Load cumulative savings from JSON file."""
    try:
        if os.path.exists(STATS_FILE):
            with open(STATS_FILE, 'r') as f:
                return json.load(f)
        return {'total_savings': 0.0, 'run_count': 0}
    except Exception as e:
        print(f"Error loading stats: {e}", file=sys.stderr)
        return {'total_savings': 0.0, 'run_count': 0}

def save_stats(stats, per_run_savings):
    """Update and save cumulative savings."""
    try:
        stats['total_savings'] += per_run_savings
        stats['run_count'] += 1
        with open(STATS_FILE, 'w') as f:
            json.dump(stats, f)
    except Exception as e:
        print(f"Error saving stats: {e}", file=sys.stderr)

def load_models(args):
    """Initialize models only when needed."""
    global compressor, embedder, bertscore
    try:
        # Device selection: prefer CUDA, then MPS, else CPU; allow --cpu to force CPU
        device_map = 'cpu'
        torch = get_torch()
        if not getattr(args, 'cpu', False):
            if torch.cuda.is_available():
                device_map = 'cuda'
            elif getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
                device_map = 'mps'
        if compressor is None:
            silence_transformers_warnings()
            PromptCompressor = get_prompt_compressor()
            compressor = PromptCompressor(
                model_name=args.model,
                use_llmlingua2=not getattr(args, 'algo', None) == 'v1',
                device_map=device_map,
                llmlingua2_config=(
                    {'max_batch_size': getattr(args, 'max_batch_size', 50), 'max_force_token': getattr(args, 'max_force_token', 100)}
                    if getattr(args, 'max_batch_size', None) or getattr(args, 'max_force_token', None)
                    else {}
                ),
            )
            # Optionally reduce max sequence length for speed
            try:
                if getattr(args, 'max_seq_len', None):
                    setattr(compressor, 'max_seq_len', int(args.max_seq_len))
            except Exception:
                pass
        if args.query and embedder is None:
            SentenceTransformer = get_sentence_transformer()
            embedder = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
        if args.verbose and bertscore is None:
            try:
                eval_load = get_evaluate_load()
                bertscore = eval_load('bertscore')
            except Exception:
                # Optional dependency; skip metrics if unavailable
                print("Warning: BERTScore metric unavailable; skipping verbose metrics", file=sys.stderr)
                bertscore = None
    except Exception as e:
        print(f"Error loading models: {e}", file=sys.stderr)
        sys.exit(1)

def compress_prompt(input_text, args):
    """Compress using LLMLingua-2 with optional query-aware scoring."""
    try:
        torch = get_torch()
        with torch.no_grad():
            is_v2 = not getattr(args, 'algo', None) == 'v1'
            # Compute effective rate when target_tokens is specified
            if args.target_tokens is not None:
                orig_tokens = len(get_encoder().encode(input_text))
                # Avoid division by zero; clamp to sane range
                if args.target_tokens <= 0:
                    rate = 0.05
                else:
                    rate = max(0.01, min(0.99, args.target_tokens / max(1, orig_tokens)))
            else:
                rate = args.ratio

            # Disk cache
            use_cache = not getattr(args, 'no_cache', False)
            if use_cache:
                cache_dir = Path(getattr(args, 'cache_dir', None) or DEFAULT_CACHE_DIR) / 'results'
                try:
                    cache_dir.mkdir(parents=True, exist_ok=True)
                except Exception:
                    pass
                key_src = json.dumps({
                    'model': getattr(args, 'model', ''),
                    'query': getattr(args, 'query', None),
                    'rate': rate,
                    'text': input_text,
                }, sort_keys=True).encode('utf-8')
                key = hashlib.sha256(key_src).hexdigest()
                cache_path = cache_dir / f"{key}.json"
                if cache_path.exists():
                    try:
                        data = json.loads(cache_path.read_text(encoding='utf-8'))
                        if isinstance(data, dict) and 'compressed_prompt' in data:
                            return data['compressed_prompt']
                    except Exception:
                        pass

            if is_v2:
                # Always pass string to v2; library handles both but string path is standard
                try:
                    if args.query:
                        result = compressor.compress_prompt(
                            input_text,
                            rate=rate,
                            question_list=[args.query],
                            condition_in_question='after_condition'
                        )
                    else:
                        result = compressor.compress_prompt(input_text, rate=rate)
                except Exception:
                    # Fallback: reinitialize a clean v2 compressor and retry
                    PromptCompressor = get_prompt_compressor()
                    torch = get_torch()
                    device_map = 'cpu'
                    if not getattr(args, 'cpu', False):
                        if torch.cuda.is_available():
                            device_map = 'cuda'
                        elif getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
                            device_map = 'mps'
                    silence_transformers_warnings()
                    tmp = PromptCompressor(
                        model_name=PRESETS['base'],
                        use_llmlingua2=True,
                        device_map=device_map,
                    )
                    try:
                        if getattr(args, 'max_seq_len', None):
                            setattr(tmp, 'max_seq_len', int(args.max_seq_len))
                    except Exception:
                        pass
                    if args.query:
                        result = tmp.compress_prompt(
                            input_text,
                            rate=rate,
                            question_list=[args.query],
                            condition_in_question='after_condition'
                        )
                    else:
                        result = tmp.compress_prompt(input_text, rate=rate)
            else:
                # v1 expects a list of contexts and optional question/target_token
                target_token = -1
                if args.target_tokens is not None:
                    target_token = int(max(1, args.target_tokens))
                else:
                    # derive target from rate and tokenized length for v1
                    orig_tokens = len(get_encoder().encode(input_text))
                    target_token = int(max(1, round(rate * max(1, orig_tokens))))
                result = compressor.compress_prompt(
                    [input_text],
                    question=(args.query or ''),
                    rate=rate,
                    target_token=target_token,
                    use_context_level_filter=True,
                    use_token_level_filter=True,
                    strict_preserve_uncompressed=False,
                )
                # v1 may return list field; ensure we extract the string
            compressed = result['compressed_prompt']

            # If v1 produced no reduction, fallback to v2 once for meaningful compression
            if not is_v2:
                try:
                    orig_tokens = len(get_encoder().encode(input_text))
                    comp_tokens = len(get_encoder().encode(compressed)) if compressed else 0
                    if comp_tokens >= orig_tokens:
                        # Build a temporary v2 compressor with a known working model
                        torch = get_torch()
                        device_map = 'cpu'
                        if not getattr(args, 'cpu', False):
                            if torch.cuda.is_available():
                                device_map = 'cuda'
                            elif getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
                                device_map = 'mps'
                        PromptCompressor = get_prompt_compressor()
                        silence_transformers_warnings()
                        tmp = PromptCompressor(
                            model_name=PRESETS['base'],
                            use_llmlingua2=True,
                            device_map=device_map,
                        )
                        if args.query:
                            res2 = tmp.compress_prompt(
                                input_text,
                                rate=rate,
                                question_list=[args.query],
                                condition_in_question='after_condition'
                            )
                        else:
                            res2 = tmp.compress_prompt(input_text, rate=rate)
                        compressed2 = res2['compressed_prompt']
                        if len(get_encoder().encode(compressed2)) < orig_tokens:
                            compressed = compressed2
                except Exception:
                    pass

            # Write to cache
            if use_cache:
                try:
                    cache_path.write_text(json.dumps({'compressed_prompt': compressed}), encoding='utf-8')
                except Exception:
                    pass
            return compressed
    except Exception as e:
        if getattr(args, 'debug', False):
            import traceback
            traceback.print_exc()
        else:
            print(f"Compression error: {e}", file=sys.stderr)
        sys.exit(1)

def compute_metrics(original, compressed, args):
    """Calculate token counts, ratio, savings, and optional BERTScore."""
    orig_tokens = len(get_encoder().encode(original))
    comp_tokens = len(get_encoder().encode(compressed)) if compressed else 0
    ratio = orig_tokens / comp_tokens if comp_tokens > 0 else 0
    per_run_savings = (orig_tokens - comp_tokens) * args.price_per_1k / 1000 if args.cost_savings else 0
    stats = {
        'original_tokens': orig_tokens,
        'compressed_tokens': comp_tokens,
        'compression_ratio': round(ratio, 2),
        'per_run_savings_usd': round(per_run_savings, 4) if args.cost_savings else None
    }
    if args.cost_savings:
        stats_data = load_stats()
        stats['total_savings_usd'] = round(stats_data['total_savings'] + per_run_savings, 4)
        stats['run_count'] = stats_data['run_count'] + 1
    if args.verbose and bertscore:
        try:
            bs = bertscore.compute(predictions=[compressed], references=[original], lang='en')
            stats['bertscore_f1'] = round(bs['f1'][0], 3)
        except Exception as e:
            print(f"BERTScore error: {e}", file=sys.stderr)
    return stats, per_run_savings

def run_server(args):
    # Ensure models are loaded once
    load_models(args)

    class Handler(BaseHTTPRequestHandler):
        def _send(self, code, obj):
            body = json.dumps(obj).encode('utf-8')
            self.send_response(code)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_POST(self):
            if self.path != '/compress':
                return self._send(404, {'error': 'not found'})
            try:
                length = int(self.headers.get('Content-Length', '0'))
                payload = json.loads(self.rfile.read(length).decode('utf-8')) if length > 0 else {}
                text = str(payload.get('text', '')).strip()
                if not text:
                    return self._send(400, {'error': 'empty text'})

                # Derive an args-like namespace for compression
                local_args = SimpleNamespace(**{
                    'ratio': float(payload.get('ratio')) if payload.get('ratio') is not None else args.ratio,
                    'target_tokens': payload.get('target_tokens'),
                    'query': payload.get('query'),
                    'format': payload.get('format') or 'text',
                    'verbose': False,
                    'cost_savings': False,
                    'price_per_1k': args.price_per_1k,
                })

                compressed = compress_prompt(text, local_args)
                out = {'compressed_prompt': compressed}
                return self._send(200, out)
            except Exception as e:
                return self._send(500, {'error': f'{e}'})

    class ReusableThreadingHTTPServer(ThreadingHTTPServer):
        allow_reuse_address = True

    try:
        httpd = ReusableThreadingHTTPServer(('0.0.0.0', args.port), Handler)
    except OSError as e:
        print(f"Server error: {e}. If the port is in use, try --port with a different value.", file=sys.stderr)
        sys.exit(2)
    print(f"tinyprompt server listening on http://0.0.0.0:{args.port}", file=sys.stderr)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        httpd.server_close()
    return 0

def main():
    parser = argparse.ArgumentParser(
        prog='tinyprompt',
        description='Semantic prompt compressor for LLMs. Reads from stdin if piped.',
        epilog='Pipe-friendly: cat prompt.txt | tinyprompt | llm -m gpt-4o'
    )
    parser.add_argument('input', nargs='?', help='Input text or file path (use - for stdin)')
    parser.add_argument('--input', dest='input_flag', help='Input text or file path (use - for stdin)')
    parser.add_argument('--ratio', type=float, default=0.7, help='Compression ratio (0.05-0.9, default 0.7; override via TINYPROMPT_RATIO)')
    parser.add_argument('--target-tokens', type=int, dest='target_tokens', help='Desired compressed token count (alternative to --ratio)')
    parser.add_argument('--query', type=str, help='Optional query for relevance boosting')
    parser.add_argument('--model', type=str, default='microsoft/llmlingua-2-xlm-roberta-large-meetingbank', help='Model identifier (HF id)')
    parser.add_argument('--preset', type=str, choices=list(PRESETS.keys()), help='Model preset (maps to a known HF id)')
    parser.add_argument('--algo', type=str, choices=['v1','v2'], help='Compression algorithm: v1 (faster), v2 (default)')
    parser.add_argument('--serve', action='store_true', help='Run as a local HTTP server to keep model warm')
    parser.add_argument('--port', type=int, default=int(os.environ.get('TINYPROMPT_PORT', '8008')), help='Port for --serve (default 8008)')
    parser.add_argument('--server-url', type=str, dest='server_url', help='Forward request to a running server (e.g., http://127.0.0.1:8008)')
    parser.add_argument('--threads', type=int, help='Number of CPU threads to use')
    parser.add_argument('--cache-dir', type=str, dest='cache_dir', help='Hugging Face cache directory')
    parser.add_argument('--no-cache', action='store_true', help='Disable tinyprompt result caching')
    parser.add_argument('--max-seq-len', type=int, dest='max_seq_len', help='Limit max sequence length for faster processing')
    parser.add_argument('--cpu', action='store_true', help='Force CPU even if CUDA/MPS available')
    parser.add_argument('--fast', action='store_true', help='Use faster defaults (algo v1, lower force tokens)')
    parser.add_argument('--max-batch-size', type=int, dest='max_batch_size', help='llmlingua2: limit batch size for speed/memory')
    parser.add_argument('--max-force-token', type=int, dest='max_force_token', help='llmlingua2: reduce force-preserved tokens for speed')
    parser.add_argument('--offline', action='store_true', help='Use local cache only; do not download models')
    parser.add_argument('--debug', action='store_true', help='Print full tracebacks on errors')
    parser.add_argument('--verbose', action='store_true', help='Output metrics to stderr')
    parser.add_argument('--format', type=str, default='text', choices=['text', 'json'], help='Output: text or json')
    parser.add_argument('--cost-savings', action='store_true', help='Show per-run and total cost savings in USD')
    parser.add_argument('--price-per-1k', type=float, default=0.01, help='API price per 1k input tokens (default 0.01)')
    parser.add_argument('--reset-stats', action='store_true', help='Reset cumulative savings counter')

    args = parser.parse_args()

    # Resolve model based on CLI, env, and preset
    env_model = os.environ.get('TINYPROMPT_MODEL')
    if env_model:
        args.model = env_model
    elif args.preset:
        args.model = PRESETS.get(args.preset, args.model)

    # Ratio env override
    env_ratio = os.environ.get('TINYPROMPT_RATIO')
    if env_ratio:
        try:
            args.ratio = float(env_ratio)
        except Exception:
            print("Error: TINYPROMPT_RATIO must be a number", file=sys.stderr)
            sys.exit(1)

    # Configure CPU threading if requested
    if args.threads and args.threads > 0:
        try:
            torch.set_num_threads(args.threads)
        except Exception:
            pass
        os.environ['OMP_NUM_THREADS'] = str(args.threads)
        os.environ['MKL_NUM_THREADS'] = str(args.threads)
        os.environ['NUMEXPR_NUM_THREADS'] = str(args.threads)
        os.environ['TOKENIZERS_PARALLELISM'] = 'false'

    # Configure cache and offline mode
    if args.cache_dir:
        try:
            os.makedirs(args.cache_dir, exist_ok=True)
        except Exception:
            pass
        os.environ['HF_HOME'] = args.cache_dir
        os.environ['TRANSFORMERS_CACHE'] = args.cache_dir
    if args.offline:
        os.environ['HF_HUB_OFFLINE'] = '1'
        os.environ['TRANSFORMERS_OFFLINE'] = '1'

    # Fast path defaults (apply BEFORE possible server start)
    if args.fast:
        # Keep current algorithm unless explicitly set; just tune for speed
        args.max_force_token = args.max_force_token or 50
        if args.max_seq_len is None:
            args.max_seq_len = 256
    # Ensure model is compatible with selected algorithm
    if (args.algo == 'v1') and not env_model:
        model_lower = (args.model or '').lower()
        if 'llmlingua-2' in model_lower or 'roberta' in model_lower or 'xlm' in model_lower:
            # Use a small causal LM for v1 path
            args.model = 'distilgpt2'

    # If running in server mode, start server (keeps process warm)
    if args.serve:
        return run_server(args)

    # Validate args (ratio bounds apply only when --target-tokens not given)
    if args.target_tokens is None and (args.ratio < 0.05 or args.ratio > 0.9):
        print("Error: Ratio must be between 0.05 and 0.9", file=sys.stderr)
        sys.exit(1)
    if args.price_per_1k <= 0:
        print("Error: Price per 1k tokens must be positive", file=sys.stderr)
        sys.exit(1)

    # Reset stats if requested
    if args.reset_stats:
        try:
            if os.path.exists(STATS_FILE):
                os.remove(STATS_FILE)
            print("Stats reset", file=sys.stderr)
        except Exception as e:
            print(f"Error resetting stats: {e}", file=sys.stderr)
            sys.exit(1)

    # Read input (supports stdin, file path, or inline text via --input or positional)
    input_text = ''
    try:
        provided = args.input_flag if getattr(args, 'input_flag', None) is not None else args.input
        if provided == '-' or not provided:
            if not sys.stdin.isatty():
                input_text = sys.stdin.read().strip()
            else:
                parser.print_help()
                sys.exit(1)
        else:
            if os.path.exists(provided):
                with open(provided, 'r', encoding='utf-8') as f:
                    input_text = f.read().strip()
            else:
                # Treat provided value as literal input text
                input_text = str(provided).strip()
    except Exception as e:
        print(f"Input error: {e}", file=sys.stderr)
        sys.exit(1)

    if not input_text:
        print("Error: Empty input", file=sys.stderr)
        sys.exit(1)

    # If client forwarding is requested, forward payload to server
    if args.server_url:
        try:
            payload = {
                'text': input_text,
                'ratio': args.ratio,
                'target_tokens': args.target_tokens,
                'query': args.query,
                'format': args.format,
            }
            data = json.dumps(payload).encode('utf-8')
            req = urllib.request.Request(args.server_url.rstrip('/') + '/compress', data=data, headers={'Content-Type': 'application/json'})
            with urllib.request.urlopen(req, timeout=180) as resp:
                body = resp.read().decode('utf-8')
                if args.format == 'json':
                    print(body)
                else:
                    # Expecting JSON {compressed_prompt: ...}
                    obj = json.loads(body)
                    print(obj.get('compressed_prompt', ''))
            return
        except Exception as e:
            print(f"Client error: {e}", file=sys.stderr)
            sys.exit(1)

    # Compress
    try:
        load_models(args)
        chunk_size = 50000
        if len(input_text) > chunk_size:
            chunks = [input_text[i:i+chunk_size] for i in range(0, len(input_text), chunk_size)]
            compressed_chunks = [compress_prompt(chunk, args) for chunk in chunks]
            compressed = ' '.join(compressed_chunks)
        else:
            compressed = compress_prompt(input_text, args)

        # Metrics and savings
        stats, per_run_savings = compute_metrics(input_text, compressed, args)
        if args.cost_savings or args.verbose:
            print(json.dumps(stats), file=sys.stderr)
        if args.cost_savings:
            save_stats(load_stats(), per_run_savings)

        # Output
        if args.format == 'json':
            output = {'compressed_prompt': compressed}
            if args.verbose or args.cost_savings:
                output['stats'] = stats
            print(json.dumps(output))
        else:
            print(compressed)
    except Exception as e:
        print(f"Output error: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == '__main__':
    main()