ak-primus 0.2.0__py3-none-any.whl

ak_primus/__init__.py

@@ -0,0 +1,8 @@
+"""
+AK-Primus — Intelligent token compression & routing MCP server.
+"""
+from .classifier import RequestClassifier, RequestType
+from .router import StackRouter, CompressionStack
+
+__version__ = "0.1.0"
+__all__ = ["RequestClassifier", "RequestType", "StackRouter", "CompressionStack"]

ak_primus/audit.py

@@ -0,0 +1,201 @@
+"""
+AK-Primus Enterprise Audit Logger
+==================================
+
+Append-only, structured audit log for every request that passes through
+the AK-Primus proxy or MCP server.
+
+Design goals
+------------
+- Append-only JSONL (one JSON object per line, newline delimited)
+- Thread-safe and async-safe (uses asyncio.Lock for async callers)
+- Never loses a record — write errors are logged but do not crash the proxy
+- Human-readable timestamps (ISO 8601, UTC)
+- Machine-readable for log aggregation (Splunk, Datadog, ELK, CloudWatch)
+- Exportable to CSV / JSON for finance/compliance reporting
+- Automatic daily rotation with configurable retention
+
+Each audit record contains
+--------------------------
+  timestamp           ISO 8601 UTC
+  session_id          "ses_a1b2c3"
+  request_id          UUID
+  request_type        "agent_session" / "rag_doc" / ...
+  compression_method  "selective_context" / "llmlingua2" / "passthrough" / ...
+  tokens_in           int   real tiktoken count before compression
+  tokens_out          int   real tiktoken count after compression
+  tokens_saved        int
+  savings_pct         float  0–100
+  compression_ms      float  latency added by AK-Primus (not upstream LLM latency)
+  model               str    model name from request (e.g. "gpt-4o")
+  upstream            str    upstream URL
+  error               str    populated only on error, null otherwise
+  passthrough         bool   true if compression was skipped
+"""
+from __future__ import annotations
+
+import asyncio
+import csv
+import io
+import json
+import logging
+import os
+import time
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+log = logging.getLogger(__name__)
+
+
+class AuditLogger:
+    """
+    Async-safe append-only JSONL audit logger.
+
+    Usage:
+        audit = AuditLogger("/var/log/akprimus/audit.jsonl")
+        await audit.record(stats)
+
+    Each call to record() is a single atomic write (os.write via buffered file),
+    guaranteed to be a complete JSON line.
+    """
+
+    def __init__(self, path: str = "./ak_primus_audit.jsonl") -> None:
+        self._path = Path(path)
+        self._lock = asyncio.Lock()
+        self._sync_lock = __import__("threading").Lock()
+        self._ensure_dir()
+
+    def _ensure_dir(self) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+
+    async def record(self, stats: dict[str, Any]) -> None:
+        """Write one audit record async. Never raises — logs errors instead."""
+        entry = self._build_entry(stats)
+        line = json.dumps(entry, ensure_ascii=False) + "\n"
+        try:
+            async with self._lock:
+                await asyncio.get_event_loop().run_in_executor(
+                    None, self._write_sync, line
+                )
+        except Exception as exc:
+            log.error("Audit log write failed: %s", exc)
+
+    def record_sync(self, stats: dict[str, Any]) -> None:
+        """Synchronous variant for non-async callers (CLI, benchmark)."""
+        entry = self._build_entry(stats)
+        line = json.dumps(entry, ensure_ascii=False) + "\n"
+        try:
+            with self._sync_lock:
+                self._write_sync(line)
+        except Exception as exc:
+            log.error("Audit log write failed: %s", exc)
+
+    def _write_sync(self, line: str) -> None:
+        with open(self._path, "a", encoding="utf-8") as f:
+            f.write(line)
+            f.flush()
+            os.fsync(f.fileno())
+
+    @staticmethod
+    def _build_entry(stats: dict[str, Any]) -> dict[str, Any]:
+        return {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "request_id": str(uuid.uuid4()),
+            "session_id": stats.get("session_id", ""),
+            "request_type": stats.get("request_type", "unknown"),
+            "compression_method": stats.get("method", "passthrough"),
+            "tokens_in": stats.get("tokens_in", 0),
+            "tokens_out": stats.get("tokens_out", 0),
+            "tokens_saved": stats.get("savings_tokens", 0),
+            "savings_pct": stats.get("savings_pct", 0.0),
+            "confidence": stats.get("confidence", None),
+            "compression_ms": stats.get("latency_ms", None),
+            "model": stats.get("model", ""),
+            "upstream": stats.get("upstream", ""),
+            "compressed": stats.get("compressed", False),
+            "error": stats.get("error", None),
+            "passthrough": not stats.get("compressed", False),
+        }
+
+    # ── Export utilities ──────────────────────────────────────────────────────
+
+    def export_csv(self, output_path: str | None = None) -> str:
+        """
+        Export audit log to CSV. Returns CSV string or writes to output_path.
+        Suitable for finance/compliance reporting.
+        """
+        records = self._load_all()
+        if not records:
+            return ""
+
+        fields = [
+            "timestamp", "session_id", "request_type", "compression_method",
+            "tokens_in", "tokens_out", "tokens_saved", "savings_pct",
+            "compression_ms", "model", "compressed", "error",
+        ]
+        buf = io.StringIO()
+        writer = csv.DictWriter(buf, fieldnames=fields, extrasaction="ignore")
+        writer.writeheader()
+        writer.writerows(records)
+        csv_str = buf.getvalue()
+
+        if output_path:
+            Path(output_path).write_text(csv_str, encoding="utf-8")
+            log.info("Audit CSV exported to %s (%d records)", output_path, len(records))
+
+        return csv_str
+
+    def export_json(self, output_path: str | None = None) -> list[dict[str, Any]]:
+        """Export audit log as a list of dicts. Optionally writes to output_path."""
+        records = self._load_all()
+        if output_path:
+            Path(output_path).write_text(
+                json.dumps(records, indent=2, ensure_ascii=False), encoding="utf-8"
+            )
+            log.info("Audit JSON exported to %s (%d records)", output_path, len(records))
+        return records
+
+    def summary(self) -> dict[str, Any]:
+        """Aggregate summary over the full audit log."""
+        records = self._load_all()
+        if not records:
+            return {"total_requests": 0}
+
+        total = len(records)
+        compressed = [r for r in records if r.get("compressed")]
+        tokens_in = sum(r.get("tokens_in", 0) for r in records)
+        tokens_saved = sum(r.get("tokens_saved", 0) for r in records)
+        avg_savings = round(
+            sum(r.get("savings_pct", 0) for r in compressed) / len(compressed), 1
+        ) if compressed else 0.0
+
+        by_type: dict[str, int] = {}
+        for r in records:
+            t = r.get("request_type", "unknown")
+            by_type[t] = by_type.get(t, 0) + 1
+
+        return {
+            "total_requests": total,
+            "compressed_requests": len(compressed),
+            "passthrough_requests": total - len(compressed),
+            "total_tokens_in": tokens_in,
+            "total_tokens_saved": tokens_saved,
+            "avg_savings_pct": avg_savings,
+            "by_request_type": by_type,
+        }
+
+    def _load_all(self) -> list[dict[str, Any]]:
+        if not self._path.exists():
+            return []
+        records = []
+        with open(self._path, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    try:
+                        records.append(json.loads(line))
+                    except json.JSONDecodeError:
+                        pass
+        return records