synapse-orch-ai 0.1.0

package/backend/core/usage_tracker.py

@@ -0,0 +1,417 @@
+"""
+LLM Usage & Cost Tracker
+------------------------
+Persists every LLM call's token counts, context size, and estimated cost
+to data/usage_logs.json using the pricing table in data/model_pricing.json.
+
+Actual token counts are sourced from API response objects where available
+(OpenAI, Anthropic, Gemini all surface usage metadata). Bedrock and Ollama
+fall back to a character-count heuristic (len / 4).
+"""
+import json
+import os
+import threading
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+from core.config import DATA_DIR
+
+# -------------------------------------------------------------
+# Paths
+# -------------------------------------------------------------
+USAGE_LOGS_FILE = os.path.join(DATA_DIR, "usage_logs.json")
+PRICING_FILE = os.path.join(DATA_DIR, "model_pricing.json")
+
+_lock = threading.Lock()
+
+# -------------------------------------------------------------
+# Pricing lookup
+# -------------------------------------------------------------
+
+def _load_pricing() -> dict:
+    """Load the flat model_pricing.json. Returns {} on any error."""
+    try:
+        if os.path.exists(PRICING_FILE):
+            with open(PRICING_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+    except Exception as e:
+        print(f"DEBUG usage_tracker: could not load pricing: {e}")
+    return {}
+
+
+def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    """Return estimated USD cost for a call. Returns 0.0 for unknown models."""
+    pricing = _load_pricing()
+
+    # Try the exact model string, then progressively shorter prefixes
+    # so e.g. "claude-sonnet-4-20250514" matches its exact key.
+    entry = pricing.get(model)
+    if entry is None:
+        # Fuzzy prefix match -- pick the longest prefix that fits
+        for key in sorted(pricing.keys(), key=len, reverse=True):
+            if model.startswith(key) or key.startswith(model.split("-")[0]):
+                entry = pricing[key]
+                break
+
+    if not entry:
+        return 0.0
+
+    input_cost = (input_tokens / 1_000_000) * entry.get("input_per_1m", 0.0)
+    output_cost = (output_tokens / 1_000_000) * entry.get("output_per_1m", 0.0)
+    return round(input_cost + output_cost, 8)
+
+
+def get_pricing_table() -> dict:
+    """Return the raw pricing table for the API."""
+    return _load_pricing()
+
+
+def save_pricing_table(table: dict) -> None:
+    """Overwrite model_pricing.json with an updated table."""
+    with _lock:
+        with open(PRICING_FILE, "w", encoding="utf-8") as f:
+            json.dump(table, f, indent=4)
+    print(f"DEBUG usage_tracker: pricing table saved ({len(table)} entries)", flush=True)
+
+
+# -------------------------------------------------------------
+# Token estimation fallback
+# -------------------------------------------------------------
+
+def estimate_tokens_from_text(text: str) -> int:
+    """Rough heuristic: 1 token ? 4 characters. Used when the API doesn't return usage."""
+    if not text:
+        return 0
+    return max(1, len(text) // 4)
+
+
+# -------------------------------------------------------------
+# Usage log persistence
+# -------------------------------------------------------------
+
+def _load_logs() -> list:
+    if not os.path.exists(USAGE_LOGS_FILE):
+        return []
+    try:
+        with open(USAGE_LOGS_FILE, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return []
+
+
+def _save_logs(logs: list):
+    with open(USAGE_LOGS_FILE, "w", encoding="utf-8") as f:
+        json.dump(logs, f, indent=2, ensure_ascii=False)
+
+
+def log_usage(
+    *,
+    model: str,
+    provider: str,
+    input_tokens: int,
+    output_tokens: int,
+    context_chars: int,
+    session_id: Optional[str] = None,
+    agent_id: Optional[str] = None,
+    source: str = "chat",           # "chat" | "orchestration"
+    run_id: Optional[str] = None,   # orchestration run id
+    tool_name: Optional[str] = None,  # tool called on this turn (if any)
+    latency_seconds: float = 0.0,
+):
+    """Append a usage record to usage_logs.json (thread-safe)."""
+    estimated_cost = calculate_cost(model, input_tokens, output_tokens)
+    record = {
+        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z",
+        "model": model,
+        "provider": provider,
+        "session_id": session_id or "unknown",
+        "agent_id": agent_id or "unknown",
+        "source": source,
+        "run_id": run_id,
+        "tool_name": tool_name,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": input_tokens + output_tokens,
+        "context_chars": context_chars,
+        "estimated_cost": estimated_cost,
+        "latency_seconds": round(latency_seconds, 2),
+    }
+    with _lock:
+        logs = _load_logs()
+        logs.append(record)
+        _save_logs(logs)
+    print(
+        f"DEBUG usage: {model} in={input_tokens} out={output_tokens} "
+        f"cost=${estimated_cost:.6f} session={session_id}",
+        flush=True,
+    )
+
+
+def log_compaction_event(
+    *,
+    stage: str,                        # "trim" | "llm_summary"
+    chars_before: int,
+    chars_after: int,
+    session_id: Optional[str] = None,
+    agent_id: Optional[str] = None,
+    archive_path: Optional[str] = None,
+    model: str = "",
+):
+    """Append a compaction event to usage_logs.json.
+
+    Tokens and cost are zero — the Stage-2 LLM call is already captured by the
+    regular log_usage() call inside llm_providers. This record exists solely for
+    observability: when did compaction fire, how much was saved, where is the archive.
+    """
+    chars_saved = chars_before - chars_after
+    reduction_pct = round(chars_saved / chars_before * 100) if chars_before > 0 else 0
+    record = {
+        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z",
+        "event_type": "compaction",
+        "source": "compaction",
+        "stage": stage,
+        "session_id": session_id or "unknown",
+        "agent_id": agent_id or "unknown",
+        "chars_before": chars_before,
+        "chars_after": chars_after,
+        "chars_saved": chars_saved,
+        "reduction_pct": reduction_pct,
+        "archive_path": archive_path,
+        "model": model,
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "total_tokens": 0,
+        "estimated_cost": 0.0,
+        "latency_seconds": 0.0,
+    }
+    with _lock:
+        logs = _load_logs()
+        logs.append(record)
+        _save_logs(logs)
+    archive_note = f" → archived: {archive_path}" if archive_path else ""
+    print(
+        f"DEBUG usage: [compaction/{stage}] "
+        f"{chars_before:,} → {chars_after:,} chars (-{reduction_pct}%){archive_note}",
+        flush=True,
+    )
+
+
+# -------------------------------------------------------------
+# Query helpers
+# -------------------------------------------------------------
+
+def get_usage_logs(
+    limit: int = 100,
+    offset: int = 0,
+    session_id: Optional[str] = None,
+    source: Optional[str] = None,
+    run_id: Optional[str] = None,
+) -> list:
+    """Return paginated usage records.
+    - When filtering by session_id or run_id: oldest-first (for per-turn context delta display).
+    - Otherwise: newest-first.
+    """
+    with _lock:
+        logs = _load_logs()
+
+    # Filter
+    if session_id:
+        logs = [r for r in logs if r.get("session_id") == session_id]
+    if run_id:
+        logs = [r for r in logs if r.get("run_id") == run_id]
+    if source:
+        logs = [r for r in logs if r.get("source") == source]
+
+    # Ordering: per-session/run ? chronological (oldest first); global ? newest first
+    if not session_id and not run_id:
+        logs = list(reversed(logs))
+
+    return logs[offset: offset + limit]
+
+
+def get_usage_summary() -> dict:
+    """Return aggregated cost/token totals grouped by model and session.
+
+    Orchestration log entries (those with a non-null run_id) are grouped
+    separately by run_id so that each orchestration run appears as a single
+    session entry regardless of how many sub-agents ran under it.
+    Chat sessions are grouped by session_id as before.
+    """
+    with _lock:
+        logs = _load_logs()
+
+    total_cost = 0.0
+    total_input = 0
+    total_output = 0
+    # Compaction events are observational (0 tokens/cost) — exclude from request count
+    total_requests = sum(1 for r in logs if r.get("event_type") != "compaction")
+    by_model: dict[str, dict] = {}
+    by_session: dict[str, dict] = {}   # keyed by session_id (chat)
+    by_run: dict[str, dict] = {}       # keyed by run_id (orchestration runs)
+    by_schedule: dict[str, dict] = {}  # keyed by run_id (schedule runs)
+
+    for r in logs:
+        if r.get("event_type") == "compaction":
+            continue  # purely observational — skip all aggregates
+        model = r.get("model", "unknown")
+        provider = r.get("provider", "unknown")
+        session = r.get("session_id", "unknown")
+        run_id = r.get("run_id")  # None for chat, e.g. "run_orch_X_ts" for orch
+        source = r.get("source", "chat")
+        cost = r.get("estimated_cost", 0.0)
+        inp = r.get("input_tokens", 0)
+        out = r.get("output_tokens", 0)
+        ctx = r.get("context_chars", 0)
+        agent_id = r.get("agent_id", "unknown")
+
+        total_cost += cost
+        total_input += inp
+        total_output += out
+
+        # By model
+        if model not in by_model:
+            by_model[model] = {
+                "model": model,
+                "provider": provider,
+                "requests": 0,
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "total_tokens": 0,
+                "estimated_cost": 0.0,
+            }
+        bm = by_model[model]
+        bm["requests"] += 1
+        bm["input_tokens"] += inp
+        bm["output_tokens"] += out
+        bm["total_tokens"] += inp + out
+        bm["estimated_cost"] = round(bm["estimated_cost"] + cost, 8)
+
+        # Schedule entries are grouped by run_id
+        if run_id and source == "schedule":
+            if run_id not in by_schedule:
+                # Extract schedule_id from run_id format: schedulerun_{schedule_id}_{ts}
+                parts = run_id.split("_")
+                schedule_id = parts[1] if len(parts) > 1 else "unknown"
+                by_schedule[run_id] = {
+                    "run_id": run_id,
+                    "schedule_id": schedule_id,
+                    "agent_id": agent_id,
+                    "agents_used": set(),
+                    "requests": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "context_chars": 0,
+                    "estimated_cost": 0.0,
+                    "models_used": set(),
+                    "first_ts": r.get("timestamp"),
+                    "last_ts": r.get("timestamp"),
+                    "source": "schedule",
+                }
+            bsch = by_schedule[run_id]
+            bsch["requests"] += 1
+            bsch["input_tokens"] += inp
+            bsch["output_tokens"] += out
+            bsch["total_tokens"] += inp + out
+            bsch["context_chars"] += ctx
+            bsch["estimated_cost"] = round(bsch["estimated_cost"] + cost, 8)
+            bsch["models_used"].add(model)
+            bsch["agents_used"].add(agent_id)
+            bsch["last_ts"] = r.get("timestamp")
+
+        # Orchestration entries are grouped by run_id
+        elif run_id and (source == "orchestration" or source.startswith("orchestration:")):
+            if run_id not in by_run:
+                by_run[run_id] = {
+                    "session_id": session,   # the chat session that spawned this run
+                    "run_id": run_id,
+                    "agent_id": agent_id,    # first agent seen (show in header)
+                    "agents_used": set(),
+                    "requests": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "context_chars": 0,
+                    "estimated_cost": 0.0,
+                    "models_used": set(),
+                    "first_ts": r.get("timestamp"),
+                    "last_ts": r.get("timestamp"),
+                    "source": "orchestration",
+                }
+            br = by_run[run_id]
+            br["requests"] += 1
+            br["input_tokens"] += inp
+            br["output_tokens"] += out
+            br["total_tokens"] += inp + out
+            br["context_chars"] += ctx
+            br["estimated_cost"] = round(br["estimated_cost"] + cost, 8)
+            br["models_used"].add(model)
+            br["agents_used"].add(agent_id)
+            br["last_ts"] = r.get("timestamp")
+        else:
+            # Chat / system-prompt-generation: group by session_id
+            if session not in by_session:
+                by_session[session] = {
+                    "session_id": session,
+                    "run_id": None,
+                    "agent_id": agent_id,
+                    "agents_used": set(),
+                    "requests": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "context_chars": 0,
+                    "estimated_cost": 0.0,
+                    "models_used": set(),
+                    "first_ts": r.get("timestamp"),
+                    "last_ts": r.get("timestamp"),
+                    "source": source,
+                }
+            bs = by_session[session]
+            bs["requests"] += 1
+            bs["input_tokens"] += inp
+            bs["output_tokens"] += out
+            bs["total_tokens"] += inp + out
+            bs["context_chars"] += ctx
+            bs["estimated_cost"] = round(bs["estimated_cost"] + cost, 8)
+            bs["models_used"].add(model)
+            bs["agents_used"].add(agent_id)
+            bs["last_ts"] = r.get("timestamp")
+
+    # Sort by cost descending
+    by_model_list = sorted(by_model.values(), key=lambda x: x["estimated_cost"], reverse=True)
+
+    # Merge chat sessions + orchestration runs, convert sets to lists
+    all_sessions = list(by_session.values()) + list(by_run.values())
+    by_session_list = []
+    for bs in sorted(all_sessions, key=lambda x: x.get("last_ts") or "", reverse=True):
+        bs["models_used"] = list(bs["models_used"])
+        bs["agents_used"] = list(bs["agents_used"])
+        by_session_list.append(bs)
+
+    # Schedule runs -- convert sets to lists, sort by last_ts descending
+    by_schedule_list = []
+    for bsch in sorted(by_schedule.values(), key=lambda x: x.get("last_ts") or "", reverse=True):
+        bsch["models_used"] = list(bsch["models_used"])
+        bsch["agents_used"] = list(bsch["agents_used"])
+        by_schedule_list.append(bsch)
+
+    return {
+        "total_cost": round(total_cost, 8),
+        "total_input_tokens": total_input,
+        "total_output_tokens": total_output,
+        "total_tokens": total_input + total_output,
+        "total_requests": total_requests,
+        "by_model": by_model_list,
+        "by_session": by_session_list,
+        "by_schedule": by_schedule_list,
+    }
+
+def clear_usage_logs() -> int:
+    """Delete all usage logs. Returns count deleted."""
+    with _lock:
+        count = len(_load_logs())
+        _save_logs([])
+    return count

package/backend/core/vault.py

@@ -0,0 +1,265 @@
+"""
+Vault: Automatically saves large tool outputs to files and provides tools to query them.
+Also resolves @[path] vault file mentions in user messages before they reach the LLM.
+
+When any tool returns more than VAULT_THRESHOLD characters, the output is saved to
+backend/data/vault/ and the LLM receives only the file path + metadata. The LLM can
+then use the Filesystem MCP tools (read_file, search_files) to access parts of the
+file without flooding its context window.
+"""
+import json
+import os
+import re
+from datetime import datetime
+from pathlib import Path
+
+from core.config import DATA_DIR
+VAULT_DIR = Path(DATA_DIR) / "vault" / "tool_outputs"
+VAULT_THRESHOLD = 100000  # characters (fallback default)
+
+
+def _make_vault_path(tool_name: str, ext: str) -> Path:
+    """Generate a unique, safe vault file path."""
+    VAULT_DIR.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:20]
+    safe_name = re.sub(r"[^\w]", "_", tool_name)[:40]
+    return VAULT_DIR / f"{safe_name}_{timestamp}.{ext}"
+
+
+def maybe_vault(tool_name: str, raw_output: str) -> str:
+    """
+    If raw_output exceeds the vault threshold (from settings), persist it to vault and return a
+    compact JSON reference the LLM can act on with the vault read/search tools.
+    Returns raw_output unchanged when under the threshold or vault is disabled.
+    """
+    from core.config import load_settings
+    settings = load_settings()
+    if not settings.get("vault_enabled", True):
+        return raw_output
+    threshold = settings.get("vault_threshold", VAULT_THRESHOLD)
+    if len(raw_output) <= threshold:
+        return raw_output
+
+    # Decide extension: JSON gets pretty-validated, everything else is text.
+    try:
+        parsed = json.loads(raw_output)
+        ext = "json"
+        content = json.dumps(parsed, indent=2, ensure_ascii=False)
+    except (json.JSONDecodeError, ValueError):
+        ext = "txt"
+        content = raw_output
+
+    path = _make_vault_path(tool_name, ext)
+    path.write_text(content, encoding="utf-8")
+
+    total_lines = content.count("\n") + 1
+    return json.dumps({
+        "vault_file": str(path),
+        "file_type": ext,
+        "size_chars": len(raw_output),
+        "total_lines": total_lines,
+        "message": (
+            f"Output too large ({len(raw_output):,} chars, {total_lines:,} lines). Saved to vault at: {path}. "
+            f"Instead: use read_file_by_lines with start_line/end_line to read a slice, "
+            f"or grep to search for specific values."
+        ),
+    })
+
+
+# ---------------------------------------------------------------------------
+# Vault tool implementations — called directly by react_engine.py
+# ---------------------------------------------------------------------------
+
+_VAULT_USER_DIR = Path(DATA_DIR) / "vault"
+
+
+def expand_vault_mentions(message: str) -> str:
+    """
+    Replace every @[relative/path] vault mention in the user message with the
+    file's actual content, so the LLM receives the data inline.
+
+    Example:
+        "Log a job. Config: @[jobs-filepath.json]"
+        →
+        "Log a job. Config: @[jobs-filepath.json]\n<file content here>"
+    """
+    def _replace(match: re.Match) -> str:
+        rel = match.group(1).strip()
+        # Prevent path traversal
+        try:
+            resolved = (_VAULT_USER_DIR / rel).resolve()
+            if not str(resolved).startswith(str(_VAULT_USER_DIR.resolve())):
+                return match.group(0)  # leave as-is if traversal attempt
+        except Exception:
+            return match.group(0)
+
+        if not resolved.exists() or not resolved.is_file():
+            return match.group(0)  # leave as-is if file not found
+
+        try:
+            content = resolved.read_text(encoding="utf-8")
+        except Exception:
+            return match.group(0)
+
+        return f"@[{rel}]\n```\n{content}\n```"
+
+    return re.sub(r"@\[([^\]]+)\]", _replace, message)
+
+
+def _safe_path(path: str) -> Path:
+    """Return Path, rejecting obvious traversal attempts."""
+    p = Path(path).resolve()
+    return p
+
+
+def tool_read_file_chunk(path: str, start_line: int, end_line: int) -> str:
+    """Read lines [start_line, end_line] (1-indexed, inclusive) from any file."""
+    try:
+        p = _safe_path(path)
+        if not p.exists():
+            return json.dumps({"error": f"File not found: {path}"})
+        lines = p.read_text(encoding="utf-8").splitlines()
+        total = len(lines)
+        s = max(1, start_line) - 1      # 0-indexed
+        e = min(end_line, total)
+        chunk = lines[s:e]
+        return json.dumps({
+            "path": path,
+            "start_line": s + 1,
+            "end_line": e,
+            "total_lines": total,
+            "content": "\n".join(chunk),
+        })
+    except Exception as ex:
+        return json.dumps({"error": str(ex)})
+
+
+def tool_search_file(path: str, query: str, context_lines: int = 5) -> str:
+    """Grep-like search: returns matching lines with ±context_lines of surrounding context."""
+    try:
+        p = _safe_path(path)
+        if not p.exists():
+            return json.dumps({"error": f"File not found: {path}"})
+        lines = p.read_text(encoding="utf-8").splitlines()
+        q = query.lower()
+        results = []
+        covered: set[int] = set()
+
+        for i, line in enumerate(lines):
+            if q not in line.lower():
+                continue
+            start = max(0, i - context_lines)
+            end = min(len(lines), i + context_lines + 1)
+            if i in covered:
+                continue
+            covered.update(range(start, end))
+            block = []
+            for j in range(start, end):
+                prefix = ">>>" if j == i else "   "
+                block.append(f"{prefix} [L{j + 1}] {lines[j]}")
+            results.append({
+                "match_line": i + 1,
+                "match": line,
+                "context": "\n".join(block),
+            })
+            if len(results) >= 20:
+                break
+
+        return json.dumps({
+            "path": path,
+            "query": query,
+            "matches_found": len(results),
+            "results": results,
+        })
+    except Exception as ex:
+        return json.dumps({"error": str(ex)})
+
+
+def tool_read_json_chunk(path: str, offset: int = 0, limit: int = 50) -> str:
+    """
+    Read a slice of a JSON vault file.
+    - Array root  → returns items[offset : offset+limit]
+    - Object root → returns keys[offset : offset+limit] with their values
+    """
+    try:
+        p = _safe_path(path)
+        if not p.exists():
+            return json.dumps({"error": f"File not found: {path}"})
+        data = json.loads(p.read_text(encoding="utf-8"))
+
+        if isinstance(data, list):
+            total = len(data)
+            chunk = data[offset: offset + limit]
+            return json.dumps({
+                "path": path,
+                "root_type": "array",
+                "total_items": total,
+                "offset": offset,
+                "limit": limit,
+                "returned": len(chunk),
+                "data": chunk,
+            }, default=str)
+
+        if isinstance(data, dict):
+            keys = list(data.keys())
+            total = len(keys)
+            chunk_keys = keys[offset: offset + limit]
+            chunk = {k: data[k] for k in chunk_keys}
+            return json.dumps({
+                "path": path,
+                "root_type": "object",
+                "total_keys": total,
+                "all_keys": keys,
+                "offset": offset,
+                "limit": limit,
+                "returned": len(chunk_keys),
+                "data": chunk,
+            }, default=str)
+
+        # Scalar / other
+        return json.dumps({
+            "path": path,
+            "root_type": type(data).__name__,
+            "data": data,
+        }, default=str)
+
+    except Exception as ex:
+        return json.dumps({"error": str(ex)})
+
+
+def tool_search_json(path: str, query: str) -> str:
+    """
+    Recursively search JSON values for query string.
+    Returns up to 20 matching {json_path, value} pairs.
+    """
+    try:
+        p = _safe_path(path)
+        if not p.exists():
+            return json.dumps({"error": f"File not found: {path}"})
+        data = json.loads(p.read_text(encoding="utf-8"))
+        q = query.lower()
+        results: list[dict] = []
+
+        def _recurse(obj, jpath: str):
+            if len(results) >= 20:
+                return
+            if isinstance(obj, dict):
+                for k, v in obj.items():
+                    _recurse(v, f"{jpath}.{k}")
+            elif isinstance(obj, list):
+                for i, item in enumerate(obj):
+                    _recurse(item, f"{jpath}[{i}]")
+            else:
+                val_str = str(obj)
+                if q in val_str.lower():
+                    results.append({"json_path": jpath, "value": val_str[:500]})
+
+        _recurse(data, "$")
+        return json.dumps({
+            "path": path,
+            "query": query,
+            "matches_found": len(results),
+            "results": results,
+        })
+    except Exception as ex:
+        return json.dumps({"error": str(ex)})