diff --git a/action.yml b/action.yml index cf6382d..98f1614 100644 --- a/action.yml +++ b/action.yml @@ -73,6 +73,18 @@ inputs: description: 'Path to a Markdown style configuration file (.md) containing documentation style guidelines. If not set, auto-detects .code-to-docs/style.md in the repository root.' required: false default: '' + cost-per-1m-input: + description: 'Cost per 1M input tokens (USD). When set alongside cost-per-1m-output, an estimated cost is shown in the token usage summary. When unset, only token counts are reported.' + required: false + default: '' + cost-per-1m-output: + description: 'Cost per 1M output tokens (USD). When set alongside cost-per-1m-input, an estimated cost is shown in the token usage summary. When unset, only token counts are reported.' + required: false + default: '' + debug-artifacts: + description: 'When true, include full prompt and response text in the run log artifact. Default false.' + required: false + default: 'false' outputs: status: @@ -81,6 +93,8 @@ outputs: description: 'JSON array of modified files' pr-created: description: 'Whether a PR was created' + acceptance-rate: + description: 'Suggestion acceptance rate from previous review (e.g. "4/6")' runs: using: 'docker' @@ -104,3 +118,6 @@ runs: GOOGLE_SA_KEY: ${{ inputs.google-sa-key }} MAX_CONTEXT_CHARS: ${{ inputs.max-context-chars }} STYLE_CONFIG_PATH: ${{ inputs.style-config-path }} + COST_PER_1M_INPUT: ${{ inputs.cost-per-1m-input }} + COST_PER_1M_OUTPUT: ${{ inputs.cost-per-1m-output }} + DEBUG_ARTIFACTS: ${{ inputs.debug-artifacts }} diff --git a/src/generation.py b/src/generation.py index 26f8ba4..2d1f29b 100644 --- a/src/generation.py +++ b/src/generation.py @@ -157,6 +157,7 @@ def generate_updates_parallel( file_instructions=None, style_guidelines="", pr_description="", + usage_tracker=None, ): """ Generate documentation updates in parallel. @@ -190,6 +191,7 @@ def process_file(file_path): file_instructions=file_instructions, style_guidelines=style_guidelines, pr_description=pr_description, + usage_tracker=usage_tracker, ) if updated.strip() == "NO_UPDATE_NEEDED": @@ -240,6 +242,7 @@ def ask_ai_for_updated_content( file_instructions=None, style_guidelines="", pr_description="", + usage_tracker=None, ): is_markdown = file_path.endswith(".md") is_asciidoc = file_path.endswith(".adoc") @@ -421,6 +424,8 @@ def ask_ai_for_updated_content( {"role": "user", "content": prompt}, ], ) + if usage_tracker: + usage_tracker.record("generation", response) output = (response.choices[0].message.content or "").strip() except Exception as e: check_context_error(e) @@ -463,6 +468,8 @@ def ask_ai_for_updated_content( {"role": "user", "content": fix_prompt}, ], ) + if usage_tracker: + usage_tracker.record("format-fix", fix_response) output = (fix_response.choices[0].message.content or "").strip() output = strip_code_fences(output) if not output.endswith("\n"): diff --git a/src/run_log.py b/src/run_log.py new file mode 100644 index 0000000..c2767a3 --- /dev/null +++ b/src/run_log.py @@ -0,0 +1,58 @@ +"""Structured JSONL run log for code-to-docs. + +Writes one record per LLM call so failed runs can be debugged without +re-running. Optionally includes full prompt/response text behind the +debug-artifacts flag. +""" + +import json +import threading +import time +from pathlib import Path + +from security_utils import sanitize_output + + +class RunLog: + """Append-only JSONL log of LLM calls.""" + + def __init__(self, path="/tmp/code-to-docs-run.jsonl", include_prompts=False): + self._path = Path(path) + self._include_prompts = include_prompts + self._lock = threading.Lock() + self._path.parent.mkdir(parents=True, exist_ok=True) + if self._path.exists(): + self._path.unlink() + + @property + def path(self): + return str(self._path) + + @property + def has_entries(self): + return self._path.exists() and self._path.stat().st_size > 0 + + def record(self, stage, file_path, prompt, response_text, usage_obj, latency_ms, outcome): + """Write a single log record.""" + entry = { + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "stage": stage, + "file_path": file_path or "", + "prompt_length": len(prompt) if prompt else 0, + "response_length": len(response_text) if response_text else 0, + "prompt_tokens": getattr(usage_obj, "prompt_tokens", None) if usage_obj else None, + "completion_tokens": ( + getattr(usage_obj, "completion_tokens", None) if usage_obj else None + ), + "latency_ms": round(latency_ms), + "outcome": outcome, + } + + if self._include_prompts: + entry["prompt"] = sanitize_output(prompt or "") + entry["response"] = sanitize_output(response_text or "") + + line = json.dumps(entry, ensure_ascii=False) + with self._lock: + with open(self._path, "a", encoding="utf-8") as f: + f.write(line + "\n") diff --git a/src/suggest_docs.py b/src/suggest_docs.py index 805f09d..fe02f26 100644 --- a/src/suggest_docs.py +++ b/src/suggest_docs.py @@ -54,7 +54,9 @@ format_feature_review_section, parse_feature_command, ) +from run_log import RunLog from security_utils import run_command_safe, sanitize_output +from telemetry import UsageTracker _GITHUB_NAME_RE = re.compile(r"^[a-zA-Z0-9._-]+$") @@ -244,6 +246,18 @@ def main(): source = "MAX_CONTEXT_CHARS" if raw else "default" print(f"Context budget: {budget:,} chars (from {source})") + # Initialize token usage tracking + cost_input = os.environ.get("COST_PER_1M_INPUT", "") + cost_output = os.environ.get("COST_PER_1M_OUTPUT", "") + usage_tracker = UsageTracker( + cost_per_1m_input=float(cost_input) if cost_input else None, + cost_per_1m_output=float(cost_output) if cost_output else None, + ) + + # Initialize structured run log + debug_artifacts = os.environ.get("DEBUG_ARTIFACTS", "false").lower() == "true" + run_log = RunLog(include_prompts=debug_artifacts) + # Load persistent style guidelines from the base branch so the AI always # uses the repo's current style config, even if the PR branch predates it. style_guidelines = load_style_config_from_branch() @@ -413,6 +427,13 @@ def main(): previous_review = parse_previous_review(pr_number) if previous_review["review_found"]: + suggested = len(previous_review["accepted_files"]) + len( + previous_review["rejected_files"] + ) + accepted = len(previous_review["accepted_files"]) + if suggested > 0: + print(f"Acceptance rate: suggested={suggested} accepted={accepted}") + if previous_review["review_commit"] and commit_info: if previous_review["review_commit"] != commit_info["short_hash"]: print( @@ -542,6 +563,7 @@ def main(): file_instructions=file_instructions, style_guidelines=style_guidelines, pr_description=pr_description, + usage_tracker=usage_tracker, ) for file_path, _current, updated in files_with_content: @@ -566,6 +588,7 @@ def main(): file_instructions=file_instructions, style_guidelines=style_guidelines, pr_description=pr_description, + usage_tracker=usage_tracker, ) if updated.strip() == "NO_UPDATE_NEEDED": @@ -762,6 +785,9 @@ def main(): confirm_parts.append( "A docs PR has been created/updated with these changes." ) + if usage_tracker.has_records: + confirm_parts.append("") + confirm_parts.append(usage_tracker.format_summary()) confirm_body = "\n".join(confirm_parts) confirm_file = Path("/tmp/update_confirm.md") confirm_file.write_text(confirm_body, encoding="utf-8") @@ -788,6 +814,9 @@ def main(): else: print("All documentation is already up to date — no PR created.") + if run_log.has_entries: + print(f"Run log written to: {run_log.path}") + if __name__ == "__main__": main() diff --git a/src/telemetry.py b/src/telemetry.py new file mode 100644 index 0000000..2f68d90 --- /dev/null +++ b/src/telemetry.py @@ -0,0 +1,94 @@ +"""Token usage tracking and cost estimation for code-to-docs.""" + +import threading + + +class UsageTracker: + """Thread-safe accumulator for LLM API token usage. + + Records prompt and completion tokens per call, tagged by stage. + Optionally computes estimated cost when per-token pricing is provided. + """ + + def __init__(self, cost_per_1m_input=None, cost_per_1m_output=None): + self._records = [] + self._lock = threading.Lock() + self._cost_input = cost_per_1m_input + self._cost_output = cost_per_1m_output + + def record(self, stage, response): + """Extract usage from an OpenAI-compatible response and store it.""" + usage = getattr(response, "usage", None) + prompt_tokens = None + completion_tokens = None + if usage: + prompt_tokens = getattr(usage, "prompt_tokens", None) + completion_tokens = getattr(usage, "completion_tokens", None) + with self._lock: + self._records.append( + { + "stage": stage, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + } + ) + + @property + def has_records(self): + return len(self._records) > 0 + + def _aggregate(self): + """Aggregate token counts by stage.""" + stages = {} + for r in self._records: + s = r["stage"] + if s not in stages: + stages[s] = {"calls": 0, "prompt": 0, "completion": 0, "reported": True} + stages[s]["calls"] += 1 + if r["prompt_tokens"] is not None: + stages[s]["prompt"] += r["prompt_tokens"] + else: + stages[s]["reported"] = False + if r["completion_tokens"] is not None: + stages[s]["completion"] += r["completion_tokens"] + else: + stages[s]["reported"] = False + return stages + + def format_summary(self): + """Return a Markdown
block with per-stage token breakdown.""" + if not self._records: + return "" + + stages = self._aggregate() + total_prompt = 0 + total_completion = 0 + all_reported = True + + lines = [] + lines.append("| Stage | Calls | Input tokens | Output tokens |") + lines.append("|-------|------:|-------------:|--------------:|") + for stage, data in stages.items(): + if data["reported"]: + lines.append( + f"| {stage} | {data['calls']} | {data['prompt']:,} | {data['completion']:,} |" + ) + total_prompt += data["prompt"] + total_completion += data["completion"] + else: + lines.append(f"| {stage} | {data['calls']} | not reported | not reported |") + all_reported = False + + lines.append( + f"| **Total** | **{len(self._records)}** " + f"| **{total_prompt:,}** | **{total_completion:,}** |" + ) + + if self._cost_input and self._cost_output and all_reported: + cost = (total_prompt / 1_000_000) * self._cost_input + ( + total_completion / 1_000_000 + ) * self._cost_output + lines.append(f"\nEstimated cost: ${cost:.4f}") + + table = "\n".join(lines) + return f"
\nToken usage\n\n{table}\n\n
" diff --git a/tests/test_run_log.py b/tests/test_run_log.py new file mode 100644 index 0000000..76baef9 --- /dev/null +++ b/tests/test_run_log.py @@ -0,0 +1,48 @@ +"""Tests for run_log.py -- structured JSONL run log.""" + +import json +from unittest.mock import MagicMock + +from run_log import RunLog + + +class TestRunLog: + def test_writes_jsonl_record(self, tmp_path): + log = RunLog(path=str(tmp_path / "test.jsonl")) + usage = MagicMock() + usage.prompt_tokens = 100 + usage.completion_tokens = 50 + log.record("generation", "docs/guide.md", "prompt text", "response text", usage, 1500, "ok") + assert log.has_entries + lines = (tmp_path / "test.jsonl").read_text().strip().split("\n") + assert len(lines) == 1 + entry = json.loads(lines[0]) + assert entry["stage"] == "generation" + assert entry["file_path"] == "docs/guide.md" + assert entry["prompt_tokens"] == 100 + assert entry["latency_ms"] == 1500 + + def test_excludes_prompts_by_default(self, tmp_path): + log = RunLog(path=str(tmp_path / "test.jsonl")) + log.record("generation", "f.md", "secret prompt", "secret response", None, 100, "ok") + entry = json.loads((tmp_path / "test.jsonl").read_text().strip()) + assert "prompt" not in entry + assert "response" not in entry + + def test_includes_prompts_when_enabled(self, tmp_path): + log = RunLog(path=str(tmp_path / "test.jsonl"), include_prompts=True) + log.record("generation", "f.md", "the prompt", "the response", None, 100, "ok") + entry = json.loads((tmp_path / "test.jsonl").read_text().strip()) + assert entry["prompt"] == "the prompt" + assert entry["response"] == "the response" + + def test_empty_log_has_no_entries(self, tmp_path): + log = RunLog(path=str(tmp_path / "test.jsonl")) + assert not log.has_entries + + def test_handles_none_usage(self, tmp_path): + log = RunLog(path=str(tmp_path / "test.jsonl")) + log.record("generation", "f.md", "p", "r", None, 100, "ok") + entry = json.loads((tmp_path / "test.jsonl").read_text().strip()) + assert entry["prompt_tokens"] is None + assert entry["completion_tokens"] is None diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..b53eaf9 --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,80 @@ +"""Tests for telemetry.py -- token usage tracking.""" + +from unittest.mock import MagicMock + +from telemetry import UsageTracker + + +class TestUsageTracker: + def test_record_with_usage(self): + tracker = UsageTracker() + resp = MagicMock() + resp.usage.prompt_tokens = 100 + resp.usage.completion_tokens = 50 + tracker.record("generation", resp) + assert tracker.has_records + summary = tracker.format_summary() + assert "100" in summary + assert "50" in summary + assert "generation" in summary + + def test_record_without_usage(self): + tracker = UsageTracker() + resp = MagicMock(spec=[]) # no usage attribute + tracker.record("generation", resp) + summary = tracker.format_summary() + assert "not reported" in summary + + def test_multiple_stages(self): + tracker = UsageTracker() + for stage in ("generation", "verification", "generation"): + resp = MagicMock() + resp.usage.prompt_tokens = 100 + resp.usage.completion_tokens = 50 + tracker.record(stage, resp) + summary = tracker.format_summary() + assert "generation" in summary + assert "verification" in summary + assert "**3**" in summary # total calls + + def test_format_with_cost(self): + tracker = UsageTracker(cost_per_1m_input=3.0, cost_per_1m_output=15.0) + resp = MagicMock() + resp.usage.prompt_tokens = 1_000_000 + resp.usage.completion_tokens = 100_000 + tracker.record("generation", resp) + summary = tracker.format_summary() + assert "$" in summary + assert "4.5000" in summary # 3.0 + 1.5 + + def test_format_without_cost(self): + tracker = UsageTracker() + resp = MagicMock() + resp.usage.prompt_tokens = 500 + resp.usage.completion_tokens = 100 + tracker.record("generation", resp) + summary = tracker.format_summary() + assert "$" not in summary + + def test_empty_tracker(self): + tracker = UsageTracker() + assert not tracker.has_records + assert tracker.format_summary() == "" + + def test_cost_suppressed_when_usage_missing(self): + tracker = UsageTracker(cost_per_1m_input=3.0, cost_per_1m_output=15.0) + resp = MagicMock(spec=[]) + tracker.record("generation", resp) + summary = tracker.format_summary() + assert "$" not in summary + + def test_details_block_structure(self): + tracker = UsageTracker() + resp = MagicMock() + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + tracker.record("test", resp) + summary = tracker.format_summary() + assert summary.startswith("
") + assert "
" in summary + assert "Token usage" in summary