llmcode-cli 1.0.0__py3-none-any.whl

llm_code/runtime/cost_tracker.py

@@ -0,0 +1,126 @@
+"""Token cost tracking with model pricing — user-customizable via config."""
+from __future__ import annotations
+from dataclasses import dataclass
+
+# Built-in fallback pricing per 1M tokens [input, output]
+BUILTIN_PRICING: dict[str, tuple[float, float]] = {
+    # OpenAI
+    "gpt-4o": (2.50, 10.00),
+    "gpt-4o-mini": (0.15, 0.60),
+    "gpt-4-turbo": (10.00, 30.00),
+    "gpt-3.5-turbo": (0.50, 1.50),
+    "o3": (2.00, 8.00),
+    "o4-mini": (0.50, 2.00),
+    # Anthropic
+    "claude-opus-4-6": (15.00, 75.00),
+    "claude-sonnet-4-6": (3.00, 15.00),
+    "claude-haiku-4-5": (0.80, 4.00),
+    # Google
+    "gemini-2.5-pro": (1.25, 10.00),
+    "gemini-2.5-flash": (0.15, 0.60),
+    # DeepSeek
+    "deepseek-chat": (0.27, 1.10),
+    "deepseek-reasoner": (0.55, 2.19),
+    # xAI
+    "grok-3": (3.00, 15.00),
+    "grok-3-mini": (0.30, 0.50),
+}
+
+
+class BudgetExceededError(Exception):
+    """Raised when the accumulated cost exceeds the configured budget."""
+
+    def __init__(self, spent: float, budget: float) -> None:
+        self.spent = spent
+        self.budget = budget
+        super().__init__(f"Budget limit (${budget:.2f}) exceeded (spent ${spent:.4f})")
+
+
+@dataclass
+class CostTracker:
+    model: str = ""
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    total_cost_usd: float = 0.0
+    custom_pricing: dict | None = None  # from config.json "pricing"
+    max_budget_usd: float | None = None
+
+    def add_usage(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_creation_tokens: int = 0,
+    ) -> float:
+        """Record token usage and return the cost of this request in USD."""
+        self.total_input_tokens += input_tokens
+        self.total_output_tokens += output_tokens
+        in_price, out_price = self._get_pricing()
+        request_cost = (
+            input_tokens * in_price
+            + output_tokens * out_price
+            + cache_read_tokens * in_price * 0.10
+            + cache_creation_tokens * in_price * 1.25
+        ) / 1_000_000
+        self.total_cost_usd += request_cost
+        return request_cost
+
+    def is_budget_exceeded(self) -> bool:
+        """Return True if a budget is set and has been exceeded."""
+        if self.max_budget_usd is None:
+            return False
+        return self.total_cost_usd > self.max_budget_usd
+
+    def remaining_budget(self) -> float | None:
+        """Return remaining budget in USD, or None if no budget is set."""
+        if self.max_budget_usd is None:
+            return None
+        return max(0.0, self.max_budget_usd - self.total_cost_usd)
+
+    def check_budget(self) -> None:
+        """Raise BudgetExceededError if the budget has been exceeded."""
+        if self.is_budget_exceeded():
+            raise BudgetExceededError(
+                spent=self.total_cost_usd,
+                budget=self.max_budget_usd,  # type: ignore[arg-type]
+            )
+
+    def _get_pricing(self) -> tuple[float, float]:
+        # 1. User custom pricing (exact match)
+        if self.custom_pricing:
+            if self.model in self.custom_pricing:
+                p = self.custom_pricing[self.model]
+                return (p[0], p[1]) if isinstance(p, list) else (0.0, 0.0)
+            # Partial match in custom
+            for key, p in self.custom_pricing.items():
+                if key != "default" and key in self.model:
+                    return (p[0], p[1]) if isinstance(p, list) else (0.0, 0.0)
+            # Custom default
+            if "default" in self.custom_pricing:
+                p = self.custom_pricing["default"]
+                return (p[0], p[1]) if isinstance(p, list) else (0.0, 0.0)
+
+        # 2. Built-in pricing (exact match)
+        if self.model in BUILTIN_PRICING:
+            return BUILTIN_PRICING[self.model]
+
+        # 3. Built-in pricing (partial match)
+        model_lower = self.model.lower()
+        for key, pricing in BUILTIN_PRICING.items():
+            if key in model_lower:
+                return pricing
+
+        # 4. Unknown model = free
+        return (0.0, 0.0)
+
+    def format_cost(self) -> str:
+        lines = [f"Tokens — in: {self.total_input_tokens:,}  out: {self.total_output_tokens:,}"]
+        in_price, out_price = self._get_pricing()
+        if self.total_cost_usd > 0.0001:
+            lines.append(f"  Cost: ${self.total_cost_usd:.4f}")
+            lines.append(f"  Rate: ${in_price}/1M in · ${out_price}/1M out")
+        elif in_price == 0 and out_price == 0:
+            lines.append("  Cost: $0 (free / local model)")
+        else:
+            lines.append(f"  Cost: ${self.total_cost_usd:.6f}")
+        return "  ".join(lines)

llm_code/runtime/dream.py

@@ -0,0 +1,127 @@
+"""DreamTask — background memory consolidation via LLM summarization."""
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from filelock import FileLock
+
+from llm_code.api.types import Message, MessageRequest, TextBlock
+
+if TYPE_CHECKING:
+    from llm_code.api.provider import LLMProvider
+    from llm_code.runtime.config import RuntimeConfig
+    from llm_code.runtime.memory import MemoryStore
+    from llm_code.runtime.session import Session
+
+logger = logging.getLogger(__name__)
+
+_CONSOLIDATION_SYSTEM_PROMPT = """\
+You are a memory consolidation agent. Given a conversation transcript from a \
+coding session, produce a structured Markdown summary with these sections:
+
+## Summary
+1-3 sentence overview of what was accomplished.
+
+## Modified Files
+Bulleted list of files that were created, edited, or deleted.
+
+## Decisions
+Key architectural or design decisions made during the session.
+
+## Patterns
+Reusable patterns, idioms, or techniques worth remembering for future sessions.
+
+## Open Items
+Any unfinished work, known issues, or next steps mentioned.
+
+Be concise. Focus on facts. Do not invent information not present in the transcript.
+"""
+
+
+class DreamTask:
+    """Consolidates a session's conversation into a structured summary via LLM."""
+
+    async def consolidate(
+        self,
+        session: "Session",
+        memory_store: "MemoryStore",
+        provider: "LLMProvider",
+        config: "RuntimeConfig",
+    ) -> str:
+        """Run LLM-powered consolidation on the session.
+
+        Returns the generated summary string, or empty string if skipped/failed.
+        """
+        dream_config = config.dream
+
+        # Guard: disabled
+        if not dream_config.enabled:
+            return ""
+
+        # Guard: too few messages
+        user_messages = [m for m in session.messages if m.role == "user"]
+        if len(user_messages) < dream_config.min_turns:
+            return ""
+
+        # Build transcript from session messages
+        transcript_parts: list[str] = []
+        for msg in session.messages:
+            role_label = "User" if msg.role == "user" else "Assistant"
+            for block in msg.content:
+                if hasattr(block, "text"):
+                    transcript_parts.append(f"**{role_label}:** {block.text}")
+
+        transcript = "\n\n".join(transcript_parts)
+
+        # Call LLM
+        request = MessageRequest(
+            model=getattr(config, "model", ""),
+            messages=(
+                Message(
+                    role="user",
+                    content=(TextBlock(text=f"Consolidate this session:\n\n{transcript}"),),
+                ),
+            ),
+            system=_CONSOLIDATION_SYSTEM_PROMPT,
+            tools=(),
+            max_tokens=2048,
+            temperature=0.3,
+        )
+
+        try:
+            response = await provider.send_message(request)
+        except Exception as exc:
+            logger.warning("DreamTask consolidation failed: %s", exc)
+            return ""
+
+        # Extract text from response
+        summary_parts: list[str] = []
+        for block in response.content:
+            if hasattr(block, "text"):
+                summary_parts.append(block.text)
+        summary = "\n".join(summary_parts)
+
+        if not summary.strip():
+            return ""
+
+        # Write with file lock
+        today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
+        lock_path = memory_store.consolidated_dir / f"{today}.md.lock"
+        lock = FileLock(str(lock_path), timeout=5)
+        with lock:
+            memory_store.save_consolidated(summary, date_str=today)
+
+        # Update last-run timestamp
+        memory_store.store(
+            "_dream_last_run",
+            datetime.now(timezone.utc).isoformat(),
+        )
+
+        logger.info(
+            "DreamTask consolidated session to %s/%s.md",
+            memory_store.consolidated_dir,
+            today,
+        )
+        return summary

llm_code/runtime/file_protection.py

@@ -0,0 +1,150 @@
+"""FileProtector — guards sensitive files from accidental reads and writes."""
+from __future__ import annotations
+
+import fnmatch
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+
+# Glob patterns for dangerous / credential files.
+# Patterns are matched against the basename as well as the full path string.
+SENSITIVE_PATTERNS: tuple[str, ...] = (
+    ".env",
+    ".env.*",
+    "*.key",
+    "*.pem",
+    "*.p12",
+    "credentials.*",
+    "*secret*",
+    "id_rsa",
+    "id_rsa.*",
+    "id_ed25519",
+    "id_ed25519.*",
+    "token.json",
+    "*.keystore",
+    ".netrc",
+    ".pgpass",
+)
+
+# Patterns that are always blocked (never allow write)
+_BLOCK_PATTERNS: tuple[str, ...] = (
+    ".env",
+    ".env.*",
+    "*.key",
+    "*.pem",
+    "*.p12",
+    "id_rsa",
+    "id_rsa.*",
+    "id_ed25519",
+    "id_ed25519.*",
+    "*.keystore",
+    ".netrc",
+    ".pgpass",
+)
+
+# Path prefixes (expanded) that are always blocked
+_BLOCK_PATH_PREFIXES: tuple[str, ...] = (
+    os.path.expanduser("~/.ssh/"),
+    os.path.expanduser("~/.aws/"),
+    os.path.expanduser("~/.config/gcloud/"),
+)
+
+
+@dataclass(frozen=True)
+class FileProtectionResult:
+    """Result of a file-protection check."""
+
+    allowed: bool
+    reason: str
+    severity: Literal["block", "warn", "allow"]
+
+
+def _matches_any(name: str, patterns: tuple[str, ...]) -> bool:
+    """Return True if *name* matches at least one glob pattern."""
+    for pattern in patterns:
+        if fnmatch.fnmatch(name, pattern):
+            return True
+    return False
+
+
+def _is_under_blocked_prefix(path: str) -> bool:
+    """Return True if the resolved path sits under a always-blocked directory."""
+    try:
+        resolved = str(Path(path).resolve())
+    except Exception:
+        resolved = path
+    for prefix in _BLOCK_PATH_PREFIXES:
+        if resolved.startswith(prefix):
+            return True
+    return False
+
+
+def is_sensitive(path: str) -> bool:
+    """Return True if *path* matches any sensitive pattern.
+
+    Resolves symlinks before checking to prevent bypass via symlink indirection.
+    """
+    resolved = str(Path(path).resolve())
+    basename = Path(resolved).name
+    if _matches_any(basename, SENSITIVE_PATTERNS):
+        return True
+    if _is_under_blocked_prefix(resolved):
+        return True
+    return False
+
+
+def check_write(path: str) -> FileProtectionResult:
+    """Check whether writing to *path* should be allowed, warned about, or blocked.
+
+    Resolves symlinks before checking to prevent bypass via symlink indirection.
+
+    Rules
+    -----
+    - ``.env`` files and SSH/AWS/GCloud keys → **block** (never write secrets)
+    - Other sensitive patterns (``credentials.*``, ``*secret*``, ``token.json``, …)
+      → **warn** (needs user confirmation)
+    - Everything else → **allow**
+    """
+    resolved = str(Path(path).resolve())
+    basename = Path(resolved).name
+
+    # Blocked: critical credential / key files
+    if _matches_any(basename, _BLOCK_PATTERNS) or _is_under_blocked_prefix(resolved):
+        return FileProtectionResult(
+            allowed=False,
+            reason=f"Writing to '{path}' is blocked: file matches a sensitive credential pattern.",
+            severity="block",
+        )
+
+    # Warn: other sensitive patterns
+    if is_sensitive(path):
+        return FileProtectionResult(
+            allowed=True,
+            reason=(
+                f"Writing to '{path}' requires confirmation: "
+                "the file matches a sensitive data pattern."
+            ),
+            severity="warn",
+        )
+
+    return FileProtectionResult(allowed=True, reason="", severity="allow")
+
+
+def check_read(path: str) -> FileProtectionResult:
+    """Check whether reading *path* should be allowed or warned about.
+
+    Sensitive file reads are warned (content may leak to LLM context).
+    """
+    if is_sensitive(path):
+        return FileProtectionResult(
+            allowed=True,
+            reason=(
+                f"Reading '{path}' may expose sensitive data to the LLM context. "
+                "Proceed with caution."
+            ),
+            severity="warn",
+        )
+
+    return FileProtectionResult(allowed=True, reason="", severity="allow")

llm_code/runtime/hardware.py

@@ -0,0 +1,85 @@
+"""VRAM and system memory detection across platforms."""
+from __future__ import annotations
+
+import subprocess
+import sys
+
+
+def detect_vram_gb() -> float | None:
+    """Detect available VRAM/memory in GB.
+
+    Detection chain (first success wins):
+    1. NVIDIA GPU via nvidia-smi
+    2. Apple Silicon unified memory via sysctl (× 0.75)
+    3. Linux /proc/meminfo (× 0.5)
+    4. None if all fail
+    """
+    result = _detect_nvidia()
+    if result is not None:
+        return result
+
+    result = _detect_apple_silicon()
+    if result is not None:
+        return result
+
+    result = _detect_linux_meminfo()
+    if result is not None:
+        return result
+
+    return None
+
+
+def _detect_nvidia() -> float | None:
+    """Detect NVIDIA GPU VRAM via nvidia-smi."""
+    try:
+        proc = subprocess.run(
+            ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
+            capture_output=True,
+            text=True,
+            timeout=2,
+        )
+        if proc.returncode != 0:
+            return None
+        first_line = proc.stdout.strip().split("\n")[0].strip()
+        mib = float(first_line.replace(" MiB", ""))
+        return mib / 1024.0
+    except (FileNotFoundError, OSError, ValueError, IndexError):
+        return None
+
+
+def _detect_apple_silicon() -> float | None:
+    """Detect Apple Silicon unified memory via sysctl."""
+    if sys.platform != "darwin":
+        return None
+    try:
+        proc = subprocess.run(
+            ["sysctl", "-n", "hw.memsize"],
+            capture_output=True,
+            text=True,
+            timeout=2,
+        )
+        if proc.returncode != 0:
+            return None
+        mem_bytes = int(proc.stdout.strip())
+        gb = mem_bytes / (1024**3)
+        return gb * 0.75
+    except (FileNotFoundError, OSError, ValueError):
+        return None
+
+
+def _detect_linux_meminfo() -> float | None:
+    """Detect total RAM from /proc/meminfo on Linux."""
+    if sys.platform != "linux":
+        return None
+    try:
+        with open("/proc/meminfo") as f:
+            content = f.read()
+        for line in content.split("\n"):
+            if line.startswith("MemTotal:"):
+                parts = line.split()
+                kb = int(parts[1])
+                gb = kb / (1024**2)
+                return gb * 0.5
+        return None
+    except (FileNotFoundError, OSError, ValueError):
+        return None