synkro 0.4.5__py3-none-any.whl

synkro/pipeline/phases.py

@@ -0,0 +1,237 @@
+"""Pipeline phases for generation.
+
+Each phase is a self-contained, testable unit that handles one step
+of the generation pipeline.
+"""
+
+import asyncio
+from asyncio import Semaphore
+from typing import TYPE_CHECKING
+
+from synkro.core.policy import Policy
+from synkro.types.core import Plan, Scenario, Trace
+from synkro.generation.planner import Planner
+from synkro.generation.scenarios import ScenarioGenerator
+from synkro.generation.responses import ResponseGenerator
+from synkro.quality.grader import Grader
+from synkro.quality.refiner import Refiner
+
+if TYPE_CHECKING:
+    from synkro.generation.tool_responses import ToolCallResponseGenerator
+
+
+class PlanPhase:
+    """
+    Planning phase - analyzes policy and creates category distribution.
+    
+    This phase uses a stronger model to understand the policy and
+    determine optimal scenario distribution.
+    """
+    
+    async def execute(self, policy: Policy, traces: int, planner: Planner) -> Plan:
+        """
+        Execute the planning phase.
+        
+        Args:
+            policy: The policy to analyze
+            traces: Target number of traces
+            planner: Planner component to use
+            
+        Returns:
+            Plan with categories and trace distribution
+        """
+        return await planner.plan(policy.text, traces)
+
+
+class ScenarioPhase:
+    """
+    Scenario generation phase - creates scenarios for each category.
+    
+    Runs in parallel across categories for efficiency.
+    """
+    
+    async def execute(
+        self,
+        policy: Policy,
+        plan: Plan,
+        generator: ScenarioGenerator,
+        semaphore: Semaphore,
+    ) -> list[Scenario]:
+        """
+        Execute scenario generation.
+        
+        Args:
+            policy: The policy text
+            plan: Plan with categories
+            generator: ScenarioGenerator component
+            semaphore: Semaphore for rate limiting
+            
+        Returns:
+            List of all generated scenarios
+        """
+        async def limited_generate(category):
+            async with semaphore:
+                return await generator.generate(policy.text, category.count, category=category)
+        
+        tasks = [limited_generate(cat) for cat in plan.categories]
+        results = await asyncio.gather(*tasks)
+        
+        # Flatten results
+        return [scenario for batch in results for scenario in batch]
+
+
+class ResponsePhase:
+    """
+    Response generation phase - creates responses for each scenario.
+    
+    Runs fully parallel with semaphore control.
+    """
+    
+    async def execute(
+        self,
+        policy: Policy,
+        scenarios: list[Scenario],
+        generator: ResponseGenerator,
+        semaphore: Semaphore,
+    ) -> list[Trace]:
+        """
+        Execute response generation.
+        
+        Args:
+            policy: The policy text
+            scenarios: List of scenarios to respond to
+            generator: ResponseGenerator component
+            semaphore: Semaphore for rate limiting
+            
+        Returns:
+            List of traces with generated responses
+        """
+        async def limited_generate(scenario):
+            async with semaphore:
+                return await generator._generate_single(policy.text, scenario)
+        
+        tasks = [limited_generate(s) for s in scenarios]
+        return await asyncio.gather(*tasks)
+
+
+class GradingPhase:
+    """
+    Grading and refinement phase - evaluates and improves responses.
+    
+    Includes the refinement loop for failed traces.
+    """
+    
+    async def execute(
+        self,
+        policy: Policy,
+        traces: list[Trace],
+        grader: Grader,
+        refiner: Refiner,
+        max_iterations: int,
+        semaphore: Semaphore,
+    ) -> tuple[list[Trace], float]:
+        """
+        Execute grading and refinement.
+        
+        Args:
+            policy: The policy text
+            traces: List of traces to grade
+            grader: Grader component
+            refiner: Refiner component
+            max_iterations: Maximum refinement iterations
+            semaphore: Semaphore for rate limiting
+            
+        Returns:
+            Tuple of (graded traces, pass rate percentage)
+        """
+        async def limited_grade(trace):
+            async with semaphore:
+                return await grader.grade(trace, policy.text)
+        
+        async def limited_refine(trace, grade):
+            async with semaphore:
+                return await refiner.refine(trace, grade, policy.text)
+        
+        # Initial grading
+        grade_tasks = [limited_grade(t) for t in traces]
+        grades = await asyncio.gather(*grade_tasks)
+        
+        # Attach grades
+        final_traces = list(traces)
+        for trace, grade in zip(final_traces, grades):
+            trace.grade = grade
+        
+        # Refinement loop
+        for iteration in range(1, max_iterations):
+            failed_indices = [i for i, t in enumerate(final_traces) if not t.grade.passed]
+            
+            if not failed_indices:
+                break
+            
+            # Refine failed traces
+            refine_tasks = [
+                limited_refine(final_traces[i], final_traces[i].grade)
+                for i in failed_indices
+            ]
+            refined_traces = await asyncio.gather(*refine_tasks)
+            
+            # Preserve original scenarios and update traces
+            for idx, refined in zip(failed_indices, refined_traces):
+                refined.scenario = final_traces[idx].scenario
+                final_traces[idx] = refined
+            
+            # Re-grade refined traces
+            regrade_tasks = [limited_grade(final_traces[i]) for i in failed_indices]
+            new_grades = await asyncio.gather(*regrade_tasks)
+            
+            for idx, grade in zip(failed_indices, new_grades):
+                final_traces[idx].grade = grade
+        
+        # Calculate pass rate
+        passed_count = sum(1 for t in final_traces if t.grade and t.grade.passed)
+        pass_rate = (passed_count / len(final_traces) * 100) if final_traces else 0
+        
+        return final_traces, pass_rate
+
+
+class ToolCallResponsePhase:
+    """
+    Tool call response generation phase - creates traces with proper tool calling format.
+    
+    Uses ToolCallResponseGenerator to produce traces with:
+    - System message with tool descriptions
+    - User message
+    - Assistant message with tool_calls (or direct response)
+    - Tool response messages  
+    - Final assistant message
+    """
+    
+    async def execute(
+        self,
+        policy: Policy,
+        scenarios: list[Scenario],
+        generator: "ToolCallResponseGenerator",
+        semaphore: Semaphore,
+    ) -> list[Trace]:
+        """
+        Execute tool call response generation.
+        
+        Args:
+            policy: The policy/guidelines text
+            scenarios: List of scenarios to respond to
+            generator: ToolCallResponseGenerator component
+            semaphore: Semaphore for rate limiting
+            
+        Returns:
+            List of traces with proper tool calling format
+        """
+        async def limited_generate(scenario):
+            async with semaphore:
+                return await generator.generate_single(policy.text, scenario)
+        
+        tasks = [limited_generate(s) for s in scenarios]
+        return await asyncio.gather(*tasks)
+
+
+__all__ = ["PlanPhase", "ScenarioPhase", "ResponsePhase", "GradingPhase", "ToolCallResponsePhase"]
+

synkro/pipeline/runner.py

@@ -0,0 +1,198 @@
+"""Pipeline runner that orchestrates all phases."""
+
+import asyncio
+from datetime import datetime
+
+from synkro.core.policy import Policy
+from synkro.core.dataset import Dataset
+from synkro.factory import ComponentFactory
+from synkro.reporting import ProgressReporter
+from synkro.pipeline.phases import (
+    PlanPhase,
+    ScenarioPhase,
+    ResponsePhase,
+    GradingPhase,
+    ToolCallResponsePhase,
+)
+
+
+class GenerationPipeline:
+    """
+    Orchestrates the full generation pipeline using decomposed phases.
+    
+    This class coordinates the execution of all phases and reports
+    progress through the injected reporter.
+    
+    Supports both standard SFT/QA generation and TOOL_CALL generation
+    with proper OpenAI function calling format.
+    
+    Examples:
+        >>> pipeline = GenerationPipeline(factory, reporter, workers=10)
+        >>> dataset = await pipeline.run(policy, traces=50)
+    """
+    
+    def __init__(
+        self,
+        factory: ComponentFactory,
+        reporter: ProgressReporter,
+        workers: int,
+        max_iterations: int = 1,
+        skip_grading: bool = False,
+    ):
+        """
+        Initialize the pipeline.
+        
+        Args:
+            factory: ComponentFactory for creating pipeline components
+            reporter: ProgressReporter for reporting progress
+            workers: Number of concurrent workers (API calls)
+            max_iterations: Maximum refinement iterations
+            skip_grading: Whether to skip the grading phase
+        """
+        self.factory = factory
+        self.reporter = reporter
+        self.workers = workers
+        self.max_iterations = max_iterations
+        self.skip_grading = skip_grading
+        
+        # Phases
+        self.plan_phase = PlanPhase()
+        self.scenario_phase = ScenarioPhase()
+        self.response_phase = ResponsePhase()
+        self.grading_phase = GradingPhase()
+        self.tool_call_response_phase = ToolCallResponsePhase()
+    
+    async def run(self, policy: Policy, traces: int, model: str, dataset_type: str) -> Dataset:
+        """
+        Run the full generation pipeline.
+        
+        Args:
+            policy: The policy to generate from
+            traces: Target number of traces
+            model: Model name (for reporting)
+            dataset_type: Dataset type (for reporting)
+            
+        Returns:
+            Dataset with generated traces
+        """
+        start_time = datetime.now()
+        semaphore = asyncio.Semaphore(self.workers)
+        
+        # Check if this is a tool_call dataset
+        is_tool_call = dataset_type == "tool_call"
+        
+        # Create components via factory
+        planner = self.factory.create_planner()
+        scenario_gen = self.factory.create_scenario_generator()
+        grader = self.factory.create_grader()
+        refiner = self.factory.create_refiner()
+        
+        # Create appropriate response generator
+        if is_tool_call and self.factory.has_tools:
+            tool_call_gen = self.factory.create_tool_call_response_generator()
+        else:
+            response_gen = self.factory.create_response_generator()
+        
+        # Report start
+        self.reporter.on_start(traces, model, dataset_type)
+        
+        # Phase 1: Planning
+        plan = await self.plan_phase.execute(policy, traces, planner)
+        self.reporter.on_plan_complete(plan)
+        
+        # Phase 2: Scenario generation
+        scenarios = await self.scenario_phase.execute(policy, plan, scenario_gen, semaphore)
+        self.reporter.on_scenarios_complete(scenarios)
+        
+        # Phase 3: Response generation (different for tool_call)
+        if is_tool_call and self.factory.has_tools:
+            all_traces = await self.tool_call_response_phase.execute(
+                policy, scenarios, tool_call_gen, semaphore
+            )
+        else:
+            all_traces = await self.response_phase.execute(
+                policy, scenarios, response_gen, semaphore
+            )
+        self.reporter.on_responses_complete(list(all_traces))
+        
+        # Phase 4: Grading (optional)
+        # Note: TOOL_CALL datasets now use specialized ToolCallGrader and
+        # ToolCallRefiner that preserve the tool_calls format.
+        pass_rate: float | None = None
+        
+        if self.skip_grading:
+            final_traces = list(all_traces)
+            self.reporter.on_grading_skipped()
+        else:
+            final_traces, pass_rate = await self._run_grading_with_reporting(
+                policy, list(all_traces), grader, refiner, semaphore
+            )
+            self.reporter.on_grading_complete(final_traces, pass_rate)
+        
+        # Report completion
+        elapsed = (datetime.now() - start_time).total_seconds()
+        self.reporter.on_complete(len(final_traces), elapsed, pass_rate)
+        
+        return Dataset(traces=final_traces)
+    
+    async def _run_grading_with_reporting(
+        self,
+        policy: Policy,
+        traces: list,
+        grader,
+        refiner,
+        semaphore,
+    ) -> tuple[list, float]:
+        """Run grading phase with refinement iteration reporting."""
+        
+        async def limited_grade(trace):
+            async with semaphore:
+                return await grader.grade(trace, policy.text)
+        
+        async def limited_refine(trace, grade):
+            async with semaphore:
+                return await refiner.refine(trace, grade, policy.text)
+        
+        # Initial grading
+        grade_tasks = [limited_grade(t) for t in traces]
+        grades = await asyncio.gather(*grade_tasks)
+        
+        final_traces = list(traces)
+        for trace, grade in zip(final_traces, grades):
+            trace.grade = grade
+        
+        # Refinement loop with reporting
+        for iteration in range(1, self.max_iterations):
+            failed_indices = [i for i, t in enumerate(final_traces) if not t.grade.passed]
+            
+            if not failed_indices:
+                break
+            
+            self.reporter.on_refinement_start(iteration + 1, len(failed_indices))
+            
+            # Refine
+            refine_tasks = [
+                limited_refine(final_traces[i], final_traces[i].grade)
+                for i in failed_indices
+            ]
+            refined_traces = await asyncio.gather(*refine_tasks)
+            
+            for idx, refined in zip(failed_indices, refined_traces):
+                refined.scenario = final_traces[idx].scenario
+                final_traces[idx] = refined
+            
+            # Re-grade
+            regrade_tasks = [limited_grade(final_traces[i]) for i in failed_indices]
+            new_grades = await asyncio.gather(*regrade_tasks)
+            
+            for idx, grade in zip(failed_indices, new_grades):
+                final_traces[idx].grade = grade
+        
+        passed_count = sum(1 for t in final_traces if t.grade and t.grade.passed)
+        pass_rate = (passed_count / len(final_traces) * 100) if final_traces else 0
+        
+        return final_traces, pass_rate
+
+
+__all__ = ["GenerationPipeline"]
+

synkro/pipelines.py

@@ -0,0 +1,105 @@
+"""Pipeline creation utilities.
+
+Usage:
+    from synkro.pipelines import create_pipeline
+    from synkro.models.openai import OpenAI
+    from synkro.types import DatasetType
+
+    pipeline = create_pipeline(
+        model=OpenAI.GPT_5_MINI,
+        dataset_type=DatasetType.SFT,
+    )
+    dataset = pipeline.generate("policy text", traces=50)
+    
+    # Tool calling pipeline
+    from synkro import ToolDefinition
+    
+    web_search = ToolDefinition(
+        name="web_search",
+        description="Search the web",
+        parameters={"type": "object", "properties": {"query": {"type": "string"}}}
+    )
+    
+    pipeline = create_pipeline(
+        dataset_type=DatasetType.TOOL_CALL,
+        tools=[web_search],
+    )
+    dataset = pipeline.generate("Search guidelines", traces=50)
+"""
+
+from typing import TYPE_CHECKING
+
+from synkro.generation.generator import Generator
+from synkro.types import DatasetType
+from synkro.models import Model, OpenAI
+from synkro.reporting import ProgressReporter
+
+if TYPE_CHECKING:
+    from synkro.types.tool import ToolDefinition
+
+
+def create_pipeline(
+    model: Model = OpenAI.GPT_5_MINI,
+    dataset_type: DatasetType = DatasetType.SFT,
+    grading_model: Model = OpenAI.GPT_52,
+    max_iterations: int = 3,
+    skip_grading: bool = False,
+    reporter: ProgressReporter | None = None,
+    tools: list["ToolDefinition"] | None = None,
+) -> Generator:
+    """
+    Create a pipeline for generating training datasets.
+
+    Args:
+        model: Model enum for generation (default: OpenAI.GPT_5_MINI)
+        dataset_type: Type of dataset - QA, SFT, or TOOL_CALL (default: SFT)
+        grading_model: Model enum for grading (default: OpenAI.GPT_52)
+        max_iterations: Max refinement iterations per trace (default: 3)
+        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
+
+    Returns:
+        Generator instance ready to use
+
+    Example:
+        >>> from synkro.pipelines import create_pipeline
+        >>> from synkro.models.openai import OpenAI
+        >>> from synkro.types import DatasetType
+        >>>
+        >>> pipeline = create_pipeline(
+        ...     model=OpenAI.GPT_5_MINI,
+        ...     dataset_type=DatasetType.SFT,
+        ... )
+        >>> dataset = pipeline.generate("policy text", traces=50)
+        >>> dataset.save("training.jsonl")
+        
+        >>> # Silent mode for embedding
+        >>> from synkro.reporting import SilentReporter
+        >>> pipeline = create_pipeline(reporter=SilentReporter())
+        
+        >>> # Tool calling dataset
+        >>> from synkro import ToolDefinition
+        >>> search_tool = ToolDefinition(
+        ...     name="web_search",
+        ...     description="Search the web for information",
+        ...     parameters={"type": "object", "properties": {"query": {"type": "string"}}}
+        ... )
+        >>> pipeline = create_pipeline(
+        ...     dataset_type=DatasetType.TOOL_CALL,
+        ...     tools=[search_tool],
+        ... )
+        >>> dataset = pipeline.generate("Search guidelines", traces=50)
+    """
+    return Generator(
+        dataset_type=dataset_type,
+        generation_model=model,
+        grading_model=grading_model,
+        max_iterations=max_iterations,
+        skip_grading=skip_grading,
+        reporter=reporter,
+        tools=tools,
+    )
+
+
+__all__ = ["create_pipeline"]

synkro/prompts/__init__.py

@@ -0,0 +1,44 @@
+"""Prompt templates and customizable prompt classes for Synkro."""
+
+from synkro.prompts.base import (
+    SystemPrompt,
+    ScenarioPrompt,
+    ResponsePrompt,
+    GradePrompt,
+    RefinePrompt,
+    PlanPrompt,
+)
+from synkro.prompts.templates import (
+    SYSTEM_PROMPT,
+    SCENARIO_GENERATOR_PROMPT,
+    CATEGORY_SCENARIO_PROMPT,
+    POLICY_PLANNING_PROMPT,
+    POLICY_COMPLEXITY_PROMPT,
+    BATCHED_RESPONSE_PROMPT,
+    BATCHED_GRADER_PROMPT,
+    BATCHED_REFINER_PROMPT,
+    SINGLE_RESPONSE_PROMPT,
+    SINGLE_GRADE_PROMPT,
+)
+
+__all__ = [
+    # Prompt classes
+    "SystemPrompt",
+    "ScenarioPrompt",
+    "ResponsePrompt",
+    "GradePrompt",
+    "RefinePrompt",
+    "PlanPrompt",
+    # Raw templates
+    "SYSTEM_PROMPT",
+    "SCENARIO_GENERATOR_PROMPT",
+    "CATEGORY_SCENARIO_PROMPT",
+    "POLICY_PLANNING_PROMPT",
+    "POLICY_COMPLEXITY_PROMPT",
+    "BATCHED_RESPONSE_PROMPT",
+    "BATCHED_GRADER_PROMPT",
+    "BATCHED_REFINER_PROMPT",
+    "SINGLE_RESPONSE_PROMPT",
+    "SINGLE_GRADE_PROMPT",
+]
+