spinchain 0.1.0__py3-none-any.whl

spinchain/__init__.py

@@ -0,0 +1,3 @@
+"""SpinChain: Inference-time reasoning optimization using QUBO/Ising formulations."""
+
+__version__ = "0.1.0"

spinchain/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for python -m spinchain."""
+
+from spinchain.server import main
+
+main()

spinchain/analyze.py

@@ -0,0 +1,326 @@
+"""Trace analysis pipeline for SpinChain MCP usage.
+
+Reads JSONL traces and produces:
+- Usage summary: call count, time range, success/fallback rates
+- Latency breakdown: per-stage timing percentiles
+- Energy distribution: min/mean/spread across calls
+- Anomaly flags: unusually slow stages, empty selections, errors
+
+Run as: uv run python -m spinchain.analyze [--last N] [--json]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from spinchain.tracing import TraceLogger, DEFAULT_TRACE_DIR
+
+
+@dataclass
+class UsageSummary:
+    """High-level usage statistics."""
+
+    total_calls: int = 0
+    first_call: str | None = None
+    last_call: str | None = None
+    success_count: int = 0
+    fallback_count: int = 0
+    error_count: int = 0
+
+
+@dataclass
+class LatencyStats:
+    """Latency percentiles for a pipeline stage (in ms)."""
+
+    stage: str
+    count: int = 0
+    min_ms: float = 0.0
+    median_ms: float = 0.0
+    p95_ms: float = 0.0
+    max_ms: float = 0.0
+    mean_ms: float = 0.0
+
+
+@dataclass
+class EnergyStats:
+    """Energy distribution across SA solutions."""
+
+    count: int = 0
+    min_energy: float | None = None
+    mean_energy: float | None = None
+    max_energy: float | None = None
+
+
+@dataclass
+class Anomaly:
+    """A flagged trace with reason."""
+
+    trace_id: str
+    timestamp: str
+    reason: str
+    details: dict[str, Any] = field(default_factory=dict)
+
+
+def _percentile(values: list[float], p: float) -> float:
+    """Compute p-th percentile (0-100) from sorted values."""
+    if not values:
+        return 0.0
+    k = (len(values) - 1) * (p / 100.0)
+    f = int(k)
+    c = f + 1
+    if c >= len(values):
+        return values[f]
+    return values[f] + (k - f) * (values[c] - values[f])
+
+
+class TraceAnalyzer:
+    """Analyzes SpinChain trace records from JSONL files."""
+
+    def __init__(self, traces: list[dict[str, Any]]):
+        self.traces = traces
+
+    @classmethod
+    def from_file(cls, trace_dir: str | Path | None = None, last_n: int | None = None) -> "TraceAnalyzer":
+        """Load traces from the JSONL file."""
+        tracer = TraceLogger(trace_dir=trace_dir)
+        traces = tracer.read_traces(last_n=last_n)
+        return cls(traces)
+
+    def usage_summary(self) -> UsageSummary:
+        """Compute usage summary stats."""
+        summary = UsageSummary(total_calls=len(self.traces))
+        if not self.traces:
+            return summary
+
+        summary.first_call = self.traces[0].get("timestamp")
+        summary.last_call = self.traces[-1].get("timestamp")
+
+        for t in self.traces:
+            if t.get("error"):
+                summary.error_count += 1
+            elif t.get("output_summary", {}).get("fallback"):
+                summary.fallback_count += 1
+            else:
+                summary.success_count += 1
+
+        return summary
+
+    def latency_breakdown(self) -> list[LatencyStats]:
+        """Compute per-stage latency percentiles."""
+        stage_times: dict[str, list[float]] = {}
+
+        for t in self.traces:
+            for stage in t.get("stages", []):
+                name = stage.get("name", "unknown")
+                dur = stage.get("duration_ms", 0.0)
+                stage_times.setdefault(name, []).append(dur)
+
+        # Also compute total
+        total_times = [
+            t.get("total_duration_ms", 0.0)
+            for t in self.traces
+            if t.get("total_duration_ms")
+        ]
+        if total_times:
+            stage_times["total"] = total_times
+
+        results = []
+        for stage_name in ["fragment_extraction", "qubo_formulation", "simulated_annealing", "stability_ranking", "total"]:
+            times = sorted(stage_times.get(stage_name, []))
+            if not times:
+                continue
+            results.append(LatencyStats(
+                stage=stage_name,
+                count=len(times),
+                min_ms=round(times[0], 2),
+                median_ms=round(_percentile(times, 50), 2),
+                p95_ms=round(_percentile(times, 95), 2),
+                max_ms=round(times[-1], 2),
+                mean_ms=round(sum(times) / len(times), 2),
+            ))
+
+        return results
+
+    def energy_stats(self) -> EnergyStats:
+        """Compute energy distribution across successful calls."""
+        min_energies = []
+        for t in self.traces:
+            out = t.get("output_summary", {})
+            me = out.get("min_energy")
+            if me is not None:
+                min_energies.append(me)
+
+        if not min_energies:
+            return EnergyStats()
+
+        return EnergyStats(
+            count=len(min_energies),
+            min_energy=round(min(min_energies), 4),
+            mean_energy=round(sum(min_energies) / len(min_energies), 4),
+            max_energy=round(max(min_energies), 4),
+        )
+
+    def detect_anomalies(
+        self,
+        slow_threshold_ms: float = 10000.0,
+        slow_stage_threshold_ms: float = 5000.0,
+    ) -> list[Anomaly]:
+        """Flag traces with potential issues."""
+        anomalies = []
+
+        for t in self.traces:
+            tid = t.get("trace_id", "?")
+            ts = t.get("timestamp", "?")
+
+            # Error
+            if t.get("error"):
+                anomalies.append(Anomaly(
+                    trace_id=tid, timestamp=ts,
+                    reason="error",
+                    details={"error": t["error"]},
+                ))
+
+            # Slow total
+            total_ms = t.get("total_duration_ms", 0.0)
+            if total_ms > slow_threshold_ms:
+                anomalies.append(Anomaly(
+                    trace_id=tid, timestamp=ts,
+                    reason="slow_total",
+                    details={"total_duration_ms": total_ms},
+                ))
+
+            # Slow individual stage
+            for stage in t.get("stages", []):
+                if stage.get("duration_ms", 0) > slow_stage_threshold_ms:
+                    anomalies.append(Anomaly(
+                        trace_id=tid, timestamp=ts,
+                        reason="slow_stage",
+                        details={
+                            "stage": stage["name"],
+                            "duration_ms": stage["duration_ms"],
+                        },
+                    ))
+
+            # Empty selection (non-fallback call selected 0 fragments)
+            out = t.get("output_summary", {})
+            if not out.get("fallback") and out.get("num_selected", 1) == 0:
+                anomalies.append(Anomaly(
+                    trace_id=tid, timestamp=ts,
+                    reason="empty_selection",
+                    details={"num_fragments": out.get("num_fragments")},
+                ))
+
+        return anomalies
+
+    def full_report(self) -> dict[str, Any]:
+        """Generate complete analysis report as a dict."""
+        usage = self.usage_summary()
+        latency = self.latency_breakdown()
+        energy = self.energy_stats()
+        anomalies = self.detect_anomalies()
+
+        return {
+            "usage": {
+                "total_calls": usage.total_calls,
+                "first_call": usage.first_call,
+                "last_call": usage.last_call,
+                "success": usage.success_count,
+                "fallback": usage.fallback_count,
+                "errors": usage.error_count,
+            },
+            "latency": [
+                {
+                    "stage": s.stage,
+                    "count": s.count,
+                    "min_ms": s.min_ms,
+                    "median_ms": s.median_ms,
+                    "p95_ms": s.p95_ms,
+                    "max_ms": s.max_ms,
+                    "mean_ms": s.mean_ms,
+                }
+                for s in latency
+            ],
+            "energy": {
+                "count": energy.count,
+                "min": energy.min_energy,
+                "mean": energy.mean_energy,
+                "max": energy.max_energy,
+            },
+            "anomalies": [
+                {
+                    "trace_id": a.trace_id,
+                    "timestamp": a.timestamp,
+                    "reason": a.reason,
+                    **a.details,
+                }
+                for a in anomalies
+            ],
+        }
+
+    def print_report(self, file=sys.stdout) -> None:
+        """Print a human-readable report."""
+        usage = self.usage_summary()
+
+        print("=" * 60, file=file)
+        print("  SpinChain Trace Analysis", file=file)
+        print("=" * 60, file=file)
+
+        print(f"\n  Calls: {usage.total_calls}  "
+              f"(success: {usage.success_count}, "
+              f"fallback: {usage.fallback_count}, "
+              f"errors: {usage.error_count})", file=file)
+        if usage.first_call:
+            print(f"  Range: {usage.first_call} -> {usage.last_call}", file=file)
+
+        latency = self.latency_breakdown()
+        if latency:
+            print(f"\n  {'Stage':<25} {'Count':>5} {'Min':>8} {'Median':>8} {'P95':>8} {'Max':>8} {'Mean':>8}", file=file)
+            print(f"  {'-'*23}  {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8}", file=file)
+            for s in latency:
+                print(f"  {s.stage:<25} {s.count:>5} {s.min_ms:>7.1f}ms {s.median_ms:>7.1f}ms "
+                      f"{s.p95_ms:>7.1f}ms {s.max_ms:>7.1f}ms {s.mean_ms:>7.1f}ms", file=file)
+
+        energy = self.energy_stats()
+        if energy.count:
+            print(f"\n  Energy: min={energy.min_energy}, mean={energy.mean_energy}, "
+                  f"max={energy.max_energy} (n={energy.count})", file=file)
+
+        anomalies = self.detect_anomalies()
+        if anomalies:
+            print(f"\n  Anomalies ({len(anomalies)}):", file=file)
+            for a in anomalies:
+                print(f"    [{a.reason}] {a.trace_id} @ {a.timestamp} — {a.details}", file=file)
+        else:
+            print("\n  No anomalies detected.", file=file)
+
+        print("", file=file)
+
+
+def main():
+    """CLI entry point for trace analysis."""
+    import os
+
+    parser = argparse.ArgumentParser(description="Analyze SpinChain MCP traces")
+    parser.add_argument("--trace-dir", default=os.environ.get("SPINCHAIN_TRACE_DIR"),
+                        help="Path to trace directory (default: SPINCHAIN_TRACE_DIR or ~/.spinchain/traces)")
+    parser.add_argument("--last", type=int, default=None,
+                        help="Only analyze last N traces")
+    parser.add_argument("--json", action="store_true",
+                        help="Output as JSON instead of human-readable")
+    args = parser.parse_args()
+
+    analyzer = TraceAnalyzer.from_file(trace_dir=args.trace_dir, last_n=args.last)
+
+    if args.json:
+        print(json.dumps(analyzer.full_report(), indent=2))
+    else:
+        analyzer.print_report()
+
+
+if __name__ == "__main__":
+    main()

spinchain/formulation/__init__.py

@@ -0,0 +1,5 @@
+from spinchain.formulation.fragment_extractor import FragmentExtractor
+from spinchain.formulation.coefficient_builder import CoefficientBuilder
+from spinchain.formulation.qubo_builder import QUBOBuilder
+
+__all__ = ["FragmentExtractor", "CoefficientBuilder", "QUBOBuilder"]

spinchain/formulation/coefficient_builder.py

@@ -0,0 +1,108 @@
+"""Build HUBO/QUBO coefficients from fragment statistics.
+
+Implements the coefficient design from QCR-LLM (arXiv 2510.24509):
+- Linear: w_i = -mu * p_i + alpha * risk_i
+- Quadratic: w_ij = -beta * (corr_ij - lambda^2 * sim(i,j))
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.stats import zscore
+
+
+class CoefficientBuilder:
+    """Computes linear and quadratic coefficients for the QUBO formulation.
+
+    Args:
+        mu: Weight for fragment popularity (higher = prefer common fragments).
+        alpha: Weight for fragment stability/risk penalty.
+        beta: Weight for pairwise coherence terms.
+        lambda_sim: Semantic similarity penalty factor.
+        regularization: Small constant for z-score stability.
+    """
+
+    def __init__(
+        self,
+        mu: float = 1.0,
+        alpha: float = 0.5,
+        beta: float = 1.0,
+        lambda_sim: float = 0.3,
+        regularization: float = 1e-6,
+    ):
+        self.mu = mu
+        self.alpha = alpha
+        self.beta = beta
+        self.lambda_sim = lambda_sim
+        self.regularization = regularization
+
+    def compute_linear_weights(
+        self,
+        fragment_sources: list[set[int]],
+        num_completions: int,
+    ) -> np.ndarray:
+        """Compute linear (1-body) weights for each fragment.
+
+        w_i = -mu * p_i + alpha * risk_i
+        where p_i = popularity, risk_i = p_i * (1 - p_i)
+        """
+        r = len(fragment_sources)
+        weights = np.zeros(r)
+
+        for i in range(r):
+            p_i = len(fragment_sources[i]) / num_completions
+            risk_i = p_i * (1.0 - p_i)
+            weights[i] = -self.mu * p_i + self.alpha * risk_i
+
+        return weights
+
+    def compute_quadratic_weights(
+        self,
+        fragment_sources: list[set[int]],
+        num_completions: int,
+        embeddings: np.ndarray,
+    ) -> np.ndarray:
+        """Compute quadratic (2-body) weights for fragment pairs.
+
+        w_ij = -beta * (z_score(corr_ij) - lambda^2 * sim(i,j))
+        where corr_ij = co-occurrence correlation, sim = cosine similarity
+        """
+        r = len(fragment_sources)
+        n = num_completions
+
+        # Popularity
+        p = np.array([len(s) / n for s in fragment_sources])
+
+        # Co-occurrence matrix
+        co_occurrence = np.zeros((r, r))
+        for i in range(r):
+            for j in range(i + 1, r):
+                n_ij = len(fragment_sources[i] & fragment_sources[j])
+                co_occurrence[i, j] = n_ij / n - p[i] * p[j]
+                co_occurrence[j, i] = co_occurrence[i, j]
+
+        # Standardize correlations (z-scores)
+        upper_tri = co_occurrence[np.triu_indices(r, k=1)]
+        if len(upper_tri) > 1 and np.std(upper_tri) > self.regularization:
+            mean_c = np.mean(upper_tri)
+            std_c = np.std(upper_tri)
+            z_corr = (co_occurrence - mean_c) / (std_c + self.regularization)
+        else:
+            z_corr = co_occurrence
+
+        # Cosine similarity matrix
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        norms = np.maximum(norms, 1e-10)
+        normed = embeddings / norms
+        sim_matrix = normed @ normed.T
+
+        # Quadratic weights
+        weights = np.zeros((r, r))
+        for i in range(r):
+            for j in range(i + 1, r):
+                weights[i, j] = -self.beta * (
+                    z_corr[i, j] - self.lambda_sim**2 * sim_matrix[i, j]
+                )
+                weights[j, i] = weights[i, j]
+
+        return weights

spinchain/formulation/fragment_extractor.py

@@ -0,0 +1,115 @@
+"""Extract and deduplicate reasoning fragments from multiple LLM completions."""
+
+from __future__ import annotations
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+
+class FragmentExtractor:
+    """Extracts distinct reasoning fragments from N CoT completions.
+
+    Following QCR-LLM: splits each completion into sentence-level fragments,
+    computes embeddings, merges near-duplicates via cosine similarity threshold,
+    and returns a normalized fragment pool.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "all-MiniLM-L6-v2",
+        similarity_threshold: float = 0.85,
+    ):
+        self.model = SentenceTransformer(model_name)
+        self.similarity_threshold = similarity_threshold
+
+    def extract_fragments(self, completions: list[str]) -> list[str]:
+        """Extract and deduplicate fragments from multiple completions.
+
+        Args:
+            completions: List of N CoT completion strings from the LLM.
+
+        Returns:
+            List of R distinct reasoning fragments.
+        """
+        raw_fragments = []
+        fragment_sources: list[set[int]] = []
+
+        for comp_idx, completion in enumerate(completions):
+            sentences = self._split_into_sentences(completion)
+            for sentence in sentences:
+                cleaned = sentence.strip()
+                if len(cleaned) < 10:
+                    continue
+                raw_fragments.append(cleaned)
+                fragment_sources.append({comp_idx})
+
+        if not raw_fragments:
+            return []
+
+        embeddings = self.model.encode(raw_fragments, convert_to_numpy=True)
+        merged_fragments, merged_sources, merged_embeddings = self._merge_similar(
+            raw_fragments, fragment_sources, embeddings
+        )
+
+        self._last_fragments = merged_fragments
+        self._last_sources = merged_sources
+        self._last_embeddings = merged_embeddings
+        self._num_completions = len(completions)
+
+        return merged_fragments
+
+    @property
+    def fragment_embeddings(self) -> np.ndarray:
+        return self._last_embeddings
+
+    @property
+    def fragment_sources(self) -> list[set[int]]:
+        return self._last_sources
+
+    @property
+    def num_completions(self) -> int:
+        return self._num_completions
+
+    def _split_into_sentences(self, text: str) -> list[str]:
+        """Split text into reasoning sentences. Simple split on period/newline."""
+        import re
+
+        sentences = re.split(r"(?<=[.!?])\s+|\n+", text)
+        return [s for s in sentences if s.strip()]
+
+    def _merge_similar(
+        self,
+        fragments: list[str],
+        sources: list[set[int]],
+        embeddings: np.ndarray,
+    ) -> tuple[list[str], list[set[int]], np.ndarray]:
+        """Merge fragments whose cosine similarity exceeds threshold."""
+        n = len(fragments)
+        if n == 0:
+            return [], [], np.array([])
+
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        norms = np.maximum(norms, 1e-10)
+        normed = embeddings / norms
+        sim_matrix = normed @ normed.T
+
+        merged = [False] * n
+        result_fragments = []
+        result_sources = []
+        result_embeddings = []
+
+        for i in range(n):
+            if merged[i]:
+                continue
+            group_sources = set(sources[i])
+            for j in range(i + 1, n):
+                if merged[j]:
+                    continue
+                if sim_matrix[i, j] >= self.similarity_threshold:
+                    group_sources |= sources[j]
+                    merged[j] = True
+            result_fragments.append(fragments[i])
+            result_sources.append(group_sources)
+            result_embeddings.append(embeddings[i])
+
+        return result_fragments, result_sources, np.array(result_embeddings)

spinchain/formulation/qubo_builder.py

@@ -0,0 +1,66 @@
+"""Build QUBO from linear and quadratic weights.
+
+Converts the HUBO formulation into a QUBO compatible with any
+Ising/QUBO solver (SA, D-Wave, COBI, Quanfluence, Qiskit QAOA).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import dimod
+
+
+class QUBOBuilder:
+    """Builds a dimod BinaryQuadraticModel from fragment coefficients.
+
+    The QUBO encodes: minimize H(x) = sum_i w_i * x_i + sum_{i<j} w_ij * x_i * x_j
+    with an optional cardinality constraint to select ~K fragments.
+    """
+
+    def __init__(self, penalty_strength: float = 5.0):
+        self.penalty_strength = penalty_strength
+
+    def build(
+        self,
+        linear_weights: np.ndarray,
+        quadratic_weights: np.ndarray,
+        target_fragments: int | None = None,
+    ) -> dimod.BinaryQuadraticModel:
+        """Build BQM from weights.
+
+        Args:
+            linear_weights: 1D array of shape (R,) — linear coefficients.
+            quadratic_weights: 2D array of shape (R, R) — quadratic coefficients.
+            target_fragments: If set, add penalty to select approximately K fragments.
+
+        Returns:
+            dimod.BinaryQuadraticModel ready for any solver.
+        """
+        r = len(linear_weights)
+        linear = {}
+        quadratic = {}
+
+        for i in range(r):
+            linear[i] = float(linear_weights[i])
+
+        for i in range(r):
+            for j in range(i + 1, r):
+                if abs(quadratic_weights[i, j]) > 1e-10:
+                    quadratic[(i, j)] = float(quadratic_weights[i, j])
+
+        # Add cardinality constraint: penalty * (sum(x_i) - K)^2
+        if target_fragments is not None:
+            k = target_fragments
+            for i in range(r):
+                linear[i] = linear.get(i, 0.0) + self.penalty_strength * (1 - 2 * k)
+            for i in range(r):
+                for j in range(i + 1, r):
+                    key = (i, j)
+                    quadratic[key] = quadratic.get(key, 0.0) + 2 * self.penalty_strength
+
+        bqm = dimod.BinaryQuadraticModel(linear, quadratic, 0.0, dimod.BINARY)
+        return bqm
+
+    def bqm_to_qubo(self, bqm: dimod.BinaryQuadraticModel) -> tuple[dict, float]:
+        """Convert BQM to raw QUBO dict format (for custom solvers)."""
+        return bqm.to_qubo()

spinchain/pipeline/__init__.py

@@ -0,0 +1,2 @@
+# Pipeline module removed — SpinChain is now an MCP server.
+# See spinchain.server for the optimize_reasoning tool.