flense 0.1.0__py3-none-any.whl

flense/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

flense/app.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from .config import FlenseConfig
+from .providers import get_adapter, register_providers
+from .proxy import proxy_request
+from .telemetry import SessionStats
+
+
+@asynccontextmanager
+async def _lifespan(app: FastAPI):
+    app.state.httpx_client = httpx.AsyncClient(
+        timeout=httpx.Timeout(
+            connect=10.0,
+            read=300.0,   # AI responses can be slow
+            write=10.0,
+            pool=10.0,
+        ),
+        follow_redirects=False,
+        http2=True,
+    )
+    yield
+    await app.state.httpx_client.aclose()
+
+
+def create_app(config: FlenseConfig) -> FastAPI:
+    app = FastAPI(title="flense", lifespan=_lifespan)
+    app.state.config = config
+    app.state.session_stats = SessionStats()
+    register_providers(config)
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    @app.api_route(
+        "/{provider}/{path:path}",
+        methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+    )
+    async def proxy_route(provider: str, path: str, request: Request):
+        adapter = get_adapter(provider)
+        if adapter is None:
+            return JSONResponse(
+                {"error": f"Unknown provider: {provider}"},
+                status_code=404,
+            )
+        return await proxy_request(request, adapter)
+
+    return app

flense/cli.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+app = typer.Typer(
+    name="flense",
+    help="Lightweight reverse proxy that compresses AI API payloads.",
+    add_completion=False,
+)
+
+
+@app.command()
+def start(
+    port: Optional[int] = typer.Option(None, help="Override server port"),
+    foreground: bool = typer.Option(
+        False, "--fg", help="Run in foreground (no daemon)"
+    ),
+    headless: Optional[bool] = typer.Option(
+        None, help="Headless mode (JSON logs, no TUI)"
+    ),
+    config: Optional[Path] = typer.Option(
+        None, help="Path to flense.toml"
+    ),
+) -> None:
+    """Start the flense proxy."""
+    from .config import load_config
+
+    cfg = load_config(config)
+    effective_port = port if port is not None else cfg.server.port
+
+    # CLI flag overrides config file
+    use_headless = headless if headless is not None else cfg.server.headless
+
+    if foreground:
+        if use_headless:
+            _run_headless(cfg, effective_port)
+        else:
+            _run_tui(cfg, effective_port)
+    else:
+        from .daemon import is_running, start_daemon
+
+        if is_running():
+            from .daemon import read_pid
+
+            typer.echo(f"flense is already running (PID {read_pid()})")
+            raise typer.Exit(1)
+
+        pid = start_daemon(cfg.server.host, effective_port, config)
+        typer.echo(f"flense started on {cfg.server.host}:{effective_port} (PID {pid})")
+
+
+def _run_headless(cfg, port: int) -> None:
+    """Run the proxy in foreground with structured JSON logs (no TUI)."""
+    import uvicorn
+
+    from .app import create_app
+
+    cfg.server.headless = True
+    instance = create_app(cfg)
+    uvicorn.run(instance, host=cfg.server.host, port=port)
+
+
+def _run_tui(cfg, port: int) -> None:
+    """Run the proxy in foreground with the Textual TUI dashboard."""
+    import threading
+
+    import uvicorn
+
+    from .app import create_app
+    from .tui import FlenseDashboard
+
+    instance = create_app(cfg)
+    stats = instance.state.session_stats
+
+    # Run uvicorn in a background thread so Textual owns the main thread.
+    server = uvicorn.Server(
+        uvicorn.Config(
+            instance,
+            host=cfg.server.host,
+            port=port,
+            log_level="warning",
+        )
+    )
+    thread = threading.Thread(target=server.run, daemon=True)
+    thread.start()
+
+    dashboard = FlenseDashboard(
+        session_stats=stats,
+        server_host=cfg.server.host,
+        server_port=port,
+    )
+    dashboard.run()
+
+    # TUI exited, shut down the server.
+    server.should_exit = True
+    thread.join(timeout=5)
+
+
+@app.command()
+def stop() -> None:
+    """Stop the flense proxy."""
+    from .daemon import stop_daemon
+
+    try:
+        stop_daemon()
+    except RuntimeError as exc:
+        typer.echo(str(exc))
+        raise typer.Exit(1)
+
+    typer.echo("flense stopped.")
+
+
+@app.command()
+def status() -> None:
+    """Check whether flense is running."""
+    from .daemon import is_running, read_pid
+
+    if is_running():
+        typer.echo(f"flense is running (PID {read_pid()})")
+    else:
+        typer.echo("flense is not running.")

flense/code_writer.py

@@ -0,0 +1,230 @@
+"""Code-writer bypass: route code generation tasks to a cheap model.
+
+Triggered by X-Flense-Strategy: code-writer. Sends the request to a cheap
+cloud model, writes the output to disk, and returns a synthetic ACK to the
+client. The frontier model is never called.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import time
+from pathlib import Path
+
+import httpx
+
+from .config import CodeWriterConfig, FlenseConfig
+
+logger = logging.getLogger(__name__)
+
+_ANTHROPIC_UPSTREAM = "https://api.anthropic.com"
+_OPENAI_UPSTREAM = "https://api.openai.com"
+
+
+def _infer_provider(model: str) -> str:
+    """Infer provider from model name prefix."""
+    if model.startswith("claude"):
+        return "anthropic"
+    return "openai"
+
+
+def _get_api_key(provider: str, headers: dict[str, str]) -> str | None:
+    """Extract API key for the given provider from request headers, then env vars."""
+    if provider == "anthropic":
+        return headers.get("x-api-key") or os.environ.get("ANTHROPIC_API_KEY")
+    if provider == "openai":
+        auth = headers.get("authorization", "")
+        key = auth.removeprefix("Bearer ").strip() or None
+        return key or os.environ.get("OPENAI_API_KEY")
+    return None
+
+
+def _extract_spec(body: bytes) -> str:
+    """Extract the spec from the last user message in the request body."""
+    try:
+        payload = json.loads(body)
+        messages = payload.get("messages", [])
+        for msg in reversed(messages):
+            if msg.get("role") == "user":
+                content = msg.get("content", "")
+                if isinstance(content, str):
+                    return content
+                if isinstance(content, list):
+                    parts = [
+                        block.get("text", "")
+                        for block in content
+                        if isinstance(block, dict) and block.get("type") == "text"
+                    ]
+                    return "\n".join(parts)
+    except Exception:
+        pass
+    return ""
+
+
+def _build_prompt(spec: str, ref_file_path: str, ref_content: str) -> str:
+    return (
+        f"You are a code generator. Write code following the exact style and "
+        f"conventions shown in the reference file.\n\n"
+        f"Reference file ({ref_file_path}):\n"
+        f"```\n{ref_content}\n```\n\n"
+        f"Task: {spec}\n\n"
+        f"Output only the raw code. No markdown formatting, no explanations, no preamble."
+    )
+
+
+def _resolve_output_path(
+    output_file: str | None,
+    ref_file_path: str | None,
+    output_dir: str,
+) -> Path:
+    if output_file:
+        p = Path(output_file)
+        if not p.is_absolute():
+            p = Path(output_dir) / p
+        return p
+
+    if ref_file_path:
+        ref = Path(ref_file_path)
+        suffix = ref.suffix or ".txt"
+        timestamp = int(time.time())
+        return Path(output_dir) / f"{ref.stem}_generated_{timestamp}{suffix}"
+
+    timestamp = int(time.time())
+    return Path(output_dir) / f"generated_{timestamp}.txt"
+
+
+async def _call_anthropic(
+    model: str,
+    prompt: str,
+    api_key: str,
+    client: httpx.AsyncClient,
+    upstream: str,
+) -> str:
+    resp = await client.post(
+        f"{upstream}/v1/messages",
+        headers={
+            "x-api-key": api_key,
+            "anthropic-version": "2023-06-01",
+            "content-type": "application/json",
+        },
+        content=json.dumps({
+            "model": model,
+            "max_tokens": 8192,
+            "messages": [{"role": "user", "content": prompt}],
+        }).encode(),
+        timeout=120.0,
+    )
+    resp.raise_for_status()
+    data = resp.json()
+    return data["content"][0]["text"]
+
+
+async def _call_openai(
+    model: str,
+    prompt: str,
+    api_key: str,
+    client: httpx.AsyncClient,
+    upstream: str,
+) -> str:
+    resp = await client.post(
+        f"{upstream}/v1/chat/completions",
+        headers={
+            "authorization": f"Bearer {api_key}",
+            "content-type": "application/json",
+        },
+        content=json.dumps({
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+        }).encode(),
+        timeout=120.0,
+    )
+    resp.raise_for_status()
+    data = resp.json()
+    return data["choices"][0]["message"]["content"]
+
+
+async def _call_model(
+    model: str,
+    prompt: str,
+    headers: dict[str, str],
+    client: httpx.AsyncClient,
+    config: FlenseConfig,
+) -> str:
+    provider = _infer_provider(model)
+    api_key = _get_api_key(provider, headers)
+    if not api_key:
+        raise ValueError(
+            f"No API key found for provider '{provider}' (model '{model}'). "
+            f"Pass the key via request headers or set the appropriate env var."
+        )
+
+    provider_config = config.providers.get(provider)
+    if provider == "anthropic":
+        upstream = (
+            provider_config.upstream
+            if provider_config and provider_config.upstream
+            else _ANTHROPIC_UPSTREAM
+        )
+        return await _call_anthropic(model, prompt, api_key, client, upstream)
+
+    upstream = (
+        provider_config.upstream
+        if provider_config and provider_config.upstream
+        else _OPENAI_UPSTREAM
+    )
+    return await _call_openai(model, prompt, api_key, client, upstream)
+
+
+async def handle_code_writer(
+    body: bytes,
+    headers: dict[str, str],
+    ref_file_header: str | None,
+    output_file_header: str | None,
+    cw_config: CodeWriterConfig,
+    config: FlenseConfig,
+    httpx_client: httpx.AsyncClient,
+) -> tuple[str, str]:
+    """Execute the code-writer bypass.
+
+    Returns (generated_code, output_path_str). Raises on error.
+    """
+    if not ref_file_header:
+        raise ValueError(
+            "X-Flense-Ref-File header is required for the code-writer strategy"
+        )
+
+    ref_path = Path(ref_file_header)
+    if not ref_path.is_file():
+        raise FileNotFoundError(f"Reference file not found: {ref_file_header}")
+    ref_content = ref_path.read_text(errors="replace")
+
+    spec = _extract_spec(body)
+    if not spec:
+        raise ValueError("Could not extract a spec from the request messages")
+
+    prompt = _build_prompt(spec, ref_file_header, ref_content)
+
+    generated_code: str | None = None
+    last_error: Exception | None = None
+    for model in [cw_config.model, cw_config.fallback]:
+        if not model:
+            continue
+        try:
+            generated_code = await _call_model(model, prompt, headers, httpx_client, config)
+            logger.info("Code-writer used model %s", model)
+            break
+        except Exception as exc:
+            logger.warning("Code-writer model %s failed: %s", model, exc)
+            last_error = exc
+
+    if generated_code is None:
+        raise RuntimeError(f"All code-writer models failed. Last error: {last_error}")
+
+    output_path = _resolve_output_path(output_file_header, ref_file_header, cw_config.output_dir)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(generated_code)
+    logger.info("Code-writer wrote output to %s", output_path)
+
+    return generated_code, str(output_path)

flense/compression/__init__.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import copy
+import json
+import logging
+import re
+import time
+
+from .classifier import detect_signals, resolve_strategy
+from .ctags import compress_with_ctags
+from .estimator import estimate_messages_tokens, estimate_tokens
+from .language import detect_language
+from .regex import compress_with_regex
+from .treesitter import compress_with_treesitter
+from .types import CodeBlock, CompressedBlock, CompressionResult, Strategy
+
+logger = logging.getLogger(__name__)
+
+_CODE_BLOCK_RE = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL)
+_FILE_HINT_RE = re.compile(r"(?:File|Path|file|path):\s*(\S+)")
+
+
+def compress_payload(
+    body: bytes,
+    provider: str,
+    threshold: int = 5000,
+    header_strategy: str | None = None,
+    provider_strategy: str | None = None,
+    global_strategy: str = "auto",
+) -> CompressionResult:
+    """Main entry point for the compression pipeline.
+
+    Takes the raw request body (bytes) and returns a CompressionResult
+    containing the (possibly compressed) body and metadata.
+    """
+    start_time = time.monotonic()
+
+    try:
+        payload = json.loads(body)
+    except (json.JSONDecodeError, UnicodeDecodeError):
+        return _passthrough(body, start_time)
+
+    if not isinstance(payload, dict):
+        return _passthrough(body, start_time)
+
+    messages = payload.get("messages", [])
+    if not messages:
+        return _passthrough(body, start_time)
+
+    system = payload.get("system", "")
+    if isinstance(system, list):
+        system = " ".join(
+            b.get("text", "") for b in system if isinstance(b, dict)
+        )
+
+    # Token estimation
+    tokens_before = estimate_messages_tokens(messages, provider)
+    if system:
+        tokens_before += estimate_tokens(system, provider)
+
+    if tokens_before < threshold:
+        return _passthrough(body, start_time, tokens_before=tokens_before)
+
+    # Classification
+    signals = detect_signals(messages, system if system else None)
+    strategy = resolve_strategy(
+        header_strategy, provider_strategy, global_strategy, signals,
+    )
+
+    if strategy == Strategy.PASSTHROUGH:
+        return _passthrough(body, start_time, tokens_before=tokens_before)
+
+    # Extract code blocks
+    code_blocks = _extract_code_blocks(messages)
+    if not code_blocks:
+        return _passthrough(body, start_time, tokens_before=tokens_before)
+
+    # Compress each block through the fallback chain
+    compressed_blocks: list[CompressedBlock] = []
+    for block in code_blocks:
+        result = _compress_block(block, strategy)
+        if result is not None:
+            compressed_blocks.append(result)
+
+    if not compressed_blocks:
+        return _passthrough(body, start_time, tokens_before=tokens_before)
+
+    # Reconstruct the payload with compressed blocks
+    compressed_payload = _reconstruct_payload(payload, compressed_blocks)
+    compressed_body = json.dumps(compressed_payload).encode("utf-8")
+
+    # Estimate tokens after compression
+    compressed_messages = compressed_payload.get("messages", [])
+    tokens_after = estimate_messages_tokens(compressed_messages, provider)
+    if system:
+        tokens_after += estimate_tokens(system, provider)
+
+    elapsed_ms = (time.monotonic() - start_time) * 1000
+
+    return CompressionResult(
+        original_body=body,
+        compressed_body=compressed_body,
+        strategy_applied=strategy,
+        tokens_before=tokens_before,
+        tokens_after=tokens_after,
+        blocks_compressed=len(compressed_blocks),
+        blocks_total=len(code_blocks),
+        compression_time_ms=elapsed_ms,
+        was_compressed=True,
+    )
+
+
+def _compress_block(block: CodeBlock, strategy: Strategy) -> CompressedBlock | None:
+    """Apply the fallback chain: tree-sitter -> ctags -> regex."""
+    if strategy in (Strategy.AST, Strategy.AUTO):
+        result = compress_with_treesitter(block)
+        if result is not None:
+            return result
+
+    if strategy in (Strategy.CTAGS, Strategy.AST, Strategy.AUTO):
+        result = compress_with_ctags(block)
+        if result is not None:
+            return result
+
+    # Regex as last resort for any non-passthrough strategy
+    result = compress_with_regex(block)
+    if result is not None:
+        return result
+
+    return None
+
+
+def _extract_code_blocks(messages: list[dict]) -> list[CodeBlock]:
+    """Extract fenced code blocks from all messages."""
+    blocks: list[CodeBlock] = []
+    for msg in messages:
+        text_parts = _get_text_parts(msg)
+        for text in text_parts:
+            for match in _CODE_BLOCK_RE.finditer(text):
+                fence_tag = match.group(1) or None
+                content = match.group(2)
+
+                # Look for a file path hint in the ~200 chars before the fence
+                pre_text = text[max(0, match.start() - 200):match.start()]
+                path_match = _FILE_HINT_RE.search(pre_text)
+                source_path = path_match.group(1) if path_match else None
+
+                lang = detect_language(
+                    content, fence_tag=fence_tag, file_path=source_path,
+                )
+                blocks.append(CodeBlock(
+                    content=content,
+                    language=lang,
+                    fence_tag=fence_tag,
+                    source_path=source_path,
+                ))
+    return blocks
+
+
+def _get_text_parts(msg: dict) -> list[str]:
+    """Extract all text content from a message."""
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return [content]
+    parts: list[str] = []
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict):
+                if block.get("text"):
+                    parts.append(block["text"])
+                inner = block.get("content", "")
+                if isinstance(inner, str) and inner:
+                    parts.append(inner)
+                elif isinstance(inner, list):
+                    for sub in inner:
+                        if isinstance(sub, dict) and sub.get("text"):
+                            parts.append(sub["text"])
+    return parts
+
+
+def _reconstruct_payload(
+    payload: dict,
+    compressed_blocks: list[CompressedBlock],
+) -> dict:
+    """Replace original code blocks in the payload with compressed versions."""
+    replacements: dict[str, str] = {}
+    for cb in compressed_blocks:
+        tag = cb.original.fence_tag or ""
+        original_fenced = f"```{tag}\n{cb.original.content}```"
+        compressed_fenced = (
+            f"```{tag}\n"
+            f"# [flense: compressed with {cb.method}]\n"
+            f"{cb.compressed}\n```"
+        )
+        replacements[original_fenced] = compressed_fenced
+
+    new_payload = copy.deepcopy(payload)
+    for msg in new_payload.get("messages", []):
+        _apply_replacements(msg, replacements)
+
+    return new_payload
+
+
+def _apply_replacements(msg: dict, replacements: dict[str, str]) -> None:
+    """Apply text replacements to a message's content in place."""
+    content = msg.get("content")
+    if isinstance(content, str):
+        for original, compressed in replacements.items():
+            content = content.replace(original, compressed)
+        msg["content"] = content
+    elif isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict):
+                if "text" in block and isinstance(block["text"], str):
+                    for original, compressed in replacements.items():
+                        block["text"] = block["text"].replace(original, compressed)
+                inner = block.get("content")
+                if isinstance(inner, str):
+                    for original, compressed in replacements.items():
+                        inner = inner.replace(original, compressed)
+                    block["content"] = inner
+                elif isinstance(inner, list):
+                    for sub in inner:
+                        if isinstance(sub, dict) and "text" in sub:
+                            for original, compressed in replacements.items():
+                                sub["text"] = sub["text"].replace(original, compressed)
+
+
+def _passthrough(
+    body: bytes,
+    start_time: float,
+    tokens_before: int = 0,
+) -> CompressionResult:
+    return CompressionResult(
+        original_body=body,
+        compressed_body=body,
+        strategy_applied=Strategy.PASSTHROUGH,
+        tokens_before=tokens_before,
+        tokens_after=tokens_before,
+        blocks_compressed=0,
+        blocks_total=0,
+        compression_time_ms=(time.monotonic() - start_time) * 1000,
+        was_compressed=False,
+    )