synkro 0.4.36__py3-none-any.whl

synkro/generation/generator.py

@@ -0,0 +1,314 @@
+"""Main Generator class orchestrating the full trace generation pipeline."""
+
+import asyncio
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from synkro.llm.client import LLM
+from synkro.llm.rate_limits import auto_workers
+from synkro.models import Model, OpenAI
+from synkro.types.dataset_type import DatasetType
+from synkro.core.policy import Policy
+from synkro.core.dataset import Dataset
+from synkro.core.checkpoint import CheckpointManager, hash_policy
+from synkro.modes.config import get_mode_config
+from synkro.errors import handle_error
+from synkro.factory import ComponentFactory
+from synkro.reporting import ProgressReporter, RichReporter
+from synkro.pipeline.runner import GenerationPipeline, GenerationResult, ScenariosResult
+
+if TYPE_CHECKING:
+    from synkro.types.tool import ToolDefinition
+
+
+class Generator:
+    """
+    Main orchestrator for generating training datasets.
+
+    The Generator handles the full pipeline:
+    1. Plan: Analyze policy and create category distribution
+    2. Generate: Create scenarios and responses
+    3. Grade: Evaluate response quality
+    4. Refine: Fix failed responses
+    5. Return: Dataset of passing traces
+
+    Examples:
+        >>> generator = Generator()
+        >>> dataset = generator.generate(policy, traces=20)
+
+        >>> # Conversation dataset (default, multi-turn)
+        >>> generator = Generator(dataset_type=DatasetType.CONVERSATION)
+        >>> dataset = generator.generate(policy)
+
+        >>> # Instruction dataset (single-turn)
+        >>> generator = Generator(dataset_type=DatasetType.INSTRUCTION)
+        >>> dataset = generator.generate(policy)
+
+        >>> # Silent mode (no console output)
+        >>> from synkro.reporting import SilentReporter
+        >>> generator = Generator(reporter=SilentReporter())
+        >>> dataset = generator.generate(policy)
+        
+        >>> # Tool call dataset
+        >>> from synkro import ToolDefinition
+        >>> tools = [ToolDefinition(name="search", description="...", parameters={})]
+        >>> generator = Generator(dataset_type=DatasetType.TOOL_CALL, tools=tools)
+        >>> dataset = generator.generate("Usage guidelines", traces=20)
+
+        >>> # Eval dataset with low temperature for deterministic outputs
+        >>> generator = Generator(dataset_type=DatasetType.EVALUATION, temperature=0.2)
+        >>> dataset = generator.generate(policy, traces=50)
+    """
+
+    def __init__(
+        self,
+        dataset_type: DatasetType = DatasetType.CONVERSATION,
+        generation_model: Model = OpenAI.GPT_4O_MINI,
+        grading_model: Model = OpenAI.GPT_4O,
+        max_iterations: int = 1,
+        skip_grading: bool = False,
+        reporter: ProgressReporter | None = None,
+        tools: list["ToolDefinition"] | None = None,
+        turns: int | str = "auto",
+        checkpoint_dir: str | Path | None = None,
+        enable_hitl: bool = True,
+        base_url: str | None = None,
+        thinking: bool = False,
+        temperature: float = 0.7,
+    ):
+        """
+        Initialize the Generator.
+
+        Args:
+            dataset_type: Type of dataset to generate (CONVERSATION, INSTRUCTION, or TOOL_CALL)
+            generation_model: Model for scenarios/responses (default: gpt-4o-mini)
+            grading_model: Model for grading (default: gpt-4o, recommend stronger)
+            max_iterations: Max refinement iterations per trace (default: 1, no retries)
+            skip_grading: Skip grading phase for faster generation (default: False)
+            reporter: Progress reporter (default: RichReporter for console output)
+            tools: List of ToolDefinition for TOOL_CALL dataset type
+            turns: Conversation turns per trace. Use int for fixed turns, or "auto"
+                for policy complexity-driven turns (Simple=1-2, Conditional=3, Complex=5+)
+            checkpoint_dir: Directory for checkpoints. If provided, enables resumable
+                generation. Progress is saved after each stage.
+            enable_hitl: Enable Human-in-the-Loop Logic Map editing. When enabled,
+                pauses after Logic Map extraction to allow interactive refinement.
+            base_url: Optional API base URL for local LLM providers (Ollama, vLLM, etc.)
+            thinking: Enable thinking mode with <think> tags in responses (default: False).
+                When enabled, assistant responses will include reasoning wrapped in
+                <think>...</think> tags, compatible with Qwen3 and DeepSeek-R1 formats.
+            temperature: Sampling temperature for generation (0.0-2.0, default: 0.7).
+                Lower values (0.1-0.3) produce more deterministic outputs for eval datasets.
+                Higher values (0.7-1.0) produce more diverse outputs for training data.
+        """
+        self.dataset_type = dataset_type
+        self.mode_config = get_mode_config(dataset_type)
+        self.max_iterations = max_iterations
+        self.skip_grading = skip_grading
+        self.tools = tools
+        self.turns = turns
+        self.thinking = thinking
+        self.checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None
+
+        # Create checkpoint manager if checkpointing enabled
+        self.checkpoint_manager = (
+            CheckpointManager(self.checkpoint_dir) if self.checkpoint_dir else None
+        )
+
+        # HITL configuration
+        self.enable_hitl = enable_hitl
+
+        # Validate tools for TOOL_CALL dataset type
+        if dataset_type == DatasetType.TOOL_CALL and not tools:
+            raise ValueError("TOOL_CALL dataset type requires tools parameter")
+
+        # Force turns=1 for INSTRUCTION and EVALUATION types
+        if dataset_type in (DatasetType.INSTRUCTION, DatasetType.EVALUATION):
+            self.turns = 1
+
+        # Store model info for reporting
+        self.generation_model = generation_model
+        self.grading_model = grading_model
+
+        # Create LLM clients
+        self.generation_llm = LLM(model=generation_model, base_url=base_url, temperature=temperature)
+        self.grading_llm = LLM(model=grading_model, base_url=base_url)
+        
+        # Create factory for component creation
+        self.factory = ComponentFactory(
+            generation_llm=self.generation_llm,
+            grading_llm=self.grading_llm,
+            mode_config=self.mode_config,
+            tools=tools,
+            thinking=thinking,
+        )
+
+        # Reporter for progress output
+        self.reporter = reporter or RichReporter()
+
+        # Auto-scale workers based on provider
+        model_str = generation_model.value if isinstance(generation_model, Enum) else str(generation_model)
+        self.workers = auto_workers(model_str)
+
+        # Create HITL editors if enabled
+        hitl_editor = self.factory.create_logic_map_editor() if enable_hitl else None
+        scenario_editor = self.factory.create_scenario_editor() if enable_hitl else None
+
+        # Create pipeline
+        self.pipeline = GenerationPipeline(
+            factory=self.factory,
+            reporter=self.reporter,
+            workers=self.workers,
+            max_iterations=max_iterations,
+            skip_grading=skip_grading,
+            checkpoint_manager=self.checkpoint_manager,
+            enable_hitl=enable_hitl,
+            hitl_editor=hitl_editor,
+            scenario_editor=scenario_editor,
+        )
+
+    @handle_error
+    def generate(
+        self,
+        policy: Policy | str,
+        traces: int = 20,
+        return_logic_map: bool = False,
+    ) -> Dataset | GenerationResult:
+        """
+        Generate a training dataset from a policy.
+
+        Args:
+            policy: Policy object or text string
+            traces: Target number of traces to generate (default: 20)
+            return_logic_map: If True, return GenerationResult with access to
+                the Logic Map, scenarios, and distribution (default: False)
+
+        Returns:
+            Dataset (default) or GenerationResult if return_logic_map=True
+
+        Examples:
+            >>> # Standard usage
+            >>> dataset = generator.generate(policy, traces=50)
+
+            >>> # Access Logic Map for inspection
+            >>> result = generator.generate(policy, return_logic_map=True)
+            >>> print(result.logic_map.rules)  # See extracted rules
+            >>> print(result.distribution)     # See scenario type counts
+            >>> dataset = result.dataset       # Get the dataset
+        """
+        if isinstance(policy, str):
+            policy = Policy(text=policy)
+
+        # Validate policy has enough content
+        policy.validate_length()
+
+        return asyncio.run(self._generate_async(policy, traces, return_logic_map))
+
+    async def _generate_async(
+        self,
+        policy: Policy,
+        traces: int,
+        return_logic_map: bool = False,
+    ) -> Dataset | GenerationResult:
+        """Async implementation of generation pipeline."""
+        model_str = self.generation_model.value if isinstance(self.generation_model, Enum) else str(self.generation_model)
+
+        return await self.pipeline.run(
+            policy=policy,
+            traces=traces,
+            model=model_str,
+            dataset_type=self.dataset_type.value,
+            turns=self.turns,
+            return_result=return_logic_map,
+        )
+
+    async def generate_async(
+        self,
+        policy: Policy | str,
+        traces: int = 20,
+        return_logic_map: bool = False,
+    ) -> Dataset | GenerationResult:
+        """
+        Async version of generate for use in async contexts.
+
+        Args:
+            policy: Policy object or text string
+            traces: Target number of traces to generate (default: 20)
+            return_logic_map: If True, return GenerationResult with Logic Map access
+
+        Returns:
+            Dataset (default) or GenerationResult if return_logic_map=True
+        """
+        if isinstance(policy, str):
+            policy = Policy(text=policy)
+
+        return await self._generate_async(policy, traces, return_logic_map)
+
+    @handle_error
+    def generate_scenarios(
+        self,
+        policy: Policy | str,
+        count: int = 20,
+    ) -> ScenariosResult:
+        """
+        Generate eval scenarios without synthetic responses.
+
+        This runs stages 0-2 of the pipeline (planning, logic extraction,
+        scenario synthesis) but skips response generation. Use this for
+        creating eval datasets where you want to test your own model.
+
+        Args:
+            policy: Policy object or text string
+            count: Target number of scenarios to generate (default: 20)
+
+        Returns:
+            ScenariosResult with scenarios, logic_map, and distribution
+
+        Examples:
+            >>> result = generator.generate_scenarios(policy, count=100)
+            >>> for scenario in result.scenarios:
+            ...     response = my_model(scenario.user_message)
+            ...     grade = synkro.grade(response, scenario, policy)
+        """
+        if isinstance(policy, str):
+            policy = Policy(text=policy)
+
+        # Validate policy has enough content
+        policy.validate_length()
+
+        return asyncio.run(self._generate_scenarios_async(policy, count))
+
+    async def _generate_scenarios_async(
+        self,
+        policy: Policy,
+        count: int,
+    ) -> ScenariosResult:
+        """Async implementation of scenario-only generation."""
+        model_str = self.generation_model.value if isinstance(self.generation_model, Enum) else str(self.generation_model)
+
+        return await self.pipeline.run_scenarios_only(
+            policy=policy,
+            count=count,
+            model=model_str,
+        )
+
+    async def generate_scenarios_async(
+        self,
+        policy: Policy | str,
+        count: int = 20,
+    ) -> ScenariosResult:
+        """
+        Async version of generate_scenarios for use in async contexts.
+
+        Args:
+            policy: Policy object or text string
+            count: Target number of scenarios to generate (default: 20)
+
+        Returns:
+            ScenariosResult with scenarios, logic_map, and distribution
+        """
+        if isinstance(policy, str):
+            policy = Policy(text=policy)
+
+        return await self._generate_scenarios_async(policy, count)

synkro/generation/golden_responses.py

@@ -0,0 +1,269 @@
+"""Golden Response Generator - The Thinker.
+
+Generates traces with grounded Chain-of-Thought reasoning and rule citations.
+This is Stage 3 of the Golden Trace pipeline for CONVERSATION/INSTRUCTION datasets.
+"""
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.schemas import GoldenTraceOutput
+from synkro.types.core import Trace, Message, Scenario
+from synkro.types.logic_map import (
+    LogicMap,
+    GoldenScenario,
+    ReasoningStep,
+)
+from synkro.prompts.golden_templates import (
+    GOLDEN_TRACE_PROMPT,
+    GOLDEN_TRACE_MULTI_TURN_PROMPT,
+)
+
+
+class GoldenResponseGenerator:
+    """
+    The Thinker - Generates traces with grounded reasoning.
+
+    Produces traces with:
+    - Explicit Chain-of-Thought reasoning
+    - Rule citations (Rule IDs) for each reasoning step
+    - Exclusionary reasoning (why rules DON'T apply)
+    - DAG-compliant dependency order
+
+    Examples:
+        >>> generator = GoldenResponseGenerator(llm=LLM(model=OpenAI.GPT_4O_MINI))
+        >>> trace = await generator.generate_single(
+        ...     policy_text="...",
+        ...     logic_map=logic_map,
+        ...     scenario=scenario,
+        ...     target_turns=1,
+        ... )
+    """
+
+    # Instruction to inject when thinking mode is enabled
+    THINKING_INSTRUCTION = """
+THINKING MODE:
+Your assistant response MUST include reasoning wrapped in <think> and </think> tags.
+Place your step-by-step reasoning inside the think tags BEFORE your actual response.
+
+Format:
+<think>
+[Your reasoning about which rules apply, why they apply/don't apply, etc.]
+</think>
+
+[Your actual response to the user]
+"""
+
+    def __init__(
+        self,
+        llm: LLM | None = None,
+        model: Model = OpenAI.GPT_4O_MINI,
+        thinking: bool = False,
+    ):
+        """
+        Initialize the Golden Response Generator.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM
+            thinking: Enable thinking mode with <think> tags in responses
+        """
+        self.llm = llm or LLM(model=model, temperature=0.7)
+        self.thinking = thinking
+
+    async def generate_single(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        scenario: GoldenScenario,
+        target_turns: int = 1,
+    ) -> Trace:
+        """
+        Generate a single trace with grounded reasoning.
+
+        Args:
+            policy_text: The policy document text
+            logic_map: The extracted Logic Map (DAG of rules)
+            scenario: The golden scenario to respond to
+            target_turns: Number of conversation turns
+
+        Returns:
+            Trace with messages and reasoning metadata
+        """
+        if target_turns > 1:
+            return await self._generate_multi_turn(
+                policy_text, logic_map, scenario, target_turns
+            )
+
+        return await self._generate_single_turn(policy_text, logic_map, scenario)
+
+    async def _generate_single_turn(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        scenario: GoldenScenario,
+    ) -> Trace:
+        """Generate a single-turn trace."""
+        # Format Logic Map for prompt
+        logic_map_str = self._format_logic_map(logic_map)
+
+        # Build prompt
+        prompt = GOLDEN_TRACE_PROMPT.format(
+            policy_text=policy_text,
+            logic_map=logic_map_str,
+            scenario_description=scenario.description,
+            scenario_context=scenario.context,
+            target_rule_ids=", ".join(scenario.target_rule_ids),
+            scenario_type=scenario.scenario_type.value.upper(),
+            expected_outcome=scenario.expected_outcome,
+        )
+
+        # Inject thinking instruction if enabled
+        if self.thinking:
+            prompt = prompt + self.THINKING_INSTRUCTION
+
+        # Generate structured output
+        result = await self.llm.generate_structured(prompt, GoldenTraceOutput)
+
+        # Convert to Trace
+        messages = [
+            Message(role=m.role, content=m.content)
+            for m in result.messages
+        ]
+
+        # Convert GoldenScenario to base Scenario for Trace
+        base_scenario = scenario.to_base_scenario()
+
+        # Convert reasoning chain to serializable format
+        reasoning_chain = None
+        if result.reasoning_chain:
+            reasoning_chain = [
+                {
+                    "rule_id": step.rule_id,
+                    "rule_text": step.rule_text,
+                    "applies": step.applies,
+                    "reasoning": step.reasoning,
+                    "exclusions": step.exclusions,
+                }
+                for step in result.reasoning_chain
+            ]
+
+        return Trace(
+            messages=messages,
+            scenario=base_scenario,
+            reasoning_chain=reasoning_chain,
+            rules_applied=result.rules_applied,
+            rules_excluded=result.rules_excluded,
+        )
+
+    async def _generate_multi_turn(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        scenario: GoldenScenario,
+        target_turns: int,
+    ) -> Trace:
+        """Generate a multi-turn trace."""
+        # Format Logic Map for prompt
+        logic_map_str = self._format_logic_map(logic_map)
+
+        # Build prompt
+        prompt = GOLDEN_TRACE_MULTI_TURN_PROMPT.format(
+            policy_text=policy_text,
+            logic_map=logic_map_str,
+            scenario_description=scenario.description,
+            scenario_context=scenario.context,
+            target_rule_ids=", ".join(scenario.target_rule_ids),
+            scenario_type=scenario.scenario_type.value.upper(),
+            target_turns=target_turns,
+        )
+
+        # Inject thinking instruction if enabled
+        if self.thinking:
+            prompt = prompt + self.THINKING_INSTRUCTION
+
+        # Generate structured output
+        result = await self.llm.generate_structured(prompt, GoldenTraceOutput)
+
+        # Convert to Trace
+        messages = [
+            Message(role=m.role, content=m.content)
+            for m in result.messages
+        ]
+
+        # Convert GoldenScenario to base Scenario for Trace
+        base_scenario = scenario.to_base_scenario()
+
+        # Convert reasoning chain to serializable format
+        reasoning_chain = None
+        if result.reasoning_chain:
+            reasoning_chain = [
+                {
+                    "rule_id": step.rule_id,
+                    "rule_text": step.rule_text,
+                    "applies": step.applies,
+                    "reasoning": step.reasoning,
+                    "exclusions": step.exclusions,
+                }
+                for step in result.reasoning_chain
+            ]
+
+        return Trace(
+            messages=messages,
+            scenario=base_scenario,
+            reasoning_chain=reasoning_chain,
+            rules_applied=result.rules_applied,
+            rules_excluded=result.rules_excluded,
+        )
+
+    def _format_logic_map(self, logic_map: LogicMap) -> str:
+        """Format Logic Map for prompt inclusion."""
+        lines = []
+        lines.append("RULES:")
+        for rule in logic_map.rules:
+            deps = f" [depends on: {', '.join(rule.dependencies)}]" if rule.dependencies else ""
+            lines.append(
+                f"  {rule.rule_id} ({rule.category.value}): {rule.text}{deps}"
+            )
+            lines.append(f"    IF: {rule.condition}")
+            lines.append(f"    THEN: {rule.action}")
+
+        lines.append("\nDEPENDENCY ORDER (evaluate in this order):")
+        # Show topological order for root rules and their chains
+        for root_id in logic_map.root_rules:
+            chain = logic_map.get_chain(root_id)
+            if chain:
+                chain_str = " -> ".join(r.rule_id for r in chain)
+                lines.append(f"  {chain_str}")
+
+        return "\n".join(lines)
+
+    async def generate(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        scenarios: list[GoldenScenario],
+        target_turns: int = 1,
+    ) -> list[Trace]:
+        """
+        Generate traces for multiple scenarios.
+
+        Args:
+            policy_text: The policy document text
+            logic_map: The extracted Logic Map
+            scenarios: List of golden scenarios
+            target_turns: Number of conversation turns
+
+        Returns:
+            List of traces with grounded reasoning
+        """
+        tasks = [
+            self.generate_single(policy_text, logic_map, s, target_turns)
+            for s in scenarios
+        ]
+        return await asyncio.gather(*tasks)
+
+
+__all__ = ["GoldenResponseGenerator"]