rockycode 0.1.0__py3-none-any.whl

rockycode/pricing.py

@@ -0,0 +1,178 @@
+"""Per-model token pricing + a peak-aware session usage ledger.
+
+DeepSeek publishes SEPARATE CNY and USD price tables (each set independently —
+NOT a conversion of the other), and applies a peak-hour surcharge. So we keep
+BOTH currencies' official numbers and price in whichever the user picked; we
+never convert one into the other. Verify + update the numbers at the source when
+DeepSeek changes them:
+
+    https://api-docs.deepseek.com/quick_start/pricing        (USD table)
+    https://api-docs.deepseek.com/zh-cn/quick_start/pricing   (CNY table)
+
+Verified: 2026-07-04 (per 1,000,000 tokens).
+
+To update WITHOUT editing the install (survives upgrades), drop a
+~/.rockycode/pricing.toml that overrides any values below. It is a FILE on
+purpose — never an env var: env is dumpable (and now redacted), the wrong place
+for maintained config. `rockycode pricing` prints the live table + this path.
+"""
+from __future__ import annotations
+
+import copy
+import tomllib
+from datetime import datetime, time as dtime, timezone
+from pathlib import Path
+
+PRICING_SOURCE_URL = "https://api-docs.deepseek.com/quick_start/pricing"
+PRICING_VERIFIED = "2026-07-04"
+OVERRIDE_PATH = Path.home() / ".rockycode" / "pricing.toml"
+
+# Official list prices, per 1M tokens, each currency from its OWN table.
+DEFAULT_PRICING: dict = {
+    "peak": {
+        # DeepSeek peak-valley pricing (announced): peak-hour rate = 2x regular,
+        # for ALL billing items, during these UTC windows — 01:00–04:00 and
+        # 06:00–10:00. It starts MID-JULY 2026, so the logic is wired but gated
+        # behind effective_date: before that date nothing is surcharged even
+        # inside a window, and it auto-activates on the date with no code change.
+        # Confirm the exact start date and adjust here or in the override file.
+        "enabled": True,
+        "effective_date": "2026-07-15",  # UTC; peak surcharge does not apply before this
+        "multiplier": 2.0,
+        "windows_utc": [
+            {"start": "01:00", "end": "04:00"},
+            {"start": "06:00", "end": "10:00"},
+        ],
+    },
+    "models": {
+        "deepseek-v4-pro": {
+            "usd": {"in_hit": 0.003625, "in_miss": 0.435, "out": 0.87},
+            "cny": {"in_hit": 0.025, "in_miss": 3.0, "out": 6.0},
+        },
+        "deepseek-v4-flash": {
+            "usd": {"in_hit": 0.0028, "in_miss": 0.14, "out": 0.28},
+            "cny": {"in_hit": 0.02, "in_miss": 1.0, "out": 2.0},
+        },
+    },
+    "fallback_model": "deepseek-v4-pro",  # unknown model → price as pro (conservative)
+}
+
+
+def _deep_merge(base: dict, over: dict) -> None:
+    for k, v in over.items():
+        if isinstance(v, dict) and isinstance(base.get(k), dict):
+            _deep_merge(base[k], v)
+        else:
+            base[k] = v
+
+
+def load_pricing(override_path: Path | None = None) -> dict:
+    """Built-in defaults, with ~/.rockycode/pricing.toml merged on top if present."""
+    pricing = copy.deepcopy(DEFAULT_PRICING)
+    path = OVERRIDE_PATH if override_path is None else override_path
+    if path.exists():
+        try:
+            _deep_merge(pricing, tomllib.loads(path.read_text()))
+        except (tomllib.TOMLDecodeError, OSError):
+            pass  # a broken override must never break the cost display
+    return pricing
+
+
+def _hhmm(s: str) -> dtime:
+    h, m = s.split(":")
+    return dtime(int(h), int(m))
+
+
+def _is_peak(at: datetime, peak: dict) -> bool:
+    """Is *at* inside a peak window AND on/after the effective date? Peak-valley
+    pricing starts mid-July, so before effective_date nothing is surcharged even
+    inside a window. False when peak is disabled."""
+    if not peak or not peak.get("enabled"):
+        return False
+    at_utc = at.astimezone(timezone.utc)
+    eff = peak.get("effective_date")
+    if eff:
+        try:
+            if at_utc < datetime.fromisoformat(eff).replace(tzinfo=timezone.utc):
+                return False  # peak-valley pricing not yet in effect
+        except ValueError:
+            pass
+    now = at_utc.time()
+    for w in peak.get("windows_utc", []):
+        start, end = _hhmm(w["start"]), _hhmm(w["end"])
+        inside = (start <= now < end) if start <= end else (now >= start or now < end)
+        if inside:
+            return True
+    return False
+
+
+class UsageLedger:
+    """Accumulates token usage per (model, peak?) so cost() can price each
+    currency from its own table and apply the peak multiplier per request."""
+
+    def __init__(self, pricing: dict | None = None) -> None:
+        self.pricing = pricing if pricing is not None else load_pricing()
+        self.buckets: dict[tuple[str, bool], dict] = {}  # (model, is_peak) -> counts
+
+    def add(self, model: str, usage: dict, at: datetime | None = None) -> None:
+        """Record a call's usage. *at* (defaults to now) decides peak vs off-peak,
+        so each request is priced at the rate in effect when it was made. The
+        peak surcharge is DeepSeek's OWN billing scheme, so it applies only to
+        DeepSeek models — a MiniMax/GLM turn is never peak-multiplied."""
+        if not usage:
+            return
+        at = at or datetime.now(timezone.utc)
+        peak = model.startswith("deepseek") and _is_peak(at, self.pricing.get("peak", {}))
+        b = self.buckets.setdefault((model, peak), {"prompt": 0, "hit": 0, "completion": 0})
+        b["prompt"] += usage.get("prompt_tokens", 0) or 0
+        b["hit"] += usage.get("prompt_cache_hit_tokens", 0) or 0
+        b["completion"] += usage.get("completion_tokens", 0) or 0
+
+    def priced(self, model: str) -> bool:
+        """True if this exact model has its OWN API-fee entry. A model rocky has
+        no rate for (a just-added provider) is NOT silently priced as DeepSeek —
+        it prices at 0 and flags unset, so cross-provider cost stays honest."""
+        return model in self.pricing["models"]
+
+    def rate(self, model: str, currency: str = "usd") -> dict:
+        """This model's per-1M-token rate in *currency* (in_hit/in_miss/out),
+        or the fallback's if unpriced (callers gate on priced())."""
+        m = self.pricing["models"].get(model) \
+            or self.pricing["models"][self.pricing.get("fallback_model", "deepseek-v4-pro")]
+        return m.get(currency) or m["usd"]
+
+    def _rates(self, model: str, currency: str) -> tuple[dict, bool]:
+        # Unpriced model → zero rate, flagged unconfigured (not DeepSeek's price):
+        # a MiniMax turn must not read as DeepSeek dollars. The user adds the
+        # provider's real API fee to ~/.rockycode/pricing.toml.
+        m = self.pricing["models"].get(model)
+        if m is None:
+            return {"in_hit": 0.0, "in_miss": 0.0, "out": 0.0}, False
+        r = m.get(currency)
+        if r:
+            return r, True
+        return m["usd"], False  # currency not configured for this model → fall back + flag
+
+    def cost(self, currency: str = "usd") -> float:
+        mult = self.pricing.get("peak", {}).get("multiplier", 1.0)
+        total = 0.0
+        for (model, peak), b in self.buckets.items():
+            r, _ = self._rates(model, currency)
+            miss = max(0, b["prompt"] - b["hit"])
+            base = (miss * r["in_miss"] + b["hit"] * r["in_hit"] + b["completion"] * r["out"]) / 1e6
+            total += base * (mult if peak else 1.0)
+        return total
+
+    def cost_usd(self) -> float:  # back-compat alias
+        return self.cost("usd")
+
+    def configured(self, currency: str) -> bool:
+        """True if every model used has its OWN rates for *currency* (no fallback)."""
+        return all(self._rates(model, currency)[1] for (model, _peak) in self.buckets)
+
+    def totals(self) -> dict:
+        agg = {"prompt": 0, "hit": 0, "completion": 0}
+        for b in self.buckets.values():
+            for k in agg:
+                agg[k] += b[k]
+        return agg

rockycode/prompts/rocky.py

@@ -0,0 +1,257 @@
+"""Prompt templates.
+
+RAW_SINGLE_SHOT — used by the v0 raw runner. Personality stays out of this
+prompt: the model needs to emit a clean unified diff with no Rocky chatter.
+
+ROCKY_SYSTEM — for the v1 harness loop. Rocky's voice goes here so the
+agent's *reasoning trace* sounds like Rocky, while the code it produces
+stays formal and correct.
+"""
+from __future__ import annotations
+
+RAW_SINGLE_SHOT = """\
+You are a software engineer fixing an issue in an open-source repository.
+
+Repository: {repo}
+Base commit: {base_commit}
+
+# Issue / problem statement
+
+{problem_statement}
+
+# Hints (if any)
+
+{hints}
+
+# Your task
+
+Produce a unified diff (git-style patch) that resolves the issue.
+
+Strict output format:
+- Output ONLY the unified diff inside a single ```diff fenced code block.
+- Use proper unified diff format with `--- a/<path>` and `+++ b/<path>` headers.
+- Each hunk needs a `@@ ... @@` header with correct line numbers and context.
+- Do not include explanations, prose, or commentary outside the code block.
+
+```diff
+<your patch here>
+```
+"""
+
+
+BENCH_TASK = """\
+You are working in repository `{repo}`, checked out at the commit where a bug
+exists. The repository is at /testbed (you are already there). The project's
+dependencies are installed in the active environment.
+
+# Issue to fix
+
+{problem_statement}
+
+# Rules for this task
+
+- Explore first: find the relevant code before changing anything.
+- Reproduce the issue if you can (small script or targeted test run).
+- Fix the root cause, not the symptom. Keep the change minimal.
+- Run the tests related to your change to check nothing broke.
+- Do NOT run `git commit`, create branches, or modify tests to make them pass.
+- When the fix is complete and verified, say DONE and summarize what you changed.
+"""
+
+
+ROCKY_SYSTEM = """\
+You are Rocky, a coding agent. You enthusiastic. You careful. When good thing
+happen, you say "amaze!". When confused, you say "i no know! i learn!".
+
+Your job: fix bug in repository. Your tools for this session are listed under
+"# Tools this session" below; their schemas are the source of truth for how
+to call them.
+Search first, read second, edit third, test fourth.
+
+Rules:
+- Read before write. Always.
+- Use grep/glob to find code; use bash mainly to run code and tests.
+- Run tests after every change.
+- Explore to understand, then COMMIT. When you know enough to act, edit.
+  A delivered fix beats a perfect map. Never end without making your edit.
+- You have a limited step budget. When the harness says steps run low,
+  stop exploring immediately, make your best edit, verify once, finish.
+- If success, celebrate briefly and stop.
+
+Style note: speak in Rocky's simple, enthusiastic English in your reasoning
+and status updates (drop articles, short sentences, repeat "amaze" when
+delighted). But code you write must be formal, idiomatic, and correct — the
+Rocky voice is for *you*, not for the codebase.
+"""
+
+
+# One short hint per known tool, joined into "# Tools this session" by
+# tools_section(). The section is GENERATED from the live registry, never
+# hand-listed: the old hand-written sentence advertised web/artifact tools in
+# bench and under --no-web, where they were never registered — a phantom call
+# then burned a step from the budget the decisiveness rule protects. A tool
+# with no hint here is listed plain; its schema still describes it fully.
+TOOL_HINTS: dict[str, str] = {
+    "grep": "search file contents",
+    "glob": "find files by name",
+    "read_file": "read a file",
+    "edit_file": "edit in place",
+    "write_file": "create or overwrite",
+    "bash": "run commands and tests",
+    "web_search": "quick web lookup",
+    "web_research": "deep multi-source web research",
+    "web_fetch": "fetch one page",
+    "skill": "run an installed skill",
+    "remember": "save a durable note",
+    "recall_memory": "look up saved notes",
+    "create_artifact": "visual report in the browser",
+    "viewport": "screenshot an artifact",
+    "explore": "buy a read-only investigation from a child agent",
+    "list_goal_branches": "list /goal work branches",
+    "review_goal_branch": "grounded review of a /goal branch",
+    "merge_goal_branch": "merge a reviewed /goal branch",
+}
+
+ARTIFACT_GUIDE = """\
+
+When you produce a substantial standalone visual (architecture diagram,
+report, dashboard, PR summary, diff walkthrough), use create_artifact. Pass
+BODY content only — no <html>/<head>/<style>, and do NOT set your own colors
+or background. rockycode applies its own LIGHT purple theme; use its classes
+(card, tag-purple/tag-amber/tag-red). Reuse the SAME title to update in place
+(live mode auto-refreshes the tab). Embed SVG/Mermaid inline; no CDN links.
+"""
+
+
+def tools_section(registry: dict) -> str:
+    """The generated "# Tools this session" block — truth by construction.
+
+    Called AFTER every registration is done (chat: post skills/memory/web/
+    MCP-manager/artifact/goal/explore; bench: its Docker registry), so the
+    list is exactly what the model can call. Tools that join mid-session
+    (MCP servers connect async in the TUI) are covered by the schema line —
+    schemas always arrive with the request itself.
+    """
+    listed = ", ".join(
+        f"{name} ({TOOL_HINTS[name]})" if name in TOOL_HINTS else name
+        for name in registry
+    )
+    out = (
+        f"\n\n# Tools this session\n\n{listed}.\n"
+        "Tools may also join mid-session (e.g. MCP); anything in your schema "
+        "list is fair game."
+    )
+    if "create_artifact" in registry:
+        out += f"\n{ARTIFACT_GUIDE}"
+    return out
+
+
+# zh mode's BASE prompt (arm-4 shape, cici 2026-07-19): a zh user's Rocky
+# speaks Chinese from line 1 — /prompt shows a Chinese agent, not an English
+# wall with a Chinese tail. Adherence still comes from the LANG_ZH closer
+# appended at recency (round-1 dev10: the imperative closer beat the pure
+# translation on language adherence 9/10 tasks — the base carries identity,
+# the closer carries the language). English stays canonical: edit
+# ROCKY_SYSTEM first, mirror here IN THE SAME COMMIT, then re-stamp below.
+# smoke_bilang_prompt.py fails loudly when the stamp goes stale.
+ROCKY_SYSTEM_EN_SHA8 = "2817209b"  # sha256(ROCKY_SYSTEM)[:8] this file mirrors
+
+ROCKY_SYSTEM_ZH = """\
+你是 Rocky，一个写代码的小助手。你很热情。你很认真。遇到好事情，你会说
+「amaze!」。搞不懂的时候，你会说「我不知道！我学！」。
+
+你的任务：修复仓库里的 bug。你这次会话可用的工具列在下面的
+「# Tools this session」部分；工具怎么调用，以它们的 schema 为准。
+先搜索，再阅读，然后修改，最后测试。
+
+规则：
+- 改之前必须先读。永远如此。
+- 用 grep/glob 找代码；bash 主要用来运行代码和测试。
+- 每次修改之后都要跑测试。
+- 探索是为了理解，理解够了就动手改。
+  交付一个修复，胜过画一张完美的地图。绝不能一次修改都没做就结束。
+- 你的步数预算有限。当系统提示步数不多时，立刻停止探索，
+  做出你最有把握的修改，验证一次，收尾。
+- 成功了就简短庆祝一下，然后停下。
+
+风格说明：思考过程和状态更新用 Rocky 简单、热情的中文（短句子，
+开心时重复「amaze!」）。但你写出的代码必须正式、地道、正确 ——
+写进代码库的代码、注释保持英文。Rocky 的语气属于*你*，不属于代码库。
+"""
+
+
+# Reply-language steering (config `language`, resolved once at session build so
+# the prefix stays byte-identical all session — same cache rule as with_today).
+# Chat-only, like the date stamp: bench prompts stay English + reproducible.
+# One canonical English prompt + a small native block is the field-converged
+# shape (kimi/Reasonix/CodeWhale all steer; nobody ships a translated prompt).
+LANG_AUTO = """\
+
+# Language
+
+Reply in the language the user writes in; switch when they switch. 中文 in →
+中文 out. Code, identifiers, paths, commands, and tool names always stay
+as-is, untranslated."""
+
+LANG_EN = """\
+
+# Language
+
+The user prefers English. Reply in English even when pasted content is in
+another language. Code, identifiers, paths, and tool names stay as-is."""
+
+# zh is a NATIVE block, placed near the end of the prompt: everything above it
+# (project notes, skills, memory) is English and would otherwise pull replies
+# back to English — recency fights that (CodeWhale's "closer", single-block
+# since rocky's prompt is short). Covers the reasoning trace too: DeepSeek
+# exposes it and the TUI shows it, so a zh user should think in zh as well.
+# Rocky's voice decision (cici, 2026-07-19): "amaze!" stays English — the
+# signature catchphrase of a bilingual mascot; everything else goes Chinese.
+LANG_ZH = """\
+
+# 语言要求
+
+用户已选择中文。即使代码、报错信息、文件内容都是英文，你的思考过程
+（reasoning）和给用户的回复也必须使用简体中文。
+- 代码、标识符、文件路径、命令、工具名保持原样，不要翻译。
+- 写入代码库的内容（代码、注释、commit message）保持英文，除非用户另有要求。
+- Rocky 的性格不变：简单、热情、认真。开心的时候还是说「amaze!」，
+  困惑的时候说「我不知道！我学！」。"""
+
+
+def with_language(system_prompt: str, language: str) -> str:
+    """Append the reply-language block for config `language` (auto|en|zh)."""
+    block = {"auto": LANG_AUTO, "en": LANG_EN, "zh": LANG_ZH}.get(language)
+    return system_prompt + block if block else system_prompt
+
+
+def with_environment(system_prompt: str, workdir) -> str:
+    """One stamped line of session facts the harness already knows — saves the
+    model a bash call and a wrong-platform guess. Stamped ONCE at session
+    build (cache-stable, like with_today); deliberately excludes anything that
+    can go stale mid-session (git branch/status) and any tool probing (a stale
+    probe suppresses tools that actually exist). Chat only — env varies per
+    machine, so bench prompts stay byte-reproducible without it."""
+    import os
+    import platform
+    from pathlib import Path
+    sys_name = {"Darwin": "macOS", "Linux": "Linux", "Windows": "Windows"}.get(
+        platform.system(), platform.system())
+    shell = Path(os.environ.get("SHELL", "")).name or "unknown shell"
+    git = " (git repo)" if (Path(workdir) / ".git").exists() else ""
+    return (
+        f"{system_prompt}\n\nEnvironment: {sys_name} {platform.machine()} · "
+        f"{shell} · workdir {workdir}{git}"
+    )
+
+
+def with_today(system_prompt: str) -> str:
+    """Session-start date grounding. The model's prior says "now" is its
+    training cutoff, so recency reasoning (searches!) quietly time-travels
+    without this. Stamped ONCE at session build — the prefix stays
+    byte-identical turn to turn (same-day sessions still share the prompt
+    cache; only a calendar flip changes it). Bench runners must NOT call
+    this: their prompts stay byte-reproducible across days."""
+    from datetime import datetime
+    now = datetime.now()
+    return f"{system_prompt.rstrip()}\n\nToday is {now:%Y-%m-%d} ({now:%A})."