content-atom/stage2/dedup.py

#!/usr/bin/env python3
"""
MVP stage 2: dedup atoms across projects.

Для каждого атома ищет топ-K по cosine similarity в Qdrant.
Если similarity > threshold (0.95) → relation_type='duplicate'.
Сохраняет связи в atom_relations (Postgres).

bge-m3 (dim=1024) — единственная embed-модель.
Qdrant — эмбеддинги.
Postgres (atom-db) — atom_relations как источник связей.

Запуск:
    python3 stage2/dedup.py --threshold 0.95 --top-k 5
    python3 stage2/dedup.py --threshold 0.98 --top-k 3
"""

import argparse
import json
import os
import sys
import time
import uuid

import psycopg2
from qdrant_client import QdrantClient
from qdrant_client.models import QueryResponse

POSTGRES = {
    "host": os.environ.get("POSTGRES_HOST", "127.0.0.1"),
    "port": int(os.environ.get("POSTGRES_PORT", "5434")),
    "user": os.environ.get("POSTGRES_USER", "atom"),
    "dbname": os.environ.get("POSTGRES_DB", "content_atom"),
    "password": os.environ.get("POSTGRES_PASSWORD", "atompass123"),
}
COLLECTION = "content_atoms"
DIM = 1024


def get_conn():
    return psycopg2.connect(**POSTGRES)


def query_qdrant(client: QdrantClient, embedding: list[float], top_k: int = 5) -> list[dict]:
    """Поиск топ-K по cosine similarity."""
    try:
        resp: QueryResponse = client.query_points(
            collection_name=COLLECTION,
            query=embedding,
            limit=top_k,
            with_payload=True,
        )
    except Exception as e:
        print(f"  ✗ Qdrant query_points failed: {e}")
        return []

    results = []
    for hit in resp.points:
        payload = hit.payload
        results.append({
            "atom_id": payload.get("atom_id"),
            "similarity": float(hit.score),
            "project": payload.get("project"),
            "type": payload.get("type"),
        })
    return results


def get_all_atoms(conn) -> list[dict]:
    """Все атомы из atoms."""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, project_id, plain_text, word_count
            FROM atoms
            WHERE plain_text IS NOT NULL
            ORDER BY created_at
        """)
        cols = [d[0] for d in cur.description]
        return [dict(zip(cols, row)) for row in cur.fetchall()]


def dedup(client: QdrantClient, atoms: list[dict], threshold: float, top_k: int) -> list[dict]:
    """Дедуп: для каждого атома ищет дубликаты."""
    import urllib.request
    relations = []
    n = 0
    for atom in atoms:
        n += 1
        # embed
        body = atom["plain_text"]
        if len(body) > 4000:
            body = body[:4000]
        try:
            req = urllib.request.Request(
                "http://127.0.0.1:11434/api/embeddings",
                data=json.dumps({"model": "bge-m3", "prompt": body}).encode(),
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=60) as resp:
                emb = json.loads(resp.read().decode())["embedding"]
        except Exception as e:
            print(f"  ✗ embed failed for {atom['id']}: {e}")
            continue
        if len(emb) != DIM:
            print(f"  ✗ wrong dim {len(emb)} for {atom['id']}")
            continue

        hits = query_qdrant(client, emb, top_k)
        for hit in hits:
            if hit["atom_id"] == str(atom["id"]):
                continue
            # только atom_id из atoms (UUID), не из старых эмбеддингов спейка
            try:
                uuid.UUID(hit["atom_id"])
            except ValueError:
                continue
            if hit["similarity"] >= threshold:
                relations.append({
                    "atom_id_a": str(atom["id"]),
                    "atom_id_b": hit["atom_id"],
                    "relation_type": "duplicate",
                    "similarity": round(hit["similarity"], 4),
                    "shared_topics": [],
                    "confidence": 0.9,
                    "source": "semantic",
                })
    return relations


def insert_relations(conn, relations: list[dict]) -> int:
    """Вставить atom_relations. ON CONFLICT DO NOTHING."""
    n = 0
    with conn.cursor() as cur:
        for rel in relations:
            cur.execute("""
                INSERT INTO atom_relations
                (atom_id_a, atom_id_b, relation_type, similarity, shared_topics, confidence, source)
                VALUES (%s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (atom_id_a, atom_id_b, relation_type) DO NOTHING
            """, (
                rel["atom_id_a"], rel["atom_id_b"], rel["relation_type"],
                rel["similarity"], rel["shared_topics"], rel["confidence"], rel["source"],
            ))
            n += cur.rowcount
    conn.commit()
    return n


def main():
    ap = argparse.ArgumentParser(description="MVP stage 2: dedup atoms across projects")
    ap.add_argument("--threshold", type=float, default=0.95, help="cosine similarity threshold")
    ap.add_argument("--top-k", type=int, default=5, help="top-K for Qdrant search")
    args = ap.parse_args()

    print(f"Дедup: threshold={args.threshold}, top_k={args.top_k}")

    # 1. Атомы
    conn = get_conn()
    atoms = get_all_atoms(conn)
    print(f"  • {len(atoms)} атомов")

    # 2. Qdrant
    client = QdrantClient(host="127.0.0.1", port=6333)

    # 3. Дедup
    start = time.time()
    relations = dedup(client, atoms, args.threshold, args.top_k)
    print(f"  ✓ {len(relations)} связей за {time.time()-start:.1f}с")

    # 4. Postgres
    n = insert_relations(conn, relations)
    conn.close()
    print(f"  ✓ Postgres: {n} atom_relations")

    # 5. Итог
    print(f"\nMVP stage 2 OK: {len(atoms)} атомов → {n} дупликатов")
    return 0


if __name__ == "__main__":
    sys.exit(main())