context-compress 0.1.0__tar.gz

context_compress-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+.env
+dist/
+*.egg-info/
+__pycache__/
+.coverage
+.pytest_cache/

context_compress-0.1.0/ARCHITECTURE.md

@@ -0,0 +1,597 @@
+# Context Compression Middleware — Architecture Spec
+
+> **Purpose:** This document is the complete build spec for a coding agent (Codex/Claude Code). Everything needed to implement is here. No ambiguity, no hand-waving.
+
+---
+
+## What This Is
+
+A Python library + CLI + HTTP API that compresses LLM context into hierarchical layers, supports incremental conversation compression, and works as drop-in middleware between any application and any LLM API.
+
+**One sentence:** You send us a big context, we give you back a smaller one that works just as well.
+
+---
+
+## Project Structure
+
+```
+context-compression/
+├── pyproject.toml              # Package config (use hatchling)
+├── README.md
+├── cctx/                       # Package name: cctx (context compression)
+│   ├── __init__.py             # Exports: Compressor, ConversationManager, CCTXClient
+│   ├── compressor.py           # Core hierarchical compression engine
+│   ├── conversation.py         # Incremental delta compression for conversations
+│   ├── scorer.py               # Sentence/chunk importance scoring
+│   ├── tokenizer.py            # Token counting (wraps tiktoken)
+│   ├── client.py               # HTTP client for the API server
+│   ├── server.py               # FastAPI server (the middleware API)
+│   ├── cli.py                  # CLI entry point (click)
+│   └── types.py                # Shared dataclasses/types
+├── tests/
+│   ├── test_compressor.py
+│   ├── test_conversation.py
+│   ├── test_scorer.py
+│   └── test_server.py
+└── examples/
+    ├── basic_compression.py
+    ├── conversation_demo.py
+    └── middleware_demo.py       # Shows drop-in usage between app and OpenAI
+```
+
+---
+
+## Core Types (`cctx/types.py`)
+
+```python
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+class Layer(Enum):
+    RAW = 0       # Full text, no compression
+    FACTS = 1     # Key facts & relationships extracted
+    SUMMARY = 2   # Executive summary
+
+@dataclass
+class CompressedResult:
+    """Output of compression at all 3 layers."""
+    layers: dict[Layer, str]            # Layer -> compressed text
+    tokens: dict[Layer, int]            # Layer -> token count
+    entities: list[str]                 # Named entities found
+    sections: list[Section]             # Drill-down sections
+    metadata: dict                      # Timing, scores, etc.
+
+@dataclass
+class Section:
+    """A chunk of the original text, addressable for drill-down."""
+    id: str                             # e.g. "s0", "s1", "s2"
+    text: str                           # Full raw text of this section
+    tokens: int                         # Token count
+    summary: str                        # One-sentence summary
+    importance: float                   # 0.0 - 1.0 score
+
+@dataclass
+class ConversationState:
+    """Persistent state for an ongoing conversation."""
+    id: str                             # Conversation identifier
+    compressed_history: str             # Compressed older messages
+    compressed_tokens: int
+    recent_messages: list[Message]      # Last N messages (raw)
+    total_raw_tokens: int               # Running total of all tokens ever added
+    compression_events: int             # How many times compression has fired
+
+@dataclass  
+class Message:
+    role: str                           # "user", "assistant", "system"
+    content: str
+    tokens: int
+
+@dataclass
+class CompressRequest:
+    """API request body."""
+    text: str
+    layer: Layer = Layer.FACTS          # Target compression layer
+    l1_ratio: float = 0.3              # Fraction of content to keep at Layer 1
+    l2_max_sentences: int = 3          # Max sentences for Layer 2
+    sections: bool = True               # Whether to generate drill-down sections
+
+@dataclass
+class ConversationRequest:
+    """API request for conversation compression."""
+    conversation_id: str
+    message: Message
+    compress_threshold: int = 1000      # Tokens before compression triggers
+    keep_recent: int = 5                # Number of recent messages to keep raw
+
+@dataclass
+class DrillDownRequest:
+    """API request to expand a section."""
+    section_id: str
+    target_layer: Layer = Layer.RAW     # What layer to return
+```
+
+---
+
+## Scoring Engine (`cctx/scorer.py`)
+
+This is the brain — decides what's important. Keep it pluggable.
+
+```python
+from abc import ABC, abstractmethod
+
+class Scorer(ABC):
+    """Base class for importance scoring strategies."""
+    
+    @abstractmethod
+    def score_sentences(self, sentences: list[str], query: str | None = None) -> list[float]:
+        """Return importance score (0.0-1.0) for each sentence."""
+        ...
+
+class TextRankScorer(Scorer):
+    """Graph-based sentence ranking. No external dependencies."""
+    # Implementation: sentence similarity matrix -> power iteration
+    # This is the DEFAULT scorer. Works offline, fast, no API calls.
+
+class TFIDFScorer(Scorer):
+    """TF-IDF based scoring. Better for document-style content."""
+    # Uses scikit-learn TfidfVectorizer
+    # Score = cosine similarity to document centroid (or query if provided)
+
+class QueryAwareScorer(Scorer):
+    """Scores sentences by relevance to a specific query/question."""
+    # Wraps another scorer but boosts sentences matching query terms
+    # Critical for RAG contexts where you know what question is being answered
+
+class CompositeScorer(Scorer):
+    """Combines multiple scorers with configurable weights."""
+    def __init__(self, scorers: list[tuple[Scorer, float]]):
+        # scorers = [(TextRankScorer(), 0.6), (TFIDFScorer(), 0.4)]
+        ...
+```
+
+**Why pluggable:** Later we can add an `LLMScorer` that uses a cheap model (GPT-4o-mini) to score importance. For now, keep it algorithmic and free.
+
+---
+
+## Compressor (`cctx/compressor.py`)
+
+```python
+class Compressor:
+    def __init__(self, scorer: Scorer | None = None):
+        """Default scorer is TextRankScorer."""
+        self.scorer = scorer or TextRankScorer()
+    
+    def compress(self, request: CompressRequest) -> CompressedResult:
+        """
+        Compress text into hierarchical layers.
+        
+        Algorithm:
+        1. Split text into sections (by paragraph/double-newline)
+        2. Split each section into sentences
+        3. Score all sentences with self.scorer
+        4. Layer 1: Keep top `l1_ratio` sentences, in original order
+        5. Layer 2: Keep top `l2_max_sentences` sentences, in original order
+        6. Extract entities (regex-based NER)
+        7. Generate section summaries (top sentence per section)
+        8. Return CompressedResult with all layers
+        """
+    
+    def drill_down(self, result: CompressedResult, section_id: str, 
+                   layer: Layer = Layer.RAW) -> str:
+        """
+        Retrieve a specific section at a specific layer.
+        Layer.RAW = full text, Layer.FACTS = compressed section.
+        """
+```
+
+**Key design decision:** The compressor is stateless. Give it text, get back compressed text. Conversation state management is separate.
+
+---
+
+## Conversation Manager (`cctx/conversation.py`)
+
+```python
+class ConversationManager:
+    """Manages incremental compression for ongoing conversations."""
+    
+    def __init__(self, compressor: Compressor | None = None,
+                 compress_threshold: int = 1000,
+                 keep_recent: int = 5):
+        self.compressor = compressor or Compressor()
+        self.conversations: dict[str, ConversationState] = {}
+    
+    def add_message(self, conversation_id: str, message: Message) -> dict:
+        """
+        Add a message to a conversation. Compress if threshold exceeded.
+        
+        Algorithm:
+        1. Append message to conversation's recent_messages
+        2. Count total tokens in (compressed_history + recent_messages)
+        3. If total > compress_threshold AND len(recent_messages) > keep_recent:
+           a. Take messages[:-keep_recent] (older ones)
+           b. Combine with existing compressed_history
+           c. Compress using self.compressor at Layer.FACTS
+           d. Replace compressed_history with result
+           e. Keep only messages[-keep_recent:] as recent
+        4. Return stats dict: {action, tokens_before, tokens_after, ...}
+        """
+    
+    def get_context(self, conversation_id: str) -> str:
+        """
+        Get the current context for LLM consumption.
+        
+        Format:
+        [Prior context summary]
+        {compressed_history}
+        [Recent messages]
+        {role}: {content}
+        {role}: {content}
+        ...
+        """
+    
+    def get_context_tokens(self, conversation_id: str) -> int:
+        """Token count of get_context() output."""
+    
+    def reset(self, conversation_id: str):
+        """Clear conversation state."""
+```
+
+**How conversations connect to compression:**
+- `ConversationManager` owns a `Compressor` instance
+- When compression triggers, it calls `compressor.compress()` on the older messages
+- The compressed output becomes the new `compressed_history`
+- Recent messages stay raw for full fidelity
+- This is the "git diff" approach — you never re-compress what's already compressed
+
+---
+
+## HTTP API Server (`cctx/server.py`)
+
+FastAPI server. Three endpoints. That's it.
+
+```python
+from fastapi import FastAPI
+
+app = FastAPI(title="cctx", version="0.1.0")
+
+# Shared instances
+compressor = Compressor()
+conversation_mgr = ConversationManager(compressor)
+
+@app.post("/v1/compress")
+async def compress(request: CompressRequest) -> CompressedResult:
+    """
+    Compress text into hierarchical layers.
+    
+    Request:
+    {
+        "text": "Your long text...",
+        "layer": 1,              # 0=raw, 1=facts, 2=summary (default: 1)
+        "l1_ratio": 0.3,         # optional
+        "l2_max_sentences": 3,   # optional
+        "sections": true         # optional, include drill-down sections
+    }
+    
+    Response:
+    {
+        "layers": {
+            "0": {"text": "...", "tokens": 1674},
+            "1": {"text": "...", "tokens": 689},
+            "2": {"text": "...", "tokens": 45}
+        },
+        "entities": ["Tesla", "Salesforce", ...],
+        "sections": [
+            {"id": "s0", "summary": "...", "tokens": 120, "importance": 0.92},
+            ...
+        ],
+        "metadata": {
+            "compression_ms": 12,
+            "scorer": "textrank"
+        }
+    }
+    """
+
+@app.post("/v1/conversation")
+async def conversation(request: ConversationRequest) -> dict:
+    """
+    Add a message to a conversation with automatic incremental compression.
+    
+    Request:
+    {
+        "conversation_id": "conv_123",
+        "message": {"role": "user", "content": "..."},
+        "compress_threshold": 1000,  # optional
+        "keep_recent": 5             # optional
+    }
+    
+    Response:
+    {
+        "context": "the full context string to send to your LLM",
+        "context_tokens": 569,
+        "total_raw_tokens": 1674,
+        "compression_ratio": 2.9,
+        "action": "compressed",      # or "appended"
+        "stats": { ... }
+    }
+    """
+
+@app.post("/v1/drill-down")
+async def drill_down(request: DrillDownRequest) -> dict:
+    """
+    Expand a section from a previous compression result.
+    
+    Request:
+    {
+        "section_id": "s0",
+        "target_layer": 0           # 0=raw, 1=facts
+    }
+    
+    Response:
+    {
+        "section_id": "s0",
+        "layer": 0,
+        "text": "full raw text of this section...",
+        "tokens": 245
+    }
+    """
+```
+
+---
+
+## CLI (`cctx/cli.py`)
+
+```python
+import click
+
+@click.group()
+def cli():
+    """cctx — context compression toolkit"""
+
+@cli.command()
+@click.argument("file", type=click.Path(exists=True))
+@click.option("--layer", "-l", type=int, default=1, help="Target layer (0/1/2)")
+@click.option("--ratio", "-r", type=float, default=0.3, help="L1 keep ratio")
+@click.option("--output", "-o", type=click.Path(), help="Output file")
+def compress(file, layer, ratio, output):
+    """Compress a text file."""
+    # Read file -> compress -> print/write result
+
+@cli.command()
+@click.option("--port", "-p", type=int, default=8420)
+def serve(port):
+    """Start the HTTP API server."""
+    import uvicorn
+    uvicorn.run("cctx.server:app", host="0.0.0.0", port=port)
+
+@cli.command()
+@click.argument("file", type=click.Path(exists=True))
+def benchmark(file):
+    """Run compression benchmarks on a file."""
+    # Compress at all layers, print stats table
+```
+
+---
+
+## Tokenizer (`cctx/tokenizer.py`)
+
+```python
+import tiktoken
+
+# Single shared encoder instance
+_enc = tiktoken.get_encoding("cl100k_base")
+
+def count_tokens(text: str) -> int:
+    """Count tokens using cl100k_base (GPT-4/Claude compatible)."""
+    return len(_enc.encode(text))
+
+def truncate_to_tokens(text: str, max_tokens: int) -> str:
+    """Truncate text to fit within max_tokens."""
+    tokens = _enc.encode(text)
+    if len(tokens) <= max_tokens:
+        return text
+    return _enc.decode(tokens[:max_tokens])
+```
+
+---
+
+## How It All Connects
+
+```
+User's App
+    │
+    ▼
+┌─────────────────┐
+│  cctx Client    │  ← Python SDK or HTTP calls
+│  (client.py)    │
+└────────┬────────┘
+         │
+         ▼
+┌─────────────────┐
+│  cctx Server    │  ← FastAPI, 3 endpoints
+│  (server.py)    │
+└────────┬────────┘
+         │
+    ┌────┴────┐
+    ▼         ▼
+┌────────┐ ┌──────────────┐
+│Compress│ │Conversation  │
+│  or    │ │  Manager     │
+└───┬────┘ └──────┬───────┘
+    │              │
+    ▼              ▼
+┌─────────────────────┐
+│  Compressor         │
+│  (compressor.py)    │
+└─────────┬───────────┘
+          │
+          ▼
+┌─────────────────────┐
+│  Scorer (pluggable) │
+│  (scorer.py)        │
+└─────────┬───────────┘
+          │
+          ▼
+┌─────────────────────┐
+│  Tokenizer          │
+│  (tokenizer.py)     │
+└─────────────────────┘
+```
+
+**Data flow for `/v1/compress`:**
+1. Request comes in with text + options
+2. Server passes to `Compressor.compress()`
+3. Compressor splits text into sections → sentences
+4. Compressor calls `Scorer.score_sentences()` on all sentences
+5. Compressor builds Layer 1 (top N% by score, original order)
+6. Compressor builds Layer 2 (top 3 by score, original order)
+7. Compressor extracts entities, generates section summaries
+8. Returns `CompressedResult` → server serializes to JSON
+
+**Data flow for `/v1/conversation`:**
+1. Request comes in with conversation_id + new message
+2. Server passes to `ConversationManager.add_message()`
+3. ConversationManager checks if threshold exceeded
+4. If yes: takes older messages, calls `Compressor.compress()` on them
+5. Stores compressed output as new `compressed_history`
+6. Returns current context (compressed_history + recent raw messages)
+
+**Data flow for `/v1/drill-down`:**
+1. Request comes in with section_id
+2. Server looks up section in the last `CompressedResult` (held in memory)
+3. Returns the section at the requested layer
+
+---
+
+## Dependencies (keep minimal)
+
+```toml
+[project]
+name = "cctx"
+version = "0.1.0"
+requires-python = ">=3.11"
+dependencies = [
+    "tiktoken>=0.7",
+    "fastapi>=0.110",
+    "uvicorn>=0.29",
+    "click>=8.1",
+    "scikit-learn>=1.4",    # For TF-IDF scorer
+    "httpx>=0.27",          # For client.py
+]
+
+[project.scripts]
+cctx = "cctx.cli:cli"
+```
+
+---
+
+## Client SDK (`cctx/client.py`)
+
+```python
+import httpx
+
+class CCTXClient:
+    """Python client for the cctx API."""
+    
+    def __init__(self, base_url: str = "http://localhost:8420"):
+        self.base_url = base_url
+        self.http = httpx.Client(base_url=base_url)
+    
+    def compress(self, text: str, layer: int = 1, **kwargs) -> CompressedResult:
+        """Compress text. Returns CompressedResult."""
+        resp = self.http.post("/v1/compress", json={
+            "text": text, "layer": layer, **kwargs
+        })
+        resp.raise_for_status()
+        return CompressedResult(**resp.json())
+    
+    def add_message(self, conversation_id: str, role: str, 
+                    content: str, **kwargs) -> dict:
+        """Add a message to a conversation."""
+        resp = self.http.post("/v1/conversation", json={
+            "conversation_id": conversation_id,
+            "message": {"role": role, "content": content},
+            **kwargs
+        })
+        resp.raise_for_status()
+        return resp.json()
+    
+    def drill_down(self, section_id: str, layer: int = 0) -> dict:
+        """Get full text of a section."""
+        resp = self.http.post("/v1/drill-down", json={
+            "section_id": section_id, "target_layer": layer
+        })
+        resp.raise_for_status()
+        return resp.json()
+```
+
+---
+
+## Tests
+
+Each test file tests one module. Use pytest.
+
+**`test_compressor.py`:**
+- Test that Layer 1 has fewer tokens than Layer 0
+- Test that Layer 2 has fewer tokens than Layer 1
+- Test that sections are generated and addressable
+- Test with empty string, single sentence, very long text
+- Test that sentence order is preserved (compressed text doesn't reorder)
+
+**`test_conversation.py`:**
+- Test adding messages below threshold (no compression)
+- Test that compression triggers at threshold
+- Test that recent messages are always raw
+- Test that compressed_history grows correctly over multiple compression cycles
+- Test get_context() format
+
+**`test_scorer.py`:**
+- Test TextRankScorer returns scores for each sentence
+- Test scores sum to ~1.0 (normalized)
+- Test TFIDFScorer with a query boosts relevant sentences
+- Test CompositeScorer combines weights correctly
+
+**`test_server.py`:**
+- Test all 3 endpoints return 200
+- Test compress with different layers
+- Test conversation flow (add 10 messages, verify compression happens)
+- Test drill-down returns section text
+
+---
+
+## What NOT To Build (yet)
+
+- No LLM-based scoring (keep it algorithmic for v0.1)
+- No persistence (conversations live in memory — Redis/SQLite is v0.2)
+- No auth on the API (add later)
+- No async compression (single-threaded is fine for prototype)
+- No streaming (compress the whole thing, return it)
+- No Token Company API integration (we have the stub, plug it in later)
+
+---
+
+## Build Order
+
+1. `types.py` — all dataclasses
+2. `tokenizer.py` — token counting
+3. `scorer.py` — TextRankScorer + TFIDFScorer + CompositeScorer
+4. `compressor.py` — uses scorer + tokenizer
+5. `conversation.py` — uses compressor
+6. `server.py` — FastAPI wrapping compressor + conversation
+7. `client.py` — HTTP client
+8. `cli.py` — click CLI
+9. Tests (in parallel with each module)
+10. `examples/` — demo scripts
+
+**Each module depends only on the ones above it.** No circular imports. No god objects.
+
+---
+
+## Success Criteria
+
+The prototype is done when:
+1. `cctx compress long_document.txt` prints a compression summary with all 3 layers
+2. `cctx serve` starts an API on port 8420
+3. The `/v1/conversation` endpoint handles 50 messages and keeps context under a configurable token budget
+4. `pytest` passes all tests
+5. The middleware demo shows: App → cctx → OpenAI, with measurable token savings