poimen-memory/scripts/compact_knowledge.py

#!/usr/bin/env python3
"""
mem compact — Deduplicate and merge similar knowledge chunks.

Flow:
1. Read all chunks from JSONL event log
2. Embed each chunk via embedding API
3. Compute cosine similarity matrix
4. Group chunks with similarity > threshold
5. Send each group to reasoning model to merge into one
6. Write compacted JSONL + update markdown files

Usage:
    python3 scripts/compact_knowledge.py [--threshold 0.80] [--dry-run] [--project knowledge]
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path

import requests
import numpy as np

API_BASE = os.environ.get("MEM_API_BASE", "https://api.riotpiao.com/v1")
EMBED_MODEL = "nomic-ai/nomic-embed-text-v2-moe"
REASON_MODEL = "reasoning"


def embed_batch(texts: list[str], batch_size: int = 32) -> list[list[float]]:
    """Embed texts in batches."""
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        resp = requests.post(
            f"{API_BASE}/embeddings",
            json={"model": EMBED_MODEL, "input": batch},
        )
        resp.raise_for_status()
        data = resp.json()["data"]
        # Sort by index to maintain order
        data.sort(key=lambda x: x["index"])
        all_embeddings.extend([d["embedding"] for d in data])
        if i + batch_size < len(texts):
            sys.stderr.write(f"  Embedded {i + batch_size}/{len(texts)}...\n")
    return all_embeddings


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between two vectors."""
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-10))


def find_similar_groups(
    chunks: list[dict], embeddings: list[list[float]], threshold: float
) -> list[list[int]]:
    """Group chunk indices by cosine similarity > threshold."""
    n = len(chunks)
    vecs = np.array(embeddings)
    # Normalize
    norms = np.linalg.norm(vecs, axis=1, keepdims=True)
    vecs_norm = vecs / (norms + 1e-10)
    # Similarity matrix
    sim_matrix = vecs_norm @ vecs_norm.T

    visited = set()
    groups = []

    for i in range(n):
        if i in visited:
            continue
        group = [i]
        visited.add(i)
        for j in range(i + 1, n):
            if j in visited:
                continue
            if sim_matrix[i, j] > threshold:
                group.append(j)
                visited.add(j)
        if len(group) > 1:
            groups.append(group)

    return groups


def merge_with_llm(chunks_text: list[str], source_files: list[str]) -> str:
    """Send similar chunks to reasoning model for merging."""
    numbered = "\n\n".join(
        f"[Chunk {i+1} from {src}]:\n{text}"
        for i, (text, src) in enumerate(zip(chunks_text, source_files))
    )

    resp = requests.post(
        f"{API_BASE}/chat/completions",
        json={
            "model": REASON_MODEL,
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "You are a knowledge compactor. Merge the following similar chunks "
                        "into ONE concise chunk. Keep ALL unique facts. Remove redundancy. "
                        "Keep markdown formatting. Output ONLY the merged text, no explanation. "
                        "No meta-commentary. No 'Here is the merged version'. Just the content."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Merge these {len(chunks_text)} similar chunks:\n\n{numbered}",
                },
            ],
            "temperature": 0.1,
            "max_tokens": 1500,
        },
    )
    resp.raise_for_status()
    import re
    result = resp.json()["choices"][0]["message"]["content"]
    # Strip <think> tags from reasoning model
    result = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL).strip()
    return result


def main():
    parser = argparse.ArgumentParser(description="Compact knowledge chunks")
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.82,
        help="Cosine similarity threshold for grouping (default: 0.82)",
    )
    parser.add_argument("--dry-run", action="store_true", help="Show groups without merging")
    parser.add_argument("--project", default="knowledge", help="Project name")
    parser.add_argument(
        "--log-file",
        default="log/knowledge/learn/latest.jsonl",
        help="Path to JSONL log",
    )
    args = parser.parse_args()

    # 1. Read chunks
    log_path = Path(args.log_file)
    if not log_path.exists():
        print(f"No log file at {log_path}")
        sys.exit(1)

    records = []
    with open(log_path) as f:
        for line in f:
            records.append(json.loads(line))

    texts = [r["data"]["text"] for r in records]
    sources = [r["data"].get("source", "unknown") for r in records]
    print(f"Loaded {len(records)} chunks from {log_path}")

    # 2. Embed all chunks
    print("Embedding chunks...")
    embeddings = embed_batch(texts)
    print(f"Embedded {len(embeddings)} chunks ({len(embeddings[0])} dims)")

    # 3. Find similar groups
    print(f"Finding groups with similarity > {args.threshold}...")
    groups = find_similar_groups(records, embeddings, args.threshold)

    if not groups:
        print("✓ No similar chunks found. Knowledge is already compact.")
        sys.exit(0)

    # 4. Report
    total_mergeable = sum(len(g) for g in groups)
    savings = total_mergeable - len(groups)
    print(f"\nFound {len(groups)} groups ({total_mergeable} chunks → {len(groups)} merged)")
    print(f"Estimated savings: {savings} chunks removed\n")

    for gi, group in enumerate(groups):
        group_texts = [texts[i] for i in group]
        group_sources = [sources[i] for i in group]
        max_sim = 0
        for a in range(len(group)):
            for b in range(a + 1, len(group)):
                s = cosine_similarity(
                    np.array(embeddings[group[a]]), np.array(embeddings[group[b]])
                )
                max_sim = max(max_sim, s)

        print(f"Group {gi+1} (sim={max_sim:.3f}, {len(group)} chunks):")
        for idx in group:
            preview = texts[idx][:80].replace("\n", " ")
            src = Path(sources[idx]).stem
            print(f"  [{src}] {preview}...")

        if not args.dry_run:
            print(f"  → Merging with {REASON_MODEL}...")
            merged = merge_with_llm(group_texts, group_sources)
            print(f"  → Merged: {len(merged)} chars (was {sum(len(t) for t in group_texts)} chars)")
            # Replace first chunk with merged, mark others for removal
            records[group[0]]["data"]["text"] = merged
            records[group[0]]["data"]["merged_from"] = len(group)
            for idx in group[1:]:
                records[idx] = None  # Mark for removal
        print()

    if args.dry_run:
        print("(dry run — no changes written)")
        sys.exit(0)

    # 5. Write compacted log
    compacted = [r for r in records if r is not None]
    backup = log_path.with_suffix(".jsonl.bak")
    log_path.rename(backup)

    with open(log_path, "w") as f:
        for r in compacted:
            f.write(json.dumps(r) + "\n")

    print(f"{'─' * 50}")
    print(f"Before: {len(records)} chunks")
    print(f"After:  {len(compacted)} chunks (-{len(records) - len(compacted)})")
    print(f"Backup: {backup}")
    print(f"Written: {log_path}")


if __name__ == "__main__":
    main()