feat: add reason_first_program/steering.py

reason_first_program/steering.py

@@ -0,0 +1,438 @@
+"""
+Stage 4: Query Language and Steering Interface
+
+Provides a formal query language for navigating the program space by composing
+concepts. Users express preferences as concept coordinates, and the system
+steers LLM generation accordingly.
+
+The query language supports:
+  - Single concept selection: steer("recursive", strength=0.8)
+  - Concept composition: steer(recursive=0.8, space_efficient=0.6)
+  - Concept negation: steer(mutation=-0.5)  (avoid mutation)
+  - Region queries: select(region_where(recursive > 0.5, fast_execution > 0.7))
+  - Lattice navigation: refine(current, add_concept="memoization")
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from reason_first_program.program_space import Program, ProgramSpace
+from reason_first_program.concepts import Concept, ConceptSet
+from reason_first_program.embeddings import (
+    ConceptEmbeddingSpace,
+    GCAVEmbedding,
+    MSRSSteering,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ConceptQuery:
+    """
+    A query in the concept space.
+    
+    A query is a weighted combination of concepts that defines a target
+    region in the program space. The system either:
+      1. Selects existing programs nearest to this region, or
+      2. Steers generation toward this region.
+    
+    Formally: q = Σ_i w_i · v_i where v_i is concept i's activation vector
+    """
+
+    weights: dict[str, float] = field(default_factory=dict)
+    constraints: dict[str, tuple[str, float]] = field(default_factory=dict)
+    # constraints: {concept_name: (operator, threshold)} e.g., {"recursive": (">", 0.5)}
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __repr__(self) -> str:
+        parts = []
+        for name, weight in sorted(self.weights.items(), key=lambda x: -abs(x[1])):
+            if weight > 0:
+                parts.append(f"+{weight:.1f}·{name}")
+            else:
+                parts.append(f"{weight:.1f}·{name}")
+        for name, (op, val) in self.constraints.items():
+            parts.append(f"{name}{op}{val:.1f}")
+        return f"Query({', '.join(parts)})"
+
+    @property
+    def concept_vector(self) -> dict[str, float]:
+        """The query as a concept-space direction vector."""
+        return self.weights.copy()
+
+    def matches(self, concept_scores: dict[str, float]) -> bool:
+        """Check if a program's concept scores satisfy the query constraints."""
+        for name, (op, threshold) in self.constraints.items():
+            score = concept_scores.get(name, 0.0)
+            if op == ">" and not (score > threshold):
+                return False
+            elif op == ">=" and not (score >= threshold):
+                return False
+            elif op == "<" and not (score < threshold):
+                return False
+            elif op == "<=" and not (score <= threshold):
+                return False
+            elif op == "==" and not (abs(score - threshold) < 0.05):
+                return False
+        return True
+
+    def distance_to(self, concept_scores: dict[str, float]) -> float:
+        """
+        Compute distance from a program's concept scores to this query.
+        Lower = more aligned with query.
+        """
+        total = 0.0
+        for name, target_weight in self.weights.items():
+            actual = concept_scores.get(name, 0.0)
+            # Distance weighted by how strongly we care about this concept
+            total += abs(target_weight) * (actual - (1.0 if target_weight > 0 else 0.0)) ** 2
+        return total
+
+
+class QueryLanguage:
+    """
+    Parser and builder for concept queries.
+    
+    Supports a simple DSL:
+      "recursive > 0.5 AND fast_execution > 0.7"
+      "recursive=0.8, space_efficient=0.6, mutation=-0.3"
+      "LIKE program_id_abc123"  (find programs similar to a reference)
+      "NOT mutation"  (avoid mutation)
+    """
+
+    def __init__(self, concept_set: ConceptSet):
+        self.concept_set = concept_set
+
+    def parse(self, query_str: str) -> ConceptQuery:
+        """Parse a query string into a ConceptQuery."""
+        query = ConceptQuery()
+
+        # Handle comma-separated weight assignments: "recursive=0.8, mutation=-0.3"
+        weight_pattern = r"(\w+)\s*=\s*(-?\d+\.?\d*)"
+        for match in re.finditer(weight_pattern, query_str):
+            name = match.group(1)
+            weight = float(match.group(2))
+            if self.concept_set.get_by_name(name):
+                query.weights[name] = weight
+
+        # Handle constraint expressions: "recursive > 0.5"
+        constraint_pattern = r"(\w+)\s*(>=|<=|>|<|==)\s*(-?\d+\.?\d*)"
+        for match in re.finditer(constraint_pattern, query_str):
+            name = match.group(1)
+            op = match.group(2)
+            threshold = float(match.group(3))
+            if name not in query.weights:  # Don't double-count
+                if self.concept_set.get_by_name(name):
+                    query.constraints[name] = (op, threshold)
+
+        # Handle NOT: "NOT mutation"
+        not_pattern = r"NOT\s+(\w+)"
+        for match in re.finditer(not_pattern, query_str, re.IGNORECASE):
+            name = match.group(1)
+            if self.concept_set.get_by_name(name):
+                query.weights[name] = query.weights.get(name, -1.0)
+
+        return query
+
+    def build(self, **concept_weights: float) -> ConceptQuery:
+        """Build a query from keyword arguments."""
+        validated = {}
+        for name, weight in concept_weights.items():
+            if self.concept_set.get_by_name(name):
+                validated[name] = weight
+            else:
+                logger.warning(f"Unknown concept: {name}")
+        return ConceptQuery(weights=validated)
+
+    def constrain(self, **constraints: str) -> ConceptQuery:
+        """
+        Build a constraint query.
+        Example: constrain(recursive=">0.5", fast_execution=">=0.7")
+        """
+        query = ConceptQuery()
+        for name, expr in constraints.items():
+            if not self.concept_set.get_by_name(name):
+                logger.warning(f"Unknown concept: {name}")
+                continue
+            match = re.match(r"(>=|<=|>|<|==)?\s*(-?\d+\.?\d*)", expr)
+            if match:
+                op = match.group(1) or ">"
+                threshold = float(match.group(2))
+                query.constraints[name] = (op, threshold)
+        return query
+
+
+class SteeringEngine:
+    """
+    Engine for steering program generation using concept queries.
+    
+    Two modes:
+      1. Selection: Given a ProgramSpace and a query, rank/filter programs
+      2. Generation: Given a query, steer LLM hidden states during generation
+    
+    Selection uses concept scores directly.
+    Generation uses GCAV vectors (e' = e + ε·v) or MSRS orthogonal steering.
+    """
+
+    def __init__(
+        self,
+        concept_set: ConceptSet,
+        embedding_space: Optional[ConceptEmbeddingSpace] = None,
+        gcav: Optional[GCAVEmbedding] = None,
+        msrs: Optional[MSRSSteering] = None,
+    ):
+        self.concept_set = concept_set
+        self.embedding_space = embedding_space
+        self.gcav = gcav
+        self.msrs = msrs
+        self.query_language = QueryLanguage(concept_set)
+
+    # ---- Selection Mode ----
+
+    def select(
+        self,
+        space: ProgramSpace,
+        query: Union[ConceptQuery, str],
+        top_k: int = 10,
+    ) -> list[tuple[Program, float]]:
+        """
+        Select programs from the space that best match the query.
+        
+        Returns list of (program, relevance_score) tuples, sorted by relevance.
+        """
+        if isinstance(query, str):
+            query = self.query_language.parse(query)
+
+        scored: list[tuple[Program, float]] = []
+
+        for program in space.valid_programs:
+            concept_scores = self.concept_set.score_program(program)
+
+            # Check hard constraints
+            if not query.matches(concept_scores):
+                continue
+
+            # Compute soft relevance score
+            relevance = self._compute_relevance(concept_scores, query)
+            scored.append((program, relevance))
+
+        # Sort by relevance (higher = better match)
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return scored[:top_k]
+
+    def _compute_relevance(
+        self,
+        concept_scores: dict[str, float],
+        query: ConceptQuery,
+    ) -> float:
+        """
+        Compute relevance of a program to a query.
+        
+        For positive weights: reward high concept scores
+        For negative weights: reward low concept scores
+        """
+        relevance = 0.0
+        total_weight = 0.0
+
+        for name, target_weight in query.weights.items():
+            actual = concept_scores.get(name, 0.0)
+            if target_weight > 0:
+                relevance += target_weight * actual
+            else:
+                relevance += abs(target_weight) * (1.0 - actual)
+            total_weight += abs(target_weight)
+
+        if total_weight > 0:
+            relevance /= total_weight
+
+        return relevance
+
+    def filter(
+        self,
+        space: ProgramSpace,
+        query: Union[ConceptQuery, str],
+    ) -> ProgramSpace:
+        """Filter a ProgramSpace to programs matching the query."""
+        if isinstance(query, str):
+            query = self.query_language.parse(query)
+
+        filtered = ProgramSpace(space.stub)
+        for program in space.valid_programs:
+            concept_scores = self.concept_set.score_program(program)
+            if query.matches(concept_scores):
+                filtered.add(program)
+        return filtered
+
+    # ---- Generation Steering Mode ----
+
+    def build_steering_vector(
+        self,
+        query: Union[ConceptQuery, str],
+        method: str = "additive",
+    ) -> Optional[np.ndarray]:
+        """
+        Build a steering vector from a concept query.
+        
+        Args:
+            query: The concept query
+            method: 'additive' (GCAV sum) or 'msrs' (orthogonal subspaces)
+            
+        Returns:
+            Steering vector in activation space, or None if no GCAV available
+        """
+        if isinstance(query, str):
+            query = self.query_language.parse(query)
+
+        if method == "additive" and self.gcav is not None:
+            # Simple additive: v_steer = Σ w_i · v_i
+            steer = np.zeros_like(
+                next(iter(self.gcav.concept_vectors.values()))
+            )
+            for name, weight in query.weights.items():
+                if name in self.gcav.concept_vectors:
+                    steer += weight * self.gcav.concept_vectors[name]
+            return steer
+
+        elif method == "msrs" and self.msrs is not None:
+            # Use MSRS orthogonal steering
+            base = np.zeros(self.msrs.S_align.shape[1])
+            return self.msrs.steer(base, query.weights) - base
+
+        return None
+
+    def steer_prompt(
+        self,
+        query: Union[ConceptQuery, str],
+        base_prompt: str,
+    ) -> str:
+        """
+        Augment a generation prompt with concept-steering instructions.
+        
+        This is a lightweight steering approach that works with any LLM API
+        (no hidden state access needed). For stronger steering, use
+        build_steering_vector with activation-level intervention.
+        """
+        if isinstance(query, str):
+            query = self.query_language.parse(query)
+
+        concept_instructions = []
+        for name, weight in sorted(
+            query.weights.items(), key=lambda x: -abs(x[1])
+        ):
+            concept = self.concept_set.get_by_name(name)
+            if concept is None:
+                continue
+
+            if weight > 0.5:
+                concept_instructions.append(
+                    f"STRONGLY PREFER: {concept.description}"
+                )
+            elif weight > 0:
+                concept_instructions.append(
+                    f"PREFER: {concept.description}"
+                )
+            elif weight < -0.5:
+                concept_instructions.append(
+                    f"STRONGLY AVOID: {concept.description}"
+                )
+            elif weight < 0:
+                concept_instructions.append(
+                    f"AVOID: {concept.description}"
+                )
+
+        for name, (op, threshold) in query.constraints.items():
+            concept = self.concept_set.get_by_name(name)
+            if concept:
+                concept_instructions.append(
+                    f"CONSTRAINT: {concept.description} ({op} {threshold})"
+                )
+
+        if not concept_instructions:
+            return base_prompt
+
+        steering_block = "\n".join(
+            f"  - {inst}" for inst in concept_instructions
+        )
+        return (
+            f"{base_prompt}\n\n"
+            f"CONCEPT STEERING INSTRUCTIONS:\n{steering_block}\n\n"
+            f"Follow the above concept preferences when implementing."
+        )
+
+    # ---- Exploration Mode ----
+
+    def explore_neighbors(
+        self,
+        program: Program,
+        space: ProgramSpace,
+        n_neighbors: int = 5,
+    ) -> list[tuple[Program, float, dict[str, float]]]:
+        """
+        Find programs in the space that are conceptually nearby.
+        
+        Returns list of (program, distance, concept_diff) tuples.
+        concept_diff shows which concepts differ most.
+        """
+        ref_scores = self.concept_set.score_program(program)
+
+        neighbors: list[tuple[Program, float, dict[str, float]]] = []
+        for other in space.valid_programs:
+            if other.program_id == program.program_id:
+                continue
+
+            other_scores = self.concept_set.score_program(other)
+
+            # Euclidean distance in concept space
+            diff = {}
+            dist_sq = 0.0
+            for name in set(ref_scores) | set(other_scores):
+                d = ref_scores.get(name, 0.0) - other_scores.get(name, 0.0)
+                if abs(d) > 0.01:
+                    diff[name] = d
+                dist_sq += d ** 2
+
+            neighbors.append((other, dist_sq ** 0.5, diff))
+
+        neighbors.sort(key=lambda x: x[1])
+        return neighbors[:n_neighbors]
+
+    def concept_boundary_programs(
+        self,
+        concept_name: str,
+        space: ProgramSpace,
+        n_per_side: int = 3,
+    ) -> dict[str, list[Program]]:
+        """
+        Find programs at the boundary of a concept.
+        Returns programs just inside and just outside the concept region.
+        """
+        concept = self.concept_set.get_by_name(concept_name)
+        if concept is None:
+            return {"inside": [], "outside": []}
+
+        inside: list[tuple[Program, float]] = []
+        outside: list[tuple[Program, float]] = []
+
+        for program in space.valid_programs:
+            score = concept.score(program)
+            if score > 0.5:
+                inside.append((program, score))
+            else:
+                outside.append((program, score))
+
+        # Sort inside by score ascending (closest to boundary)
+        inside.sort(key=lambda x: x[1])
+        # Sort outside by score descending (closest to boundary)
+        outside.sort(key=lambda x: x[1], reverse=True)
+
+        return {
+            "inside": [p for p, _ in inside[:n_per_side]],
+            "outside": [p for p, _ in outside[:n_per_side]],
+        }