luckyd-code 1.2.2__py3-none-any.whl

luckyd_code/context.py

@@ -0,0 +1,214 @@
+import logging
+from typing import List, Dict, Any, Optional
+
+_log = logging.getLogger(__name__)
+
+
+def _get_accurate_token_count(text: str) -> int:
+    """Count tokens using tiktoken if available, fall back to heuristic.
+
+    Note: tiktoken uses OpenAI's cl100k_base encoding, not DeepSeek's native
+    tokenizer. DeepSeek's vocabulary is similar but not identical, so counts
+    may be slightly off. We apply a 15% safety multiplier to avoid hitting
+    the context window before compaction triggers.
+    """
+    _DEEPSEEK_SAFETY_FACTOR = 1.15
+    try:
+        import tiktoken
+        enc = tiktoken.get_encoding("cl100k_base")
+        return int(len(enc.encode(text)) * _DEEPSEEK_SAFETY_FACTOR)
+    except Exception:
+        # Code is 2-3x more token-dense than prose due to symbols/indentation
+        if any(c in text for c in '{}(\n') or text.count('    ') > 2:
+            return max(len(text) // 3, 1)
+        return max(len(text) // 4, 1)
+
+
+class ConversationContext:
+    """Manages conversation history and message structure."""
+
+    def __init__(self, system_prompt: str, max_messages: int = 100,
+                 config=None, model: str = "deepseek-v4-flash"):
+        self.messages: List[Dict[str, Any]] = [
+            {"role": "system", "content": system_prompt}
+        ]
+        self.max_messages = max_messages
+        self._config = config
+        self._model = model
+        # Auto-compact when estimated tokens exceed this threshold.
+        # DeepSeek V4 Flash and V4 Pro both support 1M context windows.
+        # We compact at 800K to leave ~200K headroom for the response and
+        # any injected tool results. Users on older models with smaller
+        # context windows can lower this via config.
+        self._token_compact_threshold = 800_000
+
+    def add_user_message(self, content: str):
+        self.messages.append({"role": "user", "content": content})
+        self._maybe_trim()
+        # Token-aware auto-compaction: prevent silent context window overflow
+        if self._config is not None and self.estimate_tokens() > self._token_compact_threshold:
+            self.compact(self._config, self._model, keep_last=8)
+
+    def add_assistant_message(self, content: Optional[str] = None, tool_calls: Optional[List[Dict[str, Any]]] = None, reasoning_content: Optional[str] = None):
+        msg: dict[str, Any] = {"role": "assistant"}
+        if content is not None:
+            msg["content"] = content
+        elif reasoning_content:
+            # DeepSeek API requires 'content' to always be present alongside
+            # 'reasoning_content' in multi-turn requests. Omitting it causes
+            # "content or tool_calls must be set". Default to empty string.
+            msg["content"] = ""
+        if tool_calls:
+            msg["tool_calls"] = tool_calls
+        if reasoning_content:
+            msg["reasoning_content"] = reasoning_content
+        self.messages.append(msg)
+        self._maybe_trim()
+
+    def add_tool_result(self, tool_call_id: str, tool_name: str, result: str):
+        self.messages.append({
+            "role": "tool",
+            "tool_call_id": tool_call_id,
+            "content": result,
+        })
+        self._maybe_trim()
+
+    def get_messages(self) -> List[Dict[str, Any]]:
+        return self.messages
+
+    @staticmethod
+    def _drop_orphaned_tool_messages(
+        messages: List[Dict[str, Any]],
+    ) -> List[Dict[str, Any]]:
+        """Remove tool-result messages whose parent assistant tool_call is absent.
+
+        This keeps the message list valid for the DeepSeek API, which requires
+        every ``role=tool`` message to be preceded by an assistant message that
+        contains a matching ``tool_call_id`` in its ``tool_calls`` list.
+        """
+        valid_parent_ids: set = set()
+        filtered: List[Dict[str, Any]] = []
+        for msg in messages:
+            role = msg.get("role")
+            if role == "assistant":
+                for tc in msg.get("tool_calls") or []:
+                    valid_parent_ids.add(tc.get("id", ""))
+                filtered.append(msg)
+            elif role == "tool":
+                if msg.get("tool_call_id") in valid_parent_ids:
+                    filtered.append(msg)
+                # orphaned tool result — silently dropped
+            else:
+                # user / system messages reset parent-id tracking
+                valid_parent_ids.clear()
+                filtered.append(msg)
+        return filtered
+
+    def _maybe_trim(self):
+        """Trim oldest messages if we exceed max_messages, keeping system prompt."""
+        if len(self.messages) > self.max_messages:
+            keep = [self.messages[0]] + self.messages[-(self.max_messages - 1):]
+            self.messages = self._drop_orphaned_tool_messages(keep)
+
+    def count_messages(self) -> int:
+        return len(self.messages)
+
+    def estimate_tokens(self) -> int:
+        """Estimate total tokens using tiktoken (accurate) with fallback to chars/4.
+
+        Accounts for all message fields the API consumes: content, tool_calls,
+        reasoning_content, and tool_call_ids.
+        """
+        total = 0
+        for msg in self.messages:
+            content = str(msg.get("content", ""))
+            total += _get_accurate_token_count(content)
+            # Account for tool_calls in assistant messages
+            for tc in msg.get("tool_calls", []):
+                fn = tc.get("function", {})
+                total += _get_accurate_token_count(str(fn.get("name", "")))
+                total += _get_accurate_token_count(str(fn.get("arguments", "")))
+            # Account for tool_call_id in tool messages
+            if msg.get("tool_call_id"):
+                total += 5
+            # DeepSeek reasoning_content (thinking blocks) — preserved across
+            # multi-turn conversations and required by the API.  These can be
+            # surprisingly large (hundreds of tokens per assistant turn) and
+            # were previously invisible to the compaction threshold.
+            reasoning = msg.get("reasoning_content", "")
+            if reasoning:
+                total += _get_accurate_token_count(str(reasoning))
+            # Role overhead (~5 tokens per message)
+            total += 5
+        return total
+
+    def reset(self, system_prompt: Optional[str] = None):
+        if system_prompt:
+            self.messages = [{"role": "system", "content": system_prompt}]
+        else:
+            system = self.messages[0]
+            self.messages = [system]
+
+    def compact(self, config, model: str, keep_last: int = 5,
+                on_compact=None) -> str:
+        """Compact conversation by summarizing older messages using the model.
+
+        If *on_compact* is a callable, it is invoked with
+        ``(summary_text, compacted_count)`` after a successful compaction.
+        """
+        if len(self.messages) <= keep_last + 1:
+            return "Nothing to compact"
+
+        system = self.messages[0]
+        keep = self.messages[-keep_last:]
+        compact_targets = self.messages[1:-keep_last]
+
+        summary_text = "\n".join(
+            f"{m['role']}: {str(m.get('content', ''))[:800]}"
+            + (
+                f"\n[key reasoning]: {str(m['reasoning_content'])[:200]}"
+                if m.get("reasoning_content") else ""
+            )
+            for m in compact_targets
+        )
+
+        summary_prompt = (
+            "Summarize the following conversation history concisely. "
+            "Capture key decisions, code changes, file paths, and the user's goals:"
+            f"\n\n{summary_text}"
+        )
+
+        try:
+            from openai import OpenAI
+            import httpx
+            # Always use Flash for compaction — it's fast, cheap, and
+            # summarisation doesn't benefit from the Pro model's extra power.
+            compact_model = "deepseek-v4-flash"
+            client = OpenAI(
+                api_key=config.api_key,
+                base_url=config.base_url,
+                http_client=httpx.Client(timeout=30),
+            )
+            response = client.chat.completions.create(
+                model=compact_model,
+                messages=[{"role": "user", "content": summary_prompt}],
+                max_tokens=1500,
+            )
+            summary = response.choices[0].message.content or ""
+        except Exception as e:
+            return f"Compaction failed: {e}"
+
+        raw = [system, {
+            "role": "system",
+            "content": f"[Compacted conversation summary]: {summary}",
+        }] + keep
+        self.messages = self._drop_orphaned_tool_messages(raw)
+
+        count = len(compact_targets)
+        if callable(on_compact):
+            try:
+                on_compact(summary, count)
+            except Exception:
+                _log.warning("Compaction callback failed", exc_info=True)
+
+        return f"Compacted {count} messages into a summary"

luckyd_code/cost_tracker.py

@@ -0,0 +1,209 @@
+"""Track API usage costs per session with persistence."""
+
+import json
+import logging
+from dataclasses import dataclass, field, fields
+from datetime import datetime
+from typing import Any, Optional
+
+from ._data_dir import data_path, legacy_path
+
+COST_FILE = data_path("costs.jsonl")  # append-only, one JSON record per line
+_LEGACY_COST_FILE = legacy_path("costs.json")  # migrated on first write
+_TOTALS_FILE = data_path("costs_total.json")  # single-float running total
+
+_logger = logging.getLogger("luckyd_code.cost_tracker")
+
+
+@dataclass
+class UsageRecord:
+    model: str
+    input_tokens: int
+    output_tokens: int
+    timestamp: str = ""
+    estimated_cost: float = 0.0
+    # Internal flag — excluded from serialization so it never appears in
+    # costs.jsonl.  repr=False keeps it out of str() output too.
+    _cost_provided: bool = field(default=False, repr=False)
+
+    def __post_init__(self):
+        if not self.timestamp:
+            self.timestamp = datetime.now().isoformat()
+        # Only recalculate when cost was NOT explicitly provided by the caller.
+        # This preserves intentional zero-cost values (e.g. cached / free responses).
+        if not self._cost_provided and (self.input_tokens > 0 or self.output_tokens > 0):
+            self.estimated_cost = self._calc_cost()
+
+    def to_dict(self) -> dict:
+        """Serialize to a plain dict, excluding internal/private fields."""
+        return {
+            f.name: getattr(self, f.name)
+            for f in fields(self)
+            if not f.name.startswith("_")
+        }
+
+    def _calc_cost(self) -> float:
+        """Estimate cost using approximate per-token rates.
+
+        Rates (per 1K tokens, USD):
+          deepseek-v4-flash : $0.000140 in / $0.000280 out
+          deepseek-v4-pro   : $0.001740 in / $0.003480 out  (75% discount until 2026-05-31)
+          deepseek-chat     : legacy alias for deepseek-v4-flash
+          deepseek-reasoner : legacy alias for deepseek-v4-flash (thinking mode)
+        Prices sourced from api-docs.deepseek.com/quick_start/pricing (2026-04-26)
+        """
+        rates = {
+            # Current V4 models
+            "deepseek-v4-flash": (0.000140, 0.000280),
+            "deepseek-v4-pro":   (0.001740, 0.003480),
+            # Legacy names — now route to deepseek-v4-flash
+            "deepseek-chat":     (0.000140, 0.000280),
+            "deepseek-reasoner": (0.000140, 0.000280),
+            # Older models (kept for historical cost records)
+            "deepseek-v3-0324":  (0.000270, 0.001100),
+            "deepseek-v3":       (0.000270, 0.001100),
+        }
+        # Default to v4-flash pricing for unknown models
+        input_rate, output_rate = rates.get(self.model, (0.000140, 0.000280))
+        return (self.input_tokens / 1000 * input_rate) + (
+            self.output_tokens / 1000 * output_rate
+        )
+
+
+class CostTracker:
+    """Records API usage costs per session with cumulative tracking."""
+
+    def __init__(self):
+        self.session_start = datetime.now()
+        self.records: list[UsageRecord] = []
+        self._written_count: int = 0  # how many records already flushed to disk
+
+    def record_usage(self, model: str, input_tokens: int, output_tokens: int,
+                     cost: Optional[float] = None) -> UsageRecord:
+        rec = UsageRecord(
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            estimated_cost=cost if cost is not None else 0.0,
+            _cost_provided=cost is not None,
+        )
+        self.records.append(rec)
+        self._append_new_records()
+        return rec
+
+    def get_session_cost(self) -> float:
+        return sum(r.estimated_cost for r in self.records)
+
+    def get_session_tokens(self) -> tuple[int, int]:
+        inp = sum(r.input_tokens for r in self.records)
+        out = sum(r.output_tokens for r in self.records)
+        return inp, out
+
+    def get_cumulative_cost(self) -> float:
+        """Return the lifetime total cost in O(1) using the sidecar totals file.
+
+        Falls back to summing the full JSONL on first run (migration) and
+        writes the total to the sidecar file for all future calls.
+        """
+        # Fast path: read the single-value sidecar
+        if _TOTALS_FILE.exists():
+            try:
+                data = json.loads(_TOTALS_FILE.read_text(encoding="utf-8"))
+                return float(data.get("total", 0.0))
+            except Exception:
+                pass
+        # Slow path (first run / migration): sum the full JSONL, then persist
+        total = sum(r.get("estimated_cost", 0) for r in self._load_all())
+        self._write_total(total)
+        return total
+
+    @staticmethod
+    def _write_total(total: float) -> None:
+        """Persist the running total to the sidecar file."""
+        try:
+            _TOTALS_FILE.parent.mkdir(parents=True, exist_ok=True)
+            _TOTALS_FILE.write_text(json.dumps({"total": total}), encoding="utf-8")
+        except Exception:
+            pass
+
+    def get_stats(self) -> str:
+        inp, out = self.get_session_tokens()
+        cost = self.get_session_cost()
+        total_cost = self.get_cumulative_cost()
+        lines = [
+            "[bold]Cost Tracking[/bold]",
+            f"  Session tokens: {inp:,} in / {out:,} out",
+            f"  Session cost: ${cost:.4f}",
+            f"  Cumulative cost: ${total_cost:.4f}",
+            f"  API calls this session: {len(self.records)}",
+        ]
+        return "\n".join(lines)
+
+    def reset_cumulative(self) -> str:
+        """Wipe the persistent costs.jsonl and reset session records."""
+        self.records.clear()
+        self._written_count = 0
+        try:
+            for f in (COST_FILE, _LEGACY_COST_FILE, _TOTALS_FILE):
+                if f.exists():
+                    f.unlink()
+            return "Cumulative cost history cleared."
+        except Exception as e:
+            return f"Failed to clear cost file: {e}"
+
+    def _append_new_records(self):
+        """Append only new records to the JSONL file — O(1) per call."""
+        new_records = self.records[self._written_count:]
+        if not new_records:
+            return
+        COST_FILE.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            self._migrate_legacy_json_once()
+            with COST_FILE.open("a", encoding="utf-8") as fh:
+                for r in new_records:
+                    fh.write(json.dumps(r.to_dict()) + "\n")
+            self._written_count = len(self.records)
+            # Keep running total sidecar in sync — O(1) increment
+            new_cost = sum(r.estimated_cost for r in new_records)
+            if new_cost:
+                current = self.get_cumulative_cost()
+                self._write_total(current + new_cost)
+        except Exception:
+            _logger.warning("Failed to persist cost records", exc_info=True)
+
+    @staticmethod
+    def _migrate_legacy_json_once():
+        """One-time migration: convert costs.json → costs.jsonl."""
+        if not _LEGACY_COST_FILE.exists() or COST_FILE.exists():
+            return
+        try:
+            records = json.loads(_LEGACY_COST_FILE.read_text(encoding="utf-8"))
+            with COST_FILE.open("w", encoding="utf-8") as fh:
+                for r in records:
+                    fh.write(json.dumps(r) + "\n")
+            _LEGACY_COST_FILE.unlink()
+        except Exception:
+            _logger.warning("Failed to migrate legacy costs.json", exc_info=True)
+
+    @staticmethod
+    def _load_all() -> list[dict[str, Any]]:
+        """Load all records from JSONL (plus legacy JSON if present)."""
+        records: list[dict[str, Any]] = []
+        # Legacy fallback — only present before first migration
+        if _LEGACY_COST_FILE.exists() and not COST_FILE.exists():
+            try:
+                records = json.loads(_LEGACY_COST_FILE.read_text(encoding="utf-8"))
+                return records
+            except Exception:
+                _logger.warning("Failed to load legacy cost records", exc_info=True)
+                return []
+        if COST_FILE.exists():
+            try:
+                with COST_FILE.open(encoding="utf-8") as fh:
+                    for line in fh:
+                        line = line.strip()
+                        if line:
+                            records.append(json.loads(line))
+            except Exception:
+                _logger.warning("Failed to load cost records", exc_info=True)
+        return records