contextpilot-ai 0.1.0__py3-none-any.whl

contextpilot/__init__.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Any
+
+from contextpilot.config import ContextPilotConfig
+from contextpilot.pipeline import Pipeline
+
+
+def wrap(client: Any, config: ContextPilotConfig | dict | None = None) -> Any:
+    """Wrap an OpenAI or Anthropic client with ContextPilot compression.
+
+    Usage:
+        from openai import OpenAI
+        import contextpilot
+        pilot = contextpilot.wrap(OpenAI())
+        response = pilot.chat.completions.create(model="gpt-4o", messages=messages)
+    """
+    if isinstance(config, dict):
+        cfg = ContextPilotConfig.model_validate(config)
+    elif config is None:
+        cfg = ContextPilotConfig.load()
+    else:
+        cfg = config
+
+    pipeline = Pipeline(cfg)
+
+    module = type(client).__module__
+    name = type(client).__name__
+
+    if "openai" in module.lower() or name in ("OpenAI", "AsyncOpenAI"):
+        from contextpilot.adapters.openai_adapter import OpenAIWrapper
+        return OpenAIWrapper(client, pipeline)
+
+    if "anthropic" in module.lower() or name in ("Anthropic", "AsyncAnthropic"):
+        from contextpilot.adapters.anthropic_adapter import AnthropicWrapper
+        return AnthropicWrapper(client, pipeline)
+
+    raise ValueError(
+        f"Unsupported client type '{name}'. Supported: openai.OpenAI, anthropic.Anthropic"
+    )
+
+
+__all__ = ["wrap", "ContextPilotConfig"]

contextpilot/adapters/anthropic_adapter.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from contextpilot.pipeline import Pipeline
+
+
+class _MessagesWrapper:
+    def __init__(self, messages_api: object, pipeline: Pipeline) -> None:
+        self._messages_api = messages_api
+        self._pipeline = pipeline
+
+    def create(self, *, model: str, messages: list[dict], **kwargs: object) -> object:
+        system: str | None = kwargs.pop("system", None)  # type: ignore[assignment]
+        optimized_msgs, optimized_sys, _ = self._pipeline.optimize(
+            messages, system=system, provider="anthropic", model=model
+        )
+        if optimized_sys is not None:
+            kwargs["system"] = optimized_sys
+        return self._messages_api.create(model=model, messages=optimized_msgs, **kwargs)  # type: ignore[union-attr]
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._messages_api, name)
+
+
+class AnthropicWrapper:
+    """FR-001: Drop-in wrapper for anthropic.Anthropic.
+
+    Intercepts messages.create() calls, runs the compression pipeline on the
+    message list and optional system prompt, and forwards the optimised payload.
+    All other attributes delegate to the original client unchanged.
+    """
+
+    def __init__(self, client: object, pipeline: Pipeline) -> None:
+        self._client = client
+        self._pipeline = pipeline
+        self.messages = _MessagesWrapper(client.messages, pipeline)  # type: ignore[union-attr]
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._client, name)

contextpilot/adapters/openai_adapter.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from contextpilot.pipeline import Pipeline
+
+
+class _CompletionsWrapper:
+    def __init__(self, completions: object, pipeline: Pipeline) -> None:
+        self._completions = completions
+        self._pipeline = pipeline
+
+    def create(self, *, model: str, messages: list[dict], **kwargs: object) -> object:
+        optimized, _, _ = self._pipeline.optimize(messages, provider="openai", model=model)
+        return self._completions.create(model=model, messages=optimized, **kwargs)  # type: ignore[union-attr]
+
+
+class _ChatWrapper:
+    def __init__(self, chat: object, pipeline: Pipeline) -> None:
+        self._chat = chat
+        self.completions = _CompletionsWrapper(chat.completions, pipeline)  # type: ignore[union-attr]
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._chat, name)
+
+
+class OpenAIWrapper:
+    """FR-001: Drop-in wrapper for openai.OpenAI.
+
+    Intercepts chat.completions.create() calls, runs the compression pipeline,
+    and forwards the optimised payload. All other attributes delegate to the
+    original client unchanged.
+    """
+
+    def __init__(self, client: object, pipeline: Pipeline) -> None:
+        self._client = client
+        self._pipeline = pipeline
+        self.chat = _ChatWrapper(client.chat, pipeline)  # type: ignore[union-attr]
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._client, name)

contextpilot/analyzer.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Sequence
+
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
+from contextpilot.config import ContextPilotConfig
+
+
+class BlockClass(str, Enum):
+    ESSENTIAL = "essential"
+    COMPRESSIBLE = "compressible"
+    DROPPABLE = "droppable"
+
+
+@dataclass
+class MessageBlock:
+    index: int
+    role: str
+    content: str
+    staleness: float       # 0.0 = fresh (recent), 1.0 = stale (old)
+    redundancy: float      # 0.0 = unique, 1.0 = duplicate of another block
+    relevance: float       # 0.0 = irrelevant, 1.0 = highly relevant to latest query
+    density: float         # 0.0 = sparse, 1.0 = information-dense
+    classification: BlockClass = BlockClass.ESSENTIAL
+    token_count: int = 0   # approximate word-based count
+
+    @property
+    def composite_score(self) -> float:
+        """Value of this block — higher means keep it. Used for triage decisions."""
+        return (
+            self.relevance * 0.4
+            + self.density * 0.3
+            + (1.0 - self.staleness) * 0.2
+            + (1.0 - self.redundancy) * 0.1
+        )
+
+
+def _word_tokens(text: str) -> list[str]:
+    return re.findall(r"\b\w+\b", text.lower())
+
+
+def _density(text: str) -> float:
+    tokens = _word_tokens(text)
+    if not tokens:
+        return 0.0
+    return len(set(tokens)) / len(tokens)
+
+
+def _word_count(text: str) -> int:
+    return len(text.split())
+
+
+class Analyzer:
+    """FR-002: Context analysis engine.
+
+    Scores each message block across four dimensions — staleness, redundancy,
+    relevance, density — and classifies it as essential / compressible / droppable.
+    Analysis target: < 50 ms for up to 100K tokens (technical doc §8).
+    """
+
+    def __init__(self, config: ContextPilotConfig) -> None:
+        self.config = config
+
+    def analyze(
+        self, messages: list[dict], system: str | None = None
+    ) -> list[MessageBlock]:
+        n = len(messages)
+        if n == 0:
+            return []
+
+        texts = [m.get("content") or "" for m in messages]
+
+        # Most recent user message drives relevance scoring
+        recent_user = next(
+            (m.get("content") or "" for m in reversed(messages) if m.get("role") == "user"),
+            "",
+        )
+
+        redundancies, relevances = self._score_tfidf(texts, recent_user, n)
+
+        blocks: list[MessageBlock] = []
+        for i, msg in enumerate(messages):
+            content = texts[i]
+            role = msg.get("role", "user")
+            # Index 0 = oldest = most stale
+            staleness = 1.0 - (i / max(n - 1, 1))
+            density = _density(content)
+            redundancy = float(redundancies[i])
+            relevance = float(relevances[i])
+
+            classification = self._classify(
+                role=role,
+                index=i,
+                total=n,
+                staleness=staleness,
+                redundancy=redundancy,
+                relevance=relevance,
+            )
+
+            blocks.append(
+                MessageBlock(
+                    index=i,
+                    role=role,
+                    content=content,
+                    staleness=staleness,
+                    redundancy=redundancy,
+                    relevance=relevance,
+                    density=density,
+                    classification=classification,
+                    token_count=_word_count(content),
+                )
+            )
+
+        return blocks
+
+    def _score_tfidf(
+        self, texts: list[str], recent_user: str, n: int
+    ) -> tuple[list[float], list[float]]:
+        corpus = texts + ([recent_user] if recent_user else [])
+        try:
+            vec = TfidfVectorizer(min_df=1, token_pattern=r"(?u)\b\w+\b")
+            tfidf = vec.fit_transform(corpus)
+            msg_vecs = tfidf[:n]
+
+            # Redundancy: max cosine similarity against all *other* messages
+            sim = cosine_similarity(msg_vecs)
+            np.fill_diagonal(sim, 0.0)
+            redundancies = sim.max(axis=1).tolist()
+
+            # Relevance: cosine similarity to most recent user message
+            if recent_user:
+                user_vec = tfidf[n]
+                relevances = cosine_similarity(msg_vecs, user_vec.reshape(1, -1)).flatten().tolist()
+            else:
+                relevances = [0.5] * n
+        except Exception:
+            redundancies = [0.0] * n
+            relevances = [0.5] * n
+
+        return redundancies, relevances
+
+    def _classify(
+        self,
+        role: str,
+        index: int,
+        total: int,
+        staleness: float,
+        redundancy: float,
+        relevance: float,
+    ) -> BlockClass:
+        # System messages and the most recent window are always kept
+        if role == "system":
+            return BlockClass.ESSENTIAL
+        window = self.config.compression.history_window
+        if index >= total - window:
+            return BlockClass.ESSENTIAL
+
+        level = self.config.compression.level
+        if level == "conservative":
+            # Only drop near-exact duplicates
+            if redundancy > 0.95:
+                return BlockClass.DROPPABLE
+            if redundancy > 0.7:
+                return BlockClass.COMPRESSIBLE
+            return BlockClass.ESSENTIAL
+
+        if level == "aggressive":
+            if redundancy > 0.6 or (staleness > 0.5 and relevance < 0.15):
+                return BlockClass.DROPPABLE
+            if staleness > 0.3 or redundancy > 0.3:
+                return BlockClass.COMPRESSIBLE
+            return BlockClass.ESSENTIAL
+
+        # balanced (default)
+        if redundancy > 0.85 or (staleness > 0.7 and relevance < 0.2):
+            return BlockClass.DROPPABLE
+        if staleness > 0.5 or redundancy > 0.5:
+            return BlockClass.COMPRESSIBLE
+        return BlockClass.ESSENTIAL

contextpilot/cli.py

@@ -0,0 +1,161 @@
+"""ContextPilot CLI — entry point for all four surfaces.
+
+Commands
+--------
+contextpilot proxy    — Surface B: local proxy server (FR-009)
+contextpilot migrate  — Surface D: AST-based migration agent (FR-011)
+contextpilot report   — Show local token savings summary
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import click
+
+_LOCAL_LOG = Path.home() / ".contextpilot" / "events.jsonl"
+
+# Rough $/1M-token rates for common models (input side).
+_PRICING: dict[str, float] = {
+    "gpt-4o": 5.00,
+    "gpt-4o-mini": 0.15,
+    "gpt-4-turbo": 10.00,
+    "gpt-4": 30.00,
+    "gpt-3.5-turbo": 0.50,
+    "claude-opus": 15.00,
+    "claude-sonnet": 3.00,
+    "claude-haiku": 0.25,
+}
+_DEFAULT_RATE = 5.00  # $/1M tokens fallback
+
+
+def _rate_for(model: str) -> float:
+    m = model.lower()
+    for key, rate in _PRICING.items():
+        if key in m:
+            return rate
+    return _DEFAULT_RATE
+
+
+@click.group()
+@click.version_option(package_name="contextpilot")
+def main() -> None:
+    """ContextPilot — Intelligent LLM context optimization middleware."""
+
+
+# ---------------------------------------------------------------------------
+# Surface B: proxy
+# ---------------------------------------------------------------------------
+
+@main.command()
+@click.option("--port", default=8432, show_default=True, help="TCP port to listen on.")
+@click.option("--host", default="127.0.0.1", show_default=True, help="Address to bind.")
+@click.option("--config", "config_path", default=None, help="Path to contextpilot.yaml.")
+def proxy(port: int, host: str, config_path: str | None) -> None:
+    """Start the local proxy server (Surface B, FR-009).
+
+    \b
+    Set one environment variable to route requests through ContextPilot:
+
+      Anthropic / Claude Code:
+        export ANTHROPIC_BASE_URL=http://localhost:8432
+
+      OpenAI / GPT Codex / Aider:
+        export OPENAI_BASE_URL=http://localhost:8432/v1
+    """
+    from contextpilot.proxy import run_proxy
+    run_proxy(host=host, port=port, config_path=config_path)
+
+
+# ---------------------------------------------------------------------------
+# Surface D: migrate
+# ---------------------------------------------------------------------------
+
+@main.command()
+@click.argument("path", default=".", metavar="PATH")
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    default=False,
+    help="Show what would change without writing files (default if neither flag is given).",
+)
+@click.option("--apply", "apply_changes", is_flag=True, default=False, help="Write changes to files.")
+@click.option("--config", "config_path", default=None, help="Path to contextpilot.yaml.")
+def migrate(path: str, dry_run: bool, apply_changes: bool, config_path: str | None) -> None:
+    """Wrap existing LLM API calls with ContextPilot (Surface D, FR-011).
+
+    \b
+    Scans PATH (file or directory) for OpenAI and Anthropic client
+    instantiations and wraps them with contextpilot.wrap():
+
+      contextpilot migrate ./src/ --dry-run   # preview changes
+      contextpilot migrate ./src/ --apply      # rewrite files
+    """
+    from contextpilot.migrate import MigrationAgent
+
+    # Default to dry-run when neither flag is explicitly given
+    effective_dry_run = dry_run or not apply_changes
+
+    agent = MigrationAgent(config_path=config_path)
+    agent.run(path=path, dry_run=effective_dry_run, apply=apply_changes)
+
+
+# ---------------------------------------------------------------------------
+# Report: local savings summary
+# ---------------------------------------------------------------------------
+
+@main.command()
+@click.option("--tail", default=0, help="Show only the last N events (0 = all).")
+def report(tail: int) -> None:
+    """Show token savings from the local event log (~/.contextpilot/events.jsonl)."""
+    if not _LOCAL_LOG.exists():
+        click.echo("No events recorded yet. Run a few API calls through ContextPilot first.")
+        return
+
+    events: list[dict] = []
+    with _LOCAL_LOG.open(encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                try:
+                    events.append(json.loads(line))
+                except json.JSONDecodeError:
+                    pass
+
+    if not events:
+        click.echo("Event log is empty.")
+        return
+
+    if tail:
+        events = events[-tail:]
+
+    total = len(events)
+    fallbacks = sum(1 for e in events if e.get("fallback_triggered"))
+    orig_tokens = sum(e.get("tokens_input_original", 0) for e in events)
+    comp_tokens = sum(e.get("tokens_input_compressed", 0) for e in events)
+    saved_tokens = orig_tokens - comp_tokens
+    avg_quality = sum(e.get("quality_score", 100) for e in events) / total
+
+    # Dollar savings estimate
+    saved_usd = sum(
+        (e.get("tokens_input_original", 0) - e.get("tokens_input_compressed", 0))
+        / 1_000_000
+        * _rate_for(e.get("model", ""))
+        for e in events
+    )
+
+    ratio = (saved_tokens / orig_tokens * 100) if orig_tokens else 0.0
+
+    click.echo()
+    click.echo("  ContextPilot — Savings Report")
+    click.echo("  " + "─" * 36)
+    click.echo(f"  Total calls logged   : {total:,}")
+    click.echo(f"  Fallback rate        : {fallbacks}/{total} ({fallbacks/total*100:.1f}%)")
+    click.echo(f"  Tokens in (original) : {orig_tokens:,}")
+    click.echo(f"  Tokens in (sent)     : {comp_tokens:,}")
+    click.echo(f"  Tokens saved         : {saved_tokens:,}  ({ratio:.1f}% reduction)")
+    click.echo(f"  Avg quality score    : {avg_quality:.1f}/100")
+    click.echo(f"  Est. cost saved      : ${saved_usd:.4f}")
+    click.echo()
+    click.echo(f"  Log: {_LOCAL_LOG}")
+    click.echo()

contextpilot/compressor.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from contextpilot.analyzer import MessageBlock
+from contextpilot.config import ContextPilotConfig
+from contextpilot.strategies.dedup import SystemPromptDeduplicator
+from contextpilot.strategies.history import summarize_old_turns
+from contextpilot.strategies.rag_pruner import prune_rag_chunks
+from contextpilot.strategies.structural import apply_structural_stripping
+
+
+class Compressor:
+    """FR-003: Multi-stage compression pipeline.
+
+    Applies strategies in order:
+      1. Conversation history summarization
+      2. RAG chunk pruning
+      3. Structural formatting stripping
+      4. System prompt deduplication (aggressive mode)
+
+    Each strategy is applied in-place on the message list. The compressor is
+    surface-agnostic — the same pipeline runs for the Python library, proxy,
+    and MCP server surfaces.
+    """
+
+    def __init__(self, config: ContextPilotConfig) -> None:
+        self.config = config
+        self._dedup = SystemPromptDeduplicator()
+
+    def compress(
+        self,
+        messages: list[dict],
+        blocks: list[MessageBlock],
+        system: str | None = None,
+    ) -> tuple[list[dict], str | None]:
+        if not messages:
+            return messages, system
+
+        query = next(
+            (m.get("content") or "" for m in reversed(messages) if m.get("role") == "user"),
+            "",
+        )
+
+        result = list(messages)
+
+        # Stage 1: history summarization
+        result = summarize_old_turns(result, blocks, self.config)
+
+        # Stage 2: RAG chunk pruning
+        if query:
+            result = prune_rag_chunks(result, query, self.config)
+
+        # Stage 3: structural stripping
+        result = apply_structural_stripping(result, self.config)
+
+        # Stage 4: system prompt dedup (aggressive only — changes provider-visible content)
+        compressed_system = system
+        if system and self.config.compression.level == "aggressive":
+            compressed_system = self._dedup.process(system, self.config)
+
+        return result, compressed_system
+
+    def reset(self) -> None:
+        """Reset stateful strategies (e.g. between independent sessions)."""
+        self._dedup.reset()

contextpilot/config.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import yaml
+from pydantic import BaseModel, Field
+
+
+class CompressionConfig(BaseModel):
+    level: str = "balanced"          # conservative | balanced | aggressive
+    quality_threshold: float = 85.0  # fallback below this score
+    history_window: int = 6          # keep last N turns verbatim
+    rag_relevance_min: float = 0.15  # drop RAG chunks below this TF-IDF score
+
+
+class ShadowTestingConfig(BaseModel):
+    enabled: bool = False
+    sample_rate: float = 0.05  # fraction of calls shadowed
+
+
+class TelemetryConfig(BaseModel):
+    enabled: bool = True
+    endpoint: str = "https://api.contextpilot.dev/v1/telemetry"
+    api_key: Optional[str] = None
+    flush_interval: int = 60   # seconds
+    flush_size: int = 100      # events
+
+
+class ContextPilotConfig(BaseModel):
+    compression: CompressionConfig = Field(default_factory=CompressionConfig)
+    shadow_testing: ShadowTestingConfig = Field(default_factory=ShadowTestingConfig)
+    telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig)
+
+    @classmethod
+    def load(cls, path: str | Path | None = None) -> "ContextPilotConfig":
+        """Load config from YAML file, then apply environment variable overrides."""
+        if path is None:
+            for candidate in ("contextpilot.yaml", "contextpilot.yml"):
+                if Path(candidate).exists():
+                    path = candidate
+                    break
+
+        data: dict = {}
+        if path and Path(str(path)).exists():
+            with open(path) as f:
+                data = yaml.safe_load(f) or {}
+
+        comp = data.setdefault("compression", {})
+        if val := os.getenv("CONTEXTPILOT_QUALITY_THRESHOLD"):
+            comp["quality_threshold"] = float(val)
+        if val := os.getenv("CONTEXTPILOT_COMPRESSION_LEVEL"):
+            comp["level"] = val
+        if val := os.getenv("CONTEXTPILOT_HISTORY_WINDOW"):
+            comp["history_window"] = int(val)
+
+        tele = data.setdefault("telemetry", {})
+        if val := os.getenv("CONTEXTPILOT_API_KEY"):
+            tele["api_key"] = val
+        if val := os.getenv("CONTEXTPILOT_TELEMETRY_ENDPOINT"):
+            tele["endpoint"] = val
+
+        return cls.model_validate(data)

contextpilot/middleware.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from contextpilot.config import ContextPilotConfig
+from contextpilot.strategies.agent_memory import compress_agent_handoff
+
+
+class AgentMemory:
+    """FR-008: Agent memory middleware for LangChain / CrewAI / AutoGen.
+
+    Compresses inter-agent context handoffs, targeting 70–90% token reduction
+    in agentic workflows (technical doc §3.5).
+
+    Usage:
+        from contextpilot.middleware import AgentMemory
+        memory = AgentMemory(compression_level="aggressive",
+                             preserve_keys=["final_answer", "tool_outputs"])
+        compressed = memory.compress_handoff(agent_a_output)
+        agent_b_output = agent_b.run(task, context=compressed)
+    """
+
+    def __init__(
+        self,
+        compression_level: str = "balanced",
+        preserve_keys: list[str] | None = None,
+        config: ContextPilotConfig | None = None,
+    ) -> None:
+        if config is None:
+            cfg_data = {"compression": {"level": compression_level}}
+            self._config = ContextPilotConfig.model_validate(cfg_data)
+        else:
+            self._config = config
+        self._preserve_keys = preserve_keys or []
+
+    def compress_handoff(self, agent_output: str) -> str:
+        """Compress an agent's output before passing it to the next agent."""
+        return compress_agent_handoff(agent_output, self._preserve_keys, self._config)