Lenses

from qlaw.lenses.semantic import SemanticInterpreter
from qlaw.lenses.decompose import Decomposer
from qlaw.lenses.ontology import OntologyArchitect
from qlaw.lenses.trajectories import TrajectoryStrategist
from qlaw.lenses.critic import Critic
from qlaw.lenses.enrich import GroundingEnricher
from qlaw.lenses.explain import Explainer

from qlaw.graph import NodeType

ALLOWED_DECOMPOSE_TYPES = {NodeType.COMPONENT, NodeType.QUESTION, NodeType.DATA}
ALLOWED_ONTOLOGY_TYPES = {NodeType.CONCEPT}

def _basic_batch_reward(kwargs, outputs, *, allowed_types, min_items=3, max_items=5):
    items = outputs.sub_nodes if hasattr(outputs, "sub_nodes") else outputs.concepts
    if not items: return 0.0
    if not (min_items <= len(items) <= max_items): return 0.0
    if any(not (getattr(i, "label", "") or "").strip() for i in items): return 0.0
    if any(not (getattr(i, "description", "") or "").strip() for i in items): return 0.0
    if any(getattr(i, "type", None) not in allowed_types for i in items): return 0.0
    return 1.0

def validate_decompose(kwargs, outputs) -> float:
    return _basic_batch_reward(kwargs, outputs, allowed_types=ALLOWED_DECOMPOSE_TYPES)

def validate_ontology(kwargs, outputs) -> float:
    return _basic_batch_reward(kwargs, outputs, allowed_types=ALLOWED_ONTOLOGY_TYPES)

import dspy
from qlaw.signatures import SemanticInterpretation

class SemanticInterpreter(dspy.Module):
    """Prompt -> root seed. Replaces analyzePrompt(); model: deepseek-v4-flash."""
    def __init__(self):
        super().__init__()
        self.predict = dspy.ChainOfThought(SemanticInterpretation)

    def forward(self, prompt: str, **kwargs):
        return self.predict(prompt=prompt)  # Prediction: root_label, root_type, intent, entities

import dspy
from qlaw.signatures import Decompose
from qlaw.lenses._validators import validate_decompose

class Decomposer(dspy.Module):
    """Break a node into 3-5 parts. Replaces decomposeNode(); model: deepseek-v4-flash.
    Refine enforces the 3-5 / valid-type contract with retry + feedback."""
    def __init__(self):
        super().__init__()
        self.predict = dspy.ChainOfThought(Decompose)
        self.refined = dspy.Refine(self.predict, N=3, reward_fn=validate_decompose, threshold=0.5)

    def forward(self, node, **kwargs):
        return self.refined(node=node)  # Prediction: sub_nodes