prompt-flamegraph 0.2.0__py3-none-any.whl

prompt_flamegraph/__init__.py

@@ -0,0 +1,32 @@
+# prompt-flamegraph - Lightweight prompt context flamegraph generator for LLMs.
+# Copyright (C) 2025  fjjjuv
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""prompt_flamegraph: lightweight prompt context flamegraph generator for LLMs."""
+
+from .core import build_tree, count_tokens, profile_prompt
+from .diff import diff_prompts
+from .waste import detect_waste, WasteReport, Finding
+
+__version__ = "0.2.0"
+__all__ = [
+    "build_tree",
+    "count_tokens",
+    "diff_prompts",
+    "detect_waste",
+    "Finding",
+    "profile_prompt",
+    "WasteReport",
+]

prompt_flamegraph/__main__.py

@@ -0,0 +1,20 @@
+# prompt-flamegraph - Lightweight prompt context flamegraph generator for LLMs.
+# Copyright (C) 2025  fjjjuv
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from .cli import main
+import sys
+
+sys.exit(main())

prompt_flamegraph/cli.py

@@ -0,0 +1,233 @@
+# prompt-flamegraph - Lightweight prompt context flamegraph generator for LLMs.
+# Copyright (C) 2025  fjjjuv
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Command-line interface for prompt-flamegraph."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+SAMPLE_PROMPT = {
+    "system_prompt": (
+        "You are a helpful coding assistant. Be concise. "
+        "Always think step by step before answering."
+    ),
+    "tools": [
+        {
+            "name": "read_file",
+            "description": "Read a file from disk. Accepts a path argument.",
+            "schema": "{\"path\": \"string\"}",
+        },
+        {
+            "name": "run_command",
+            "description": "Execute a shell command and return stdout/stderr.",
+            "schema": "{\"command\": \"string\"}",
+        },
+    ],
+    "rag_context": {
+        "doc_1.py": "import os\n\ndef hello():\n    return 'hello'\n",
+        "doc_2.py": "import sys\n\ndef world():\n    return 'world'\n",
+    },
+    "chat_history": [
+        {"role": "user", "content": "Help me profile my prompt tokens."},
+        {"role": "assistant", "content": "Sure, I can help you understand where your tokens go."},
+    ],
+}
+
+
+def _load_input(value: str) -> dict:
+    path = Path(value)
+    if path.exists():
+        raw = path.read_text(encoding="utf-8")
+    else:
+        raw = value
+    try:
+        data = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise SystemExit(f"Invalid JSON input: {exc}") from exc
+    if not isinstance(data, (dict, list)):
+        raise SystemExit("Input JSON must be a dict or a list.")
+    return data
+
+
+def _detect_output(args: argparse.Namespace) -> str:
+    if args.output:
+        return args.output
+    ext = {"html": ".html", "svg": ".svg", "md": ".md"}.get(args.format, ".html")
+    return f"prompt_flamegraph{ext}"
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="prompt-flamegraph",
+        description="Generate a lightweight flamegraph of your LLM prompt tokens.",
+    )
+    parser.add_argument(
+        "input",
+        nargs="?",
+        help="JSON file (or JSON string) representing the prompt structure.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="Output file (inferred from --format if not given).",
+    )
+    parser.add_argument(
+        "-t",
+        "--title",
+        default=None,
+        help="Title of the output.",
+    )
+    parser.add_argument(
+        "--format",
+        choices=["html", "svg", "md"],
+        default="html",
+        help="Output format (default: html).",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        default=None,
+        help="Tokenizer to use: 'tiktoken', 'words', or a callable.",
+    )
+    parser.add_argument(
+        "--cost",
+        type=float,
+        default=None,
+        help="Cost per token (e.g. 1.5e-6 for $1.5 per million tokens).",
+    )
+    parser.add_argument(
+        "--diff",
+        metavar="FILE",
+        default=None,
+        help="Compare INPUT with another JSON file and output a diff flamegraph.",
+    )
+    parser.add_argument(
+        "--terminal",
+        action="store_true",
+        help="Print a terminal bar chart instead of writing a file.",
+    )
+    parser.add_argument(
+        "--no-waste",
+        action="store_true",
+        help="Disable waste detection for HTML output.",
+    )
+    parser.add_argument(
+        "--demo",
+        action="store_true",
+        help="Use the built-in sample prompt.",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=1200,
+        help="Graph width in pixels (SVG/HTML).",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=720,
+        help="Graph max height in pixels (HTML).",
+    )
+
+    args = parser.parse_args(argv)
+
+    if args.demo:
+        data = SAMPLE_PROMPT
+    elif args.input:
+        data = _load_input(args.input)
+    else:
+        parser.print_help()
+        return 1
+
+    output = _detect_output(args)
+
+    if args.terminal:
+        from .core import build_tree
+        from .terminal import to_terminal
+
+        tree = build_tree(data, tokenizer=args.tokenizer)
+        to_terminal(tree, title=args.title, cost_per_token=args.cost)
+        return 0
+
+    if args.diff:
+        from .core import build_tree
+        from .diff import build_diff_tree
+        from .render import to_html
+        from .export import to_svg, to_markdown
+
+        v2 = _load_input(args.diff)
+        t1 = build_tree(data, tokenizer=args.tokenizer)
+        t2 = build_tree(v2, tokenizer=args.tokenizer)
+        diff_tree = build_diff_tree(t1, t2)
+
+        if args.format == "html":
+            html = to_html(diff_tree, title=args.title or "Prompt Diff", cost_per_token=args.cost, width=args.width, height=args.height)
+            with open(output, "w", encoding="utf-8") as f:
+                f.write(html)
+        elif args.format == "svg":
+            from .export import to_svg
+            svg = to_svg(diff_tree, title=args.title or "Prompt Diff", width=args.width)
+            with open(output, "w", encoding="utf-8") as f:
+                f.write(svg)
+        elif args.format == "md":
+            from .export import to_markdown
+            md = to_markdown(diff_tree, title=args.title or "Prompt Diff", cost_per_token=args.cost)
+            with open(output, "w", encoding="utf-8") as f:
+                f.write(md)
+
+        print(f"Diff written to: {output}")
+        return 0
+
+    if args.format == "html":
+        from .core import profile_prompt
+
+        profile_prompt(
+            data,
+            output=output,
+            title=args.title,
+            tokenizer=args.tokenizer,
+            cost_per_token=args.cost,
+            detect_waste=not args.no_waste,
+            width=args.width,
+            height=args.height,
+        )
+    elif args.format == "svg":
+        from .core import build_tree
+        from .export import to_svg
+
+        tree = build_tree(data, tokenizer=args.tokenizer)
+        svg = to_svg(tree, title=args.title, width=args.width)
+        with open(output, "w", encoding="utf-8") as f:
+            f.write(svg)
+    elif args.format == "md":
+        from .core import build_tree
+        from .export import to_markdown
+
+        tree = build_tree(data, tokenizer=args.tokenizer)
+        md = to_markdown(tree, title=args.title, cost_per_token=args.cost)
+        with open(output, "w", encoding="utf-8") as f:
+            f.write(md)
+
+    print(f"Flamegraph written to: {output}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

prompt_flamegraph/core.py

@@ -0,0 +1,180 @@
+# prompt-flamegraph - Lightweight prompt context flamegraph generator for LLMs.
+# Copyright (C) 2025  fjjjuv
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Core logic: token counting and tree building."""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterable
+
+
+@dataclass
+class Node:
+    name: str
+    tokens: int
+    children: list[Node] = field(default_factory=list)
+    text: str | None = None
+    change: str | None = None  # 'added', 'removed', 'same', 'changed'
+    delta: int = 0
+
+    @property
+    def is_leaf(self) -> bool:
+        return not self.children
+
+
+Tokenizer = Callable[[str], int]
+
+
+def _default_count(text: str) -> int:
+    """Fallback token estimator when tiktoken is not installed."""
+    text = text.strip()
+    if not text:
+        return 0
+    # Word-ish tokens plus isolated punctuation. Rough but stable and fast.
+    return len(re.findall(r"\w+|[^\w\s]", text, flags=re.UNICODE))
+
+
+def _load_tiktoken() -> Tokenizer | None:
+    try:
+        import tiktoken
+    except ImportError:
+        return None
+
+    enc = tiktoken.get_encoding("cl100k_base")
+
+    def count(text: str) -> int:
+        return len(enc.encode(text))
+
+    return count
+
+
+def get_tokenizer(
+    tokenizer: Tokenizer | str | None = None,
+) -> Tokenizer:
+    """Resolve a tokenizer from a name, callable, or default."""
+    if tokenizer is None:
+        tiktoken_count = _load_tiktoken()
+        if tiktoken_count is not None:
+            return tiktoken_count
+        return _default_count
+
+    if callable(tokenizer):
+        return tokenizer
+
+    if tokenizer == "tiktoken" or tokenizer.startswith("cl100k"):
+        tiktoken_count = _load_tiktoken()
+        if tiktoken_count is None:
+            raise ImportError(
+                "tiktoken is not installed. Run 'pip install prompt-flamegraph[tiktoken]'"
+            )
+        return tiktoken_count
+
+    if tokenizer == "words":
+        return _default_count
+
+    raise ValueError(f"Unknown tokenizer: {tokenizer}")
+
+
+def count_tokens(text: str, tokenizer: Tokenizer | str | None = None) -> int:
+    """Count tokens in a string using the selected tokenizer."""
+    return get_tokenizer(tokenizer)(text)
+
+
+def _guess_name(value: Any, index: int) -> str:
+    if isinstance(value, dict):
+        for key in ("name", "title", "id", "role", "filename", "file", "source"):
+            candidate = value.get(key)
+            if isinstance(candidate, str) and candidate.strip():
+                return candidate
+    return f"item_{index}"
+
+
+def _to_text(value: Any) -> str:
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, sort_keys=True)
+
+
+def build_tree(
+    data: Any,
+    name: str = "prompt",
+    tokenizer: Tokenizer | str | None = None,
+    _count_fn: Callable[[str], int] | None = None,
+) -> Node:
+    """Recursively build a token tree from a nested dict/list/string."""
+    count_fn = _count_fn or get_tokenizer(tokenizer)
+
+    if isinstance(data, dict):
+        children = [
+            build_tree(value, name=key, _count_fn=count_fn)
+            for key, value in data.items()
+        ]
+        return Node(name=name, tokens=sum(c.tokens for c in children), children=children)
+
+    if isinstance(data, (list, tuple)):
+        children = [
+            build_tree(value, name=_guess_name(value, i), _count_fn=count_fn)
+            for i, value in enumerate(data)
+        ]
+        return Node(name=name, tokens=sum(c.tokens for c in children), children=children)
+
+    text = _to_text(data)
+    return Node(name=name, tokens=count_fn(text), text=text)
+
+
+def flatten_tree(node: Node) -> list[Node]:
+    """Return a flat list of all nodes in depth-first order."""
+    out = [node]
+    for child in node.children:
+        out.extend(flatten_tree(child))
+    return out
+
+
+def profile_prompt(
+    data: Any,
+    output: str | None = "prompt_flamegraph.html",
+    title: str | None = None,
+    tokenizer: Tokenizer | str | None = None,
+    cost_per_token: float | None = None,
+    detect_waste: bool = True,
+    width: int = 1200,
+    height: int = 720,
+) -> str:
+    """Build a prompt token tree and render it to a standalone HTML flamegraph."""
+    from . import render
+
+    tree = build_tree(data, tokenizer=tokenizer)
+    waste_report = None
+    if detect_waste:
+        from .waste import detect_waste
+
+        waste_report = detect_waste(tree)
+
+    html = render.to_html(
+        tree,
+        title=title or "Prompt Flamegraph",
+        cost_per_token=cost_per_token,
+        waste_report=waste_report,
+        width=width,
+        height=height,
+    )
+    if output:
+        with open(output, "w", encoding="utf-8") as f:
+            f.write(html)
+    return html

prompt_flamegraph/diff.py

@@ -0,0 +1,131 @@
+# prompt-flamegraph - Lightweight prompt context flamegraph generator for LLMs.
+# Copyright (C) 2025  fjjjuv
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Compare two prompt trees and build a diff flamegraph."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .core import Node
+
+
+def _all_paths(node: Node, path: tuple[str, ...]) -> dict[tuple[str, ...], tuple[int, str | None, list[str]]]:
+    """Return a dict path -> (tokens, text, child_names) for every node."""
+    current = (*path, node.name)
+    result = {current: (node.tokens, node.text, [c.name for c in node.children])}
+    for child in node.children:
+        result.update(_all_paths(child, current))
+    return result
+
+
+def _build_subtree(
+    name: str,
+    path: tuple[str, ...],
+    v1_map: dict,
+    v2_map: dict,
+) -> Node:
+    """Recursively build the unified diff tree for a given path prefix."""
+    in_v1 = path in v1_map
+    in_v2 = path in v2_map
+
+    if in_v2:
+        tokens_v2, text_v2, children_v2 = v2_map[path]
+    else:
+        tokens_v2, text_v2, children_v2 = 0, None, []
+
+    if in_v1:
+        tokens_v1, text_v1, children_v1 = v1_map[path]
+    else:
+        tokens_v1, text_v1, children_v1 = 0, None, []
+
+    if in_v1 and in_v2:
+        if text_v1 == text_v2 and tokens_v1 == tokens_v2:
+            change = "same"
+        else:
+            change = "changed"
+        delta = tokens_v2 - tokens_v1
+    elif in_v2 and not in_v1:
+        change = "added"
+        delta = tokens_v2
+    elif in_v1 and not in_v2:
+        change = "removed"
+        delta = -tokens_v1
+    else:
+        # Should not happen
+        change = "same"
+        delta = 0
+
+    # Child union
+    child_names = sorted(set(children_v1) | set(children_v2))
+    children: list[Node] = []
+    for child_name in child_names:
+        child_path = (*path, child_name)
+        children.append(_build_subtree(child_name, child_path, v1_map, v2_map))
+
+    return Node(
+        name=name,
+        tokens=tokens_v2 or tokens_v1,
+        children=children,
+        text=text_v2 or text_v1,
+        change=change,
+        delta=delta,
+    )
+
+
+def build_diff_tree(v1: Node, v2: Node) -> Node:
+    """Build a unified diff tree from two prompt trees."""
+    v1_map = _all_paths(v1, ())
+    v2_map = _all_paths(v2, ())
+    root_name = v2.name if v2.name == v1.name else "diff"
+    root_path = (root_name,)
+    # Ensure both maps contain the chosen root
+    if (v1.name,) in v1_map and (v2.name,) in v2_map and v1.name == v2.name:
+        v1_map[root_path] = v1_map.pop((v1.name,))
+        v2_map[root_path] = v2_map.pop((v2.name,))
+    tree = _build_subtree(root_name, root_path, v1_map, v2_map)
+    tree.tokens = max(v1.tokens, v2.tokens, 1)
+    return tree
+
+
+def diff_prompts(
+    v1: Any,
+    v2: Any,
+    output: str | None = "prompt_diff.html",
+    title: str | None = None,
+    tokenizer: Any = None,
+    cost_per_token: float | None = None,
+    width: int = 1200,
+    height: int = 720,
+) -> str:
+    """Build and render a diff flamegraph between two prompts."""
+    from .core import build_tree
+    from .render import to_html
+
+    tree1 = build_tree(v1, name="v1", tokenizer=tokenizer)
+    tree2 = build_tree(v2, name="v2", tokenizer=tokenizer)
+    diff_tree = build_diff_tree(tree1, tree2)
+    html = to_html(
+        diff_tree,
+        title=title or "Prompt Diff",
+        cost_per_token=cost_per_token,
+        width=width,
+        height=height,
+    )
+    if output:
+        with open(output, "w", encoding="utf-8") as f:
+            f.write(html)
+    return html