claude-sql 0.4.0__py3-none-any.whl

claude_sql/cluster_worker.py

@@ -0,0 +1,208 @@
+"""Cluster message embeddings via UMAP + HDBSCAN.
+
+Pipeline
+--------
+1. Read the embeddings parquet into a numpy float32 matrix.
+2. UMAP reduce to 50d (for HDBSCAN) and 2d (for viz), both with ``seed=42``.
+3. HDBSCAN on the 50d projection → cluster_id per row, -1 for noise.
+4. Write ``clusters.parquet`` with ``(uuid, cluster_id, x, y, is_noise)``.
+
+Public API
+----------
+run_clustering(settings, *, force=False) -> dict[str, int]
+    Read parquet, compute clusters, write output parquet.  Returns stats dict.
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+from loguru import logger
+
+from claude_sql.config import Settings
+from claude_sql.parquet_shards import iter_part_files
+
+
+def _load_embeddings(path: Path) -> tuple[list[str], np.ndarray]:
+    """Read parquet → (uuid_list, embedding_matrix[float32]).  Matrix shape (N, dim).
+
+    Accepts either a legacy single-file ``embeddings.parquet`` or a sharded
+    directory containing ``part-*.parquet`` files; the union of all parts is
+    materialized in scan order.
+    """
+    parts = iter_part_files(path)
+    df = pl.read_parquet([str(p) for p in parts]) if parts else pl.read_parquet(path)
+    uuids = df["uuid"].to_list()
+    # polars Array[Float32, dim] → 2D numpy (N, dim).  to_numpy() on an Array column
+    # returns 2D when the element type is fixed-size Array; fall back to vstack if not.
+    emb = df["embedding"].to_numpy()
+    if emb.ndim == 1:
+        # Object array of 1D arrays (ragged container); stack into a 2D matrix.
+        emb = np.stack(list(emb))
+    return uuids, np.ascontiguousarray(emb, dtype=np.float32)
+
+
+def run_clustering(settings: Settings, *, force: bool = False) -> dict[str, int]:
+    """Run UMAP + HDBSCAN on the embeddings parquet.
+
+    Parameters
+    ----------
+    settings
+        Configured Settings; reads ``embeddings_parquet_path`` and writes
+        ``clusters_parquet_path``.
+    force
+        If False and clusters_parquet_path exists, return its stats without
+        recomputing.  If True, always rerun.
+
+    Returns
+    -------
+    dict
+        ``{"total": N, "clusters": K, "noise": M}`` where K excludes the
+        noise cluster (label -1).
+    """
+    out_path = settings.clusters_parquet_path
+    in_path = settings.embeddings_parquet_path
+
+    # Sharded directory or legacy single file: insist on at least one part
+    # whose footer is plausibly populated (>16 bytes).
+    in_parts = [p for p in iter_part_files(in_path) if p.stat().st_size > 16]
+    if not in_parts:
+        raise FileNotFoundError(
+            f"Embeddings parquet missing at {in_path}.  Run `claude-sql embed` first."
+        )
+
+    # Mtime-sidecar fast path: if the embeddings haven't moved since the
+    # last successful clustering, skip the ~40 s UMAP+HDBSCAN refit. The
+    # sidecar lives next to the output parquet and stores the latest
+    # part-file's nanosecond mtime; coarse hand-edits of clusters.parquet
+    # alone won't bust the cache (use ``force=True`` for that).
+    sidecar = out_path.with_suffix(out_path.suffix + ".embeddings_mtime")
+    in_mtime_ns = max(p.stat().st_mtime_ns for p in in_parts)
+    if (
+        not force
+        and out_path.exists()
+        and out_path.stat().st_size > 16
+        and sidecar.exists()
+        and sidecar.read_text().strip() == str(in_mtime_ns)
+    ):
+        logger.info("Embeddings unchanged since last cluster run; reusing {}.", out_path)
+        df = pl.read_parquet(out_path)
+        return {
+            "total": len(df),
+            "clusters": int((df["cluster_id"] >= 0).sum()),
+            "noise": int((df["cluster_id"] < 0).sum()),
+        }
+
+    # Legacy short-circuit: a clusters parquet exists but no sidecar (older
+    # install before the mtime-skip landed). Trust the parquet and stamp
+    # a sidecar so the next call hits the fast path. Forces a rebuild only
+    # when ``force=True`` is set explicitly.
+    if not force and out_path.exists() and out_path.stat().st_size > 16 and not sidecar.exists():
+        logger.info(
+            "Clusters parquet at {} predates sidecar; reusing and stamping mtime.",
+            out_path,
+        )
+        sidecar.write_text(str(in_mtime_ns))
+        df = pl.read_parquet(out_path)
+        return {
+            "total": len(df),
+            "clusters": int((df["cluster_id"] >= 0).sum()),
+            "noise": int((df["cluster_id"] < 0).sum()),
+        }
+
+    # Heavy imports inside the function so module import stays cheap.
+    import hdbscan
+    import umap
+
+    t0 = time.monotonic()
+    uuids, X = _load_embeddings(in_path)  # noqa: N806 — X follows sklearn/ML matrix convention
+    logger.info("Loaded {} embeddings, shape={}, dtype={}", len(uuids), X.shape, X.dtype)
+
+    logger.info(
+        "UMAP → {}d (clustering): n_neighbors={}, min_dist={}, metric={}",
+        settings.umap_n_components_50,
+        settings.umap_n_neighbors,
+        settings.umap_min_dist_cluster,
+        settings.umap_metric,
+    )
+    reducer_cluster = umap.UMAP(
+        n_components=settings.umap_n_components_50,
+        n_neighbors=settings.umap_n_neighbors,
+        min_dist=settings.umap_min_dist_cluster,
+        metric=settings.umap_metric,
+        random_state=settings.seed,
+    )
+    X50 = reducer_cluster.fit_transform(X)  # noqa: N806 — 50d projection matrix
+    logger.info("UMAP 50d done in {:.1f}s", time.monotonic() - t0)
+
+    t1 = time.monotonic()
+    logger.info(
+        "UMAP → {}d (viz): n_neighbors={}, min_dist={}",
+        settings.umap_n_components_2,
+        settings.umap_n_neighbors,
+        settings.umap_min_dist_viz,
+    )
+    reducer_viz = umap.UMAP(
+        n_components=settings.umap_n_components_2,
+        n_neighbors=settings.umap_n_neighbors,
+        min_dist=settings.umap_min_dist_viz,
+        metric=settings.umap_metric,
+        random_state=settings.seed,
+    )
+    X2 = reducer_viz.fit_transform(X)  # noqa: N806 — 2d viz projection matrix
+    logger.info("UMAP 2d done in {:.1f}s", time.monotonic() - t1)
+
+    t2 = time.monotonic()
+    logger.info(
+        "HDBSCAN: min_cluster_size={}, min_samples={}",
+        settings.hdbscan_min_cluster_size,
+        settings.hdbscan_min_samples,
+    )
+    clusterer = hdbscan.HDBSCAN(
+        min_cluster_size=settings.hdbscan_min_cluster_size,
+        min_samples=settings.hdbscan_min_samples,
+        metric="euclidean",  # post-UMAP space is euclidean-friendly
+        core_dist_n_jobs=-1,
+        prediction_data=False,
+    )
+    labels = clusterer.fit_predict(X50)
+    k = int(labels.max()) + 1 if labels.max() >= 0 else 0
+    noise = int((labels < 0).sum())
+    logger.info(
+        "HDBSCAN done in {:.1f}s: {} clusters, {} noise ({:.1%})",
+        time.monotonic() - t2,
+        k,
+        noise,
+        noise / len(labels) if len(labels) else 0,
+    )
+
+    df = pl.DataFrame(
+        {
+            "uuid": uuids,
+            "cluster_id": labels.astype(np.int32).tolist(),
+            "x": X2[:, 0].astype(np.float32).tolist(),
+            "y": X2[:, 1].astype(np.float32).tolist(),
+            "is_noise": (labels < 0).tolist(),
+        },
+        schema={
+            "uuid": pl.Utf8,
+            "cluster_id": pl.Int32,
+            "x": pl.Float32,
+            "y": pl.Float32,
+            "is_noise": pl.Boolean,
+        },
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    df.write_parquet(out_path)
+    sidecar.write_text(str(in_mtime_ns))
+    logger.info(
+        "Wrote {} rows to {} (total elapsed: {:.1f}s)",
+        len(df),
+        out_path,
+        time.monotonic() - t0,
+    )
+
+    return {"total": len(uuids), "clusters": k, "noise": noise}

claude_sql/community_worker.py

@@ -0,0 +1,306 @@
+"""Session-level community detection via cosine graph + Louvain.
+
+Pipeline
+--------
+1. Compute session centroid embeddings: group message embeddings by
+   ``messages.session_id`` and average (mean over rows, then L2-normalize).
+2. Pick an adaptive cosine threshold that yields a connected but not
+   dense graph (target average degree in ``[target_avg_degree_low,
+   target_avg_degree_high]``), falling back to ``louvain_edge_threshold``
+   when the adaptive fit fails.
+3. Threshold the pairwise similarity into a sparse graph.
+4. Run ``networkx.community.louvain_communities`` with seed=42.
+5. Collapse singleton communities into a single ``community_id=-1`` "noise"
+   bucket so the output parquet stays scannable.
+6. Write ``session_communities.parquet`` with
+   ``(session_id, community_id, size)``.
+
+Why adaptive thresholding
+-------------------------
+The previous design used a fixed ``louvain_edge_threshold=0.75`` absolute
+cosine cut.  Against ``int8``-quantized embeddings the pairwise cosine
+distribution is both narrower and higher-mean than against unquantized
+``float32``, so 0.75 excluded the majority of "related" session pairs and
+Louvain produced thousands of singleton communities (on the real corpus:
+1,512 communities for 6,329 sessions, with only 21 ≥ 20 sessions).
+
+The adaptive path picks the percentile-based cut that puts average degree
+in the 8–15 range (Louvain's typical sweet spot for this kind of graph),
+so the answer stabilizes whether the underlying embeddings are int8,
+float16, or float32.
+
+Public API
+----------
+run_communities(con, settings, *, force=False) -> dict[str, int]
+    Expects ``con`` to have v1 views registered (``messages`` + access to
+    the embeddings parquet via ``message_embeddings`` table or direct parquet
+    read).  Returns ``{sessions, communities, noise, threshold}``.
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+import polars as pl
+from loguru import logger
+
+from claude_sql.config import Settings
+from claude_sql.parquet_shards import iter_part_files
+
+if TYPE_CHECKING:
+    import duckdb
+
+
+#: Community id used for singleton / unclusterable sessions.  -1 is an obvious
+#: out-of-band sentinel that stays negative even after ``Int32`` serialization.
+NOISE_COMMUNITY_ID: int = -1
+
+
+def _load_session_centroids(
+    con: duckdb.DuckDBPyConnection, embeddings_parquet_path: Path
+) -> tuple[list[str], np.ndarray]:
+    """Return ``(session_ids, centroids)`` where centroids is ``(N_sessions, dim)`` float32.
+
+    Joins the embeddings parquet to the v1 ``messages`` view on uuid, then
+    aggregates inside DuckDB (unnest with position → ``avg`` per (session,
+    dim_index) → ordered ``list``). The L2-normalize step stays in numpy
+    where ``np.linalg.norm`` is faster on a contiguous (N, dim) matrix than
+    a per-row SQL norm. The previous Python double-loop over ~50K rows is
+    replaced with a single vectorized DuckDB query.
+    """
+    logger.info("Loading message embeddings and joining to sessions...")
+    parts = iter_part_files(embeddings_parquet_path)
+    if not parts:
+        raise RuntimeError(
+            f"No embeddings parquet at {embeddings_parquet_path} - run `claude-sql embed`."
+        )
+    # DDL doesn't accept prepared params for read_parquet's file list, so
+    # escape inline. Single-quotes inside paths get doubled per SQL rules.
+    path_literals = ", ".join("'{}'".format(str(p).replace("'", "''")) for p in parts)
+    sql = f"""
+        WITH joined AS (
+            SELECT CAST(m.session_id AS VARCHAR) AS session_id,
+                   e.embedding::FLOAT[] AS emb
+              FROM read_parquet([{path_literals}]) e
+              JOIN messages m
+                ON CAST(m.uuid AS VARCHAR) = e.uuid
+        ),
+        unrolled AS (
+            SELECT session_id,
+                   generate_subscripts(emb, 1) AS pos,
+                   unnest(emb) AS v
+              FROM joined
+        ),
+        agg AS (
+            SELECT session_id, pos, avg(v) AS m
+              FROM unrolled
+             GROUP BY 1, 2
+        )
+        SELECT session_id, list(m ORDER BY pos) AS centroid
+          FROM agg
+         GROUP BY 1
+         ORDER BY 1
+    """
+    df = con.execute(sql).pl()
+    if len(df) == 0:
+        raise RuntimeError(
+            "No rows returned joining embeddings to messages - check that the "
+            "embeddings parquet exists and the messages view is registered."
+        )
+
+    sids = df["session_id"].to_list()
+    centroids = np.stack([np.asarray(c, dtype=np.float32) for c in df["centroid"].to_list()])
+    norms = np.linalg.norm(centroids, axis=1, keepdims=True)
+    centroids = centroids / np.where(norms == 0, 1.0, norms)
+    logger.info("Computed {} session centroids (dim={})", len(sids), centroids.shape[1])
+    return sids, centroids
+
+
+def _pick_adaptive_threshold(
+    sim: np.ndarray,
+    *,
+    floor: float,
+    target_avg_degree_low: float,
+    target_avg_degree_high: float,
+) -> float:
+    """Pick a cosine threshold that puts the average graph degree in the target band.
+
+    Binary-searches the upper-triangular similarity quantiles.  Falls back to
+    ``floor`` when the graph is too sparse to hit the band even at q=0.
+    """
+    n = sim.shape[0]
+    if n < 2:
+        return floor
+    # Upper triangle only -- pairwise matrix is symmetric; diag already zeroed.
+    iu = np.triu_indices(n, k=1)
+    upper = sim[iu]
+    if upper.size == 0:
+        return floor
+
+    target_edges_low = n * target_avg_degree_low / 2.0
+    target_edges_high = n * target_avg_degree_high / 2.0
+
+    # Quantile search: try a handful of candidate thresholds and pick the
+    # lowest one that stays under target_edges_high, then clamp up to floor.
+    # Using np.quantile ensures the candidate values are monotonically
+    # decreasing -> edge count monotonically increasing -- no binary search
+    # machinery needed.
+    quantiles = np.linspace(0.999, 0.80, 20)
+    candidates = np.quantile(upper, quantiles)
+    for thr in candidates:
+        n_edges = int((upper >= thr).sum())
+        if target_edges_low <= n_edges <= target_edges_high:
+            return max(float(thr), floor)
+    # No quantile hit the target band: pick the tightest threshold that stays
+    # at or under the high target to avoid a hairball.
+    for thr in candidates:
+        n_edges = int((upper >= thr).sum())
+        if n_edges <= target_edges_high:
+            return max(float(thr), floor)
+    return floor
+
+
+def run_communities(
+    con: duckdb.DuckDBPyConnection, settings: Settings, *, force: bool = False
+) -> dict[str, int | float]:
+    """Run Louvain on session centroids, write parquet.
+
+    Parameters
+    ----------
+    con
+        An open DuckDB connection with the v1 ``messages`` view registered.
+    settings
+        Configured Settings.  Honors ``louvain_edge_threshold`` (the floor
+        cosine value), ``louvain_resolution``, ``louvain_target_avg_degree_low``,
+        ``louvain_target_avg_degree_high``, and ``louvain_min_community_size``.
+    force
+        If False and output parquet exists, skip recomputation.
+
+    Returns
+    -------
+    dict
+        ``{sessions, communities, noise, threshold}`` — ``communities`` counts
+        only real (>= min-size) communities; singletons are aggregated into
+        ``noise``.  ``threshold`` is the adaptive cut actually used.
+    """
+    out_path = settings.communities_parquet_path
+
+    if out_path.exists() and out_path.stat().st_size > 16 and not force:
+        logger.info("Communities parquet exists at {}; pass force=True to rebuild", out_path)
+        df = pl.read_parquet(out_path)
+        real = df.filter(pl.col("community_id") != NOISE_COMMUNITY_ID)
+        noise_n = df.height - real.height
+        n_comm = int(real["community_id"].n_unique()) if real.height else 0
+        return {
+            "sessions": int(df.height),
+            "communities": n_comm,
+            "noise": noise_n,
+            "threshold": float("nan"),
+        }
+
+    import networkx as nx
+    from scipy.sparse import csr_matrix
+
+    sids, centroids = _load_session_centroids(con, settings.embeddings_parquet_path)
+
+    t0 = time.monotonic()
+    logger.info("Computing pairwise cosine similarity over {} sessions...", len(sids))
+    sim = centroids @ centroids.T
+    np.fill_diagonal(sim, 0.0)
+
+    floor = settings.louvain_edge_threshold
+    threshold = _pick_adaptive_threshold(
+        sim,
+        floor=floor,
+        target_avg_degree_low=settings.louvain_target_avg_degree_low,
+        target_avg_degree_high=settings.louvain_target_avg_degree_high,
+    )
+
+    # Apply the chosen threshold.  Operate on a copy so we keep the raw sim
+    # matrix intact for logging if we need to probe.
+    masked = np.where(sim >= threshold, sim, 0.0).astype(np.float32, copy=False)
+    sp = csr_matrix(masked)
+    logger.info(
+        "Graph: {} nodes, {} edges (threshold={:.3f}, avg_degree={:.1f}) in {:.1f}s",
+        sp.shape[0],
+        sp.nnz // 2,
+        threshold,
+        (sp.nnz / sp.shape[0]) if sp.shape[0] else 0.0,
+        time.monotonic() - t0,
+    )
+
+    graph = nx.from_scipy_sparse_array(sp)
+    t1 = time.monotonic()
+    communities = nx.community.louvain_communities(
+        graph,
+        weight="weight",
+        resolution=settings.louvain_resolution,
+        seed=settings.seed,
+    )
+
+    min_size = settings.louvain_min_community_size
+    rows: list[tuple[str, int, int]] = []
+    real_communities = 0
+    noise_count = 0
+    # networkx.community.louvain_communities returns list[set[int]] at
+    # runtime but the default stubs only expose it as Sized.  Cast through
+    # typing.cast so the type checker agrees with reality.
+    from typing import cast
+
+    typed_communities = cast("list[set[int]]", communities)
+    sorted_communities = cast(
+        "list[set[int]]",
+        sorted(typed_communities, key=len, reverse=True),
+    )
+    for cid, comm in enumerate(sorted_communities):
+        size = len(comm)
+        if size < min_size:
+            for node in comm:
+                rows.append((sids[node], NOISE_COMMUNITY_ID, 0))
+                noise_count += 1
+            continue
+        real_communities += 1
+        # PERF401: comprehension would obscure the nested control flow above.
+        for node in comm:
+            rows.append((sids[node], cid, size))  # noqa: PERF401
+
+    logger.info(
+        "Louvain: {} raw communities ({} kept >= {} sessions, {} singletons -> noise) in {:.1f}s",
+        len(communities),
+        real_communities,
+        min_size,
+        noise_count,
+        time.monotonic() - t1,
+    )
+
+    df = pl.DataFrame(
+        rows,
+        schema={
+            "session_id": pl.Utf8,
+            "community_id": pl.Int32,
+            "size": pl.Int32,
+        },
+        orient="row",
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    df.write_parquet(out_path)
+    logger.info("Wrote {} rows to {}", len(df), out_path)
+
+    size_dist = (
+        df.filter(pl.col("community_id") != NOISE_COMMUNITY_ID)
+        .group_by("community_id")
+        .agg(pl.len().alias("n"))
+        .sort("n", descending=True)
+        .head(5)
+    )
+    logger.info("Top 5 community sizes: {}", size_dist.to_dicts())
+
+    return {
+        "sessions": int(df.height),
+        "communities": real_communities,
+        "noise": noise_count,
+        "threshold": float(threshold),
+    }