aura-state 0.2.0__py3-none-any.whl

aura_state/__init__.py

@@ -0,0 +1,126 @@
+"""
+Aura-State: LLM State Machine Compiler.
+
+Usage:
+    from aura_state import AuraEngine, Node, CompiledTransition
+"""
+from .core.engine import AuraEngine, Node, CompiledTransition
+from .execution.tracer import AuraTrace
+from .compiler.dspy_tuner import BootstrapTeleprompter
+from .execution.sandbox import SandboxedInterpreter, SandboxExecutionError
+from .memory.pruner import ContextPruner
+from .consensus.auto_vote import ConsensusStrategy
+from .core.exceptions import AuraStateError, StateTransitionError, MaxRetriesExceededError
+from .loaders.json_graph import JSONGraphLoader
+
+# ── Core Innovations ──
+from .core.adaptive_graph import AdaptiveDAG, NodeHealthMetrics, EdgeStats
+from .core.verification_loop import VerificationLoop, ReflectionMemory, Reflection
+from .core.providers import LLMProvider, CostTracker
+from .compiler.schema_compiler import compile_schema, compile_openapi_schemas, levenshtein_distance
+
+from .verification.temporal_verifier import (
+    compile_kripke, verify_engine, verify_property,
+    reachability, always_before, mutual_exclusion, eventual_completion,
+    PropertyResult, VerificationResult,
+)
+from .verification.conformal import (
+    conformal_interval, conformal_from_extractions,
+    PredictionInterval, ConformalResult,
+)
+from .verification.proof_engine import (
+    prove_extraction, prove_consistency, ProofResult,
+    prove_obligations_satisfiable, field_bounds_from_model, SatResult,
+)
+from .verification.pipeline_conformal import PipelineConformal
+from .compiler.spec_compiler import (
+    compile_contract, diff_contracts, check_faithfulness,
+    AuraContract, NodeContract, PropertyVerdict, TaintContract, TaintPath, ContractError,
+)
+from .verification.taint import (
+    analyze_taint, TaintResult, TaintViolation,
+    analyze_field_taint, FieldTaintResult, FieldTaintViolation,
+)
+from .verification.risk_control import RiskController, learn_then_test
+from .core.replan import (
+    counterexample_guided_repair, default_repair,
+    ctl_to_repair, z3_to_repair, taint_to_repair,
+    RepairSignal, ReplanResult, RepairIteration,
+)
+
+__all__ = [
+    "AuraEngine",
+    "Node",
+    "CompiledTransition",
+    "AuraTrace",
+    "BootstrapTeleprompter",
+    "AdaptiveDAG",
+    "NodeHealthMetrics",
+    "EdgeStats",
+    "VerificationLoop",
+    "ReflectionMemory",
+    "Reflection",
+    "LLMProvider",
+    "CostTracker",
+    "compile_schema",
+    "compile_openapi_schemas",
+    "levenshtein_distance",
+    "ContextPruner",
+    "ConsensusStrategy",
+    "SandboxedInterpreter",
+    "SandboxExecutionError",
+    "JSONGraphLoader",
+    "AuraStateError",
+    "StateTransitionError",
+    "MaxRetriesExceededError",
+    # Formal Verification
+    "compile_kripke",
+    "verify_engine",
+    "verify_property",
+    "reachability",
+    "always_before",
+    "mutual_exclusion",
+    "eventual_completion",
+    "PropertyResult",
+    "VerificationResult",
+    "conformal_interval",
+    "conformal_from_extractions",
+    "PredictionInterval",
+    "ConformalResult",
+    "prove_extraction",
+    "prove_consistency",
+    "ProofResult",
+    "prove_obligations_satisfiable",
+    "field_bounds_from_model",
+    "SatResult",
+    "PipelineConformal",
+    # Design→spec compiler
+    "compile_contract",
+    "diff_contracts",
+    "check_faithfulness",
+    "AuraContract",
+    "NodeContract",
+    "PropertyVerdict",
+    "TaintContract",
+    "TaintPath",
+    "ContractError",
+    # Capability-typed dataflow (injection-proof)
+    "analyze_taint",
+    "TaintResult",
+    "TaintViolation",
+    "analyze_field_taint",
+    "FieldTaintResult",
+    "FieldTaintViolation",
+    # Risk-controlled abstention
+    "RiskController",
+    "learn_then_test",
+    # Counterexample-guided replanning
+    "counterexample_guided_repair",
+    "default_repair",
+    "ctl_to_repair",
+    "z3_to_repair",
+    "taint_to_repair",
+    "RepairSignal",
+    "ReplanResult",
+    "RepairIteration",
+]

aura_state/compiler/__init__.py

@@ -0,0 +1 @@
+# Compiler sub-package

aura_state/compiler/dspy_tuner.py

@@ -0,0 +1,124 @@
+import json
+import logging
+import math
+from typing import Callable, List, Dict, Any, Optional
+
+logger = logging.getLogger("aura_state")
+
+
+def char_stub_embedder(text: str) -> List[float]:
+    """Deterministic character-based pseudo-embedding for OFFLINE TESTS ONLY.
+
+    Ranks by byte value, not semantics -- meaningless for real similarity.
+    Never used by default; must be injected explicitly via
+    ``BootstrapTeleprompter(embedder=char_stub_embedder)``.
+    """
+    vec = [0.0] * 256
+    for i, char in enumerate(text[:256]):
+        vec[i] = float(ord(char))
+    mag = math.sqrt(sum(v * v for v in vec))
+    if mag > 0:
+        vec = [v / mag for v in vec]
+    return vec
+
+
+def _openai_embedder(client, model: str = "text-embedding-3-small") -> Callable[[str], List[float]]:
+    """Build a real embedding function backed by an OpenAI client."""
+    def embed(text: str) -> List[float]:
+        resp = client.embeddings.create(model=model, input=text)
+        return list(resp.data[0].embedding)
+    return embed
+
+
+class BootstrapTeleprompter:
+    """
+    KNN-based few-shot optimizer.
+
+    Finds the K most similar successful past executions and injects them
+    as few-shot demonstrations into the system prompt.
+
+    An embedder is required to compute similarity. Provide one of:
+      - ``embedder``: any ``Callable[[str], List[float]]``
+      - ``openai_client``: a real embedding model is used
+    If neither is given, embedding fails loud (no meaningless default).
+    """
+    def __init__(
+        self,
+        k_neighbors: int = 3,
+        embedder: Optional[Callable[[str], List[float]]] = None,
+        openai_client: Optional[Any] = None,
+    ):
+        self.k = k_neighbors
+        self.successful_traces: Dict[str, List[Dict[str, Any]]] = {}
+        if embedder is not None:
+            self._embedder = embedder
+        elif openai_client is not None:
+            self._embedder = _openai_embedder(openai_client)
+        else:
+            self._embedder = None
+
+    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
+        dot_product = sum(a * b for a, b in zip(vec1, vec2))
+        magnitude1 = math.sqrt(sum(a * a for a in vec1))
+        magnitude2 = math.sqrt(sum(b * b for b in vec2))
+        if magnitude1 == 0 or magnitude2 == 0:
+            return 0.0
+        return dot_product / (magnitude1 * magnitude2)
+
+    def _embed(self, text: str) -> List[float]:
+        if self._embedder is None:
+            raise RuntimeError(
+                "BootstrapTeleprompter has no embedder. Pass embedder=... or "
+                "openai_client=... (or char_stub_embedder for offline tests). "
+                "The old ord()-based default was removed -- it ranked by byte "
+                "value, not meaning."
+            )
+        return self._embedder(text)
+
+    def _bootstrap_dataset(self, dataset: List[Dict[str, Any]]):
+        """Filters the dataset to keep only successful executions."""
+        for trace in dataset:
+            if trace.get("success"):
+                node_name = trace["node"]
+                if node_name not in self.successful_traces:
+                    self.successful_traces[node_name] = []
+                    
+                trace["embedding"] = self._embed(trace["input"])
+                self.successful_traces[node_name].append(trace)
+
+    def optimize_node(self, node_name: str, current_prompt: str, new_user_input: str) -> str:
+        """
+        Finds the K most similar past successes and appends them
+        as few-shot examples to the system prompt.
+        """
+        if node_name not in self.successful_traces or not self.successful_traces[node_name]:
+            return current_prompt
+            
+        # O(N) KNN Search
+        query_vec = self._embed(new_user_input)
+        scored_traces = []
+        for trace in self.successful_traces[node_name]:
+            score = self._cosine_similarity(query_vec, trace["embedding"])
+            scored_traces.append((score, trace))
+            
+        # Sort by similarity descending, take Top K
+        scored_traces.sort(key=lambda x: x[0], reverse=True)
+        top_k = [x[1] for x in scored_traces[:self.k]]
+        
+        # Abstract the Few-Shot injection prompt
+        few_shot_block = "\n\n--- FEW-SHOT DEMONSTRATIONS ---\n"
+        few_shot_block += "Follow the structure of these examples:\n\n"
+        for i, trace in enumerate(top_k):
+            few_shot_block += f"EXAMPLE {i+1}:\n"
+            few_shot_block += f"User Context: {trace['input']}\n"
+            few_shot_block += f"Extracted Schema Validation Response:\n{json.dumps(trace['output'], indent=2)}\n\n"
+            
+        optimized_prompt = current_prompt + few_shot_block
+        logger.info(f"Injected {len(top_k)} few-shot examples into node '{node_name}'")
+        return optimized_prompt
+        
+    def compile(self, dataset: List[Dict[str, Any]]):
+        """Loads and indexes the dataset for KNN lookup."""
+        self._bootstrap_dataset(dataset)
+        total = sum(len(v) for v in self.successful_traces.values())
+        logger.info(f"Indexed {total} successful traces across {len(self.successful_traces)} nodes")

aura_state/compiler/json_generator.py

@@ -0,0 +1,69 @@
+import json
+from typing import Dict, Type, List, Any
+from pydantic import BaseModel
+
+from ..core.engine import Node, CompiledTransition
+
+
+def _pydantic_to_json_extracts(model: Type[BaseModel]) -> Dict[str, Any]:
+    """Converts a Pydantic model into the 'extracts' dict format expected by JSONGraphLoader."""
+    extracts = {}
+    
+    type_names = {
+        str: "str",
+        int: "int",
+        float: "float",
+        bool: "bool",
+        list[str]: "list[str]",
+        list[int]: "list[int]"
+    }
+    
+    for field_name, model_field in model.model_fields.items():
+        field_type_str = type_names.get(model_field.annotation, "str")
+        extracts[field_name] = {
+            "type": field_type_str,
+            "description": model_field.description or ""
+        }
+        
+    return extracts
+
+
+def generate_flow_json(nodes: Dict[str, Type[Node]], transitions: List[CompiledTransition], output_path: str = "flow.json"):
+    """
+    The execution compiler.
+    Translates declarative Python Node classes down into the universal flow.json DAG artifact.
+    """
+    graph = {
+        "nodes": [],
+        "edges": []
+    }
+    
+    # 1. Compile Nodes
+    for name, node_cls in nodes.items():
+        node_def = {
+            "id": name,
+            "system_prompt": node_cls.system_prompt,
+        }
+        
+        if node_cls.extracts:
+            node_def["extracts"] = _pydantic_to_json_extracts(node_cls.extracts)
+            
+        if node_cls.sandbox_rule:
+            node_def["sandbox_rule"] = node_cls.sandbox_rule
+            
+        graph["nodes"].append(node_def)
+        
+    # 2. Compile Edges
+    for t in transitions:
+        edge = {
+            "from": t.from_node.__name__,
+            "to": t.to_node.__name__,
+            "condition": t.condition 
+        }
+        graph["edges"].append(edge)
+        
+    # 3. Serialize to disk
+    with open(output_path, "w") as f:
+        json.dump(graph, f, indent=4)
+        
+    return graph

aura_state/compiler/schema_compiler.py

@@ -0,0 +1,221 @@
+"""
+Schema Compiler: JSON Schema → Node classes with Pydantic extraction models.
+
+Given a JSON Schema or OpenAPI spec, generates Node subclasses with
+auto-generated system prompts and Levenshtein fuzzy matching for field correction.
+"""
+import logging
+from typing import Any, Dict, List, Optional, Type
+from pydantic import BaseModel, Field, create_model
+
+logger = logging.getLogger("aura_state")
+
+
+# --- Levenshtein Distance ---
+
+def levenshtein_distance(a: str, b: str) -> int:
+    """Compute Levenshtein edit distance between two strings."""
+    la, lb = len(a), len(b)
+    if la == 0:
+        return lb
+    if lb == 0:
+        return la
+    
+    prev = list(range(lb + 1))
+    curr = [0] * (lb + 1)
+    
+    for i in range(1, la + 1):
+        curr[0] = i
+        for j in range(1, lb + 1):
+            cost = 0 if a[i - 1] == b[j - 1] else 1
+            ins = prev[j] + 1
+            delete = curr[j - 1] + 1
+            sub = prev[j - 1] + cost
+            curr[j] = min(ins, delete, sub)
+        prev, curr = curr, prev
+    
+    return prev[lb]
+
+
+def suggest_field(name: str, available: List[str], max_distance: int = 3) -> Optional[str]:
+    """
+    Find the closest field name using Levenshtein distance.
+    Returns None if no match within max_distance.
+    """
+    if not available:
+        return None
+    
+    best_dist = -1
+    best_name = ""
+    
+    for field_name in available:
+        d = levenshtein_distance(name.lower(), field_name.lower())
+        if best_dist < 0 or d < best_dist:
+            best_dist = d
+            best_name = field_name
+    
+    if 0 <= best_dist <= max_distance:
+        return best_name
+    return None
+
+
+# --- JSON Schema → Pydantic Model ---
+
+# JSON Schema types → Python types
+SCHEMA_TYPE_MAP = {
+    "string": str,
+    "integer": int,
+    "number": float,
+    "boolean": bool,
+    "array": list,
+}
+
+
+def _resolve_type(prop: Dict[str, Any]) -> type:
+    """Resolve a JSON Schema property to a Python type."""
+    schema_type = prop.get("type", "string")
+    
+    if schema_type == "array":
+        items_type = prop.get("items", {}).get("type", "string")
+        inner = SCHEMA_TYPE_MAP.get(items_type, str)
+        return List[inner]
+    
+    return SCHEMA_TYPE_MAP.get(schema_type, str)
+
+
+def compile_pydantic_model(
+    model_name: str,
+    schema: Dict[str, Any],
+) -> Type[BaseModel]:
+    """
+    Compile a JSON Schema object into a Pydantic model.
+    
+    Supports:
+    - type mapping (string, integer, number, boolean, array)
+    - descriptions → Field(description=...)
+    - required fields
+    - default values
+    - enum constraints
+    - minimum/maximum validation
+    """
+    properties = schema.get("properties", {})
+    required = set(schema.get("required", []))
+    
+    fields = {}
+    for field_name, prop in properties.items():
+        python_type = _resolve_type(prop)
+        description = prop.get("description", "")
+        default_value = prop.get("default", ...)
+        
+        # Build Field kwargs
+        field_kwargs = {"description": description}
+        
+        if "enum" in prop:
+            field_kwargs["json_schema_extra"] = {"enum": prop["enum"]}
+        
+        if "minimum" in prop:
+            field_kwargs["ge"] = prop["minimum"]
+        if "maximum" in prop:
+            field_kwargs["le"] = prop["maximum"]
+        
+        if field_name in required:
+            fields[field_name] = (python_type, Field(default=..., **field_kwargs))
+        elif default_value is not ...:
+            fields[field_name] = (python_type, Field(default=default_value, **field_kwargs))
+        else:
+            fields[field_name] = (Optional[python_type], Field(default=None, **field_kwargs))
+    
+    return create_model(model_name, **fields)
+
+
+# --- Schema → Node ---
+
+def _generate_system_prompt(model_name: str, schema: Dict[str, Any]) -> str:
+    """Auto-generate a system prompt from JSON Schema field descriptions."""
+    properties = schema.get("properties", {})
+    title = schema.get("title", model_name)
+    desc = schema.get("description", f"Extract data for {title}")
+    
+    lines = [desc, "", "Extract the following fields from the user input:"]
+    for field_name, prop in properties.items():
+        field_desc = prop.get("description", "")
+        field_type = prop.get("type", "string")
+        constraint = ""
+        
+        if "enum" in prop:
+            constraint = f" (one of: {', '.join(str(v) for v in prop['enum'])})"
+        if "minimum" in prop and "maximum" in prop:
+            constraint = f" (between {prop['minimum']} and {prop['maximum']})"
+        
+        lines.append(f"- {field_name} ({field_type}): {field_desc}{constraint}")
+    
+    return "\n".join(lines)
+
+
+def compile_schema(
+    schema: Dict[str, Any],
+    node_name: Optional[str] = None,
+) -> type:
+    """
+    Compile a JSON Schema into a Node subclass with:
+    - Auto-generated Pydantic extraction model
+    - Auto-generated system prompt from field descriptions
+    - Levenshtein fuzzy matching for field correction
+    """
+    # Import here to avoid circular dependency
+    from ..core.engine import Node
+    
+    name = node_name or schema.get("title", "CompiledNode")
+    model = compile_pydantic_model(f"{name}Extraction", schema)
+    prompt = _generate_system_prompt(name, schema)
+    
+    expected_fields = list(schema.get("properties", {}).keys())
+    
+    def _handle(self, user_text, extracted_data=None, memory=None):
+        if extracted_data:
+            data = extracted_data.model_dump()
+            # Fuzzy field correction: if LLM returned unexpected key names,
+            # try Levenshtein matching to recover
+            corrected = {}
+            for key, value in data.items():
+                if key in expected_fields:
+                    corrected[key] = value
+                else:
+                    suggestion = suggest_field(key, expected_fields)
+                    if suggestion:
+                        logger.info(f"[SchemaCompiler] Fuzzy match: '{key}' → '{suggestion}'")
+                        corrected[suggestion] = value
+                    else:
+                        corrected[key] = value
+            
+            return "END", {"extracted": corrected, "source": "schema_compiled"}
+        return "END", {"extracted": {}, "source": "schema_compiled"}
+    
+    node_cls = type(name, (Node,), {
+        "system_prompt": prompt,
+        "extracts": model,
+        "handle": _handle,
+    })
+    
+    return node_cls
+
+
+def compile_openapi_schemas(
+    spec: Dict[str, Any],
+) -> List[type]:
+    """
+    Compile all schemas from an OpenAPI spec into Node subclasses.
+    
+    Each schema in components/schemas becomes a Node with an
+    auto-generated extraction model and system prompt.
+    """
+    schemas = spec.get("components", {}).get("schemas", {})
+    nodes = []
+    
+    for name, schema in schemas.items():
+        if schema.get("type") == "object" and "properties" in schema:
+            node = compile_schema(schema, node_name=name)
+            nodes.append(node)
+            logger.info(f"[SchemaCompiler] Compiled OpenAPI schema '{name}' → Node class")
+    
+    return nodes