synkro 0.4.12__py3-none-any.whl

synkro/generation/logic_extractor.py

@@ -0,0 +1,126 @@
+"""Logic Extractor - The Cartographer.
+
+Extracts a Logic Map (DAG of rules) from a policy document.
+This is Stage 1 of the Golden Trace pipeline.
+"""
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.schemas import LogicMapOutput
+from synkro.types.logic_map import LogicMap, Rule, RuleCategory
+from synkro.prompts.golden_templates import LOGIC_EXTRACTION_PROMPT
+
+
+class LogicExtractor:
+    """
+    The Cartographer - Extracts a Logic Map from policy documents.
+
+    The Logic Map is a Directed Acyclic Graph (DAG) of rules that enables:
+    - Grounded scenario generation with rule references
+    - Chain-of-thought reasoning with rule citations
+    - Verification that traces don't skip or hallucinate rules
+
+    Examples:
+        >>> extractor = LogicExtractor(llm=LLM(model=OpenAI.GPT_4O))
+        >>> logic_map = await extractor.extract(policy_text)
+        >>> print(f"Extracted {len(logic_map.rules)} rules")
+    """
+
+    def __init__(
+        self,
+        llm: LLM | None = None,
+        model: Model = OpenAI.GPT_4O,
+    ):
+        """
+        Initialize the Logic Extractor.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM (default: GPT-4O for accuracy)
+        """
+        self.llm = llm or LLM(model=model, temperature=0.3)
+
+    async def extract(self, policy_text: str) -> LogicMap:
+        """
+        Extract a Logic Map from a policy document.
+
+        Args:
+            policy_text: The policy document text
+
+        Returns:
+            LogicMap with extracted rules as a DAG
+
+        Raises:
+            ValueError: If extraction fails or produces invalid DAG
+        """
+        # Format the prompt
+        prompt = LOGIC_EXTRACTION_PROMPT.format(policy_text=policy_text)
+
+        # Generate structured output
+        result = await self.llm.generate_structured(prompt, LogicMapOutput)
+
+        # Convert to domain model
+        logic_map = self._convert_to_logic_map(result)
+
+        # Validate DAG properties
+        if not logic_map.validate_dag():
+            raise ValueError("Extracted rules contain circular dependencies")
+
+        return logic_map
+
+    def _convert_to_logic_map(self, output: LogicMapOutput) -> LogicMap:
+        """Convert schema output to domain model."""
+        rules = []
+        for rule_out in output.rules:
+            # Convert category string to enum
+            category = RuleCategory(rule_out.category)
+
+            rule = Rule(
+                rule_id=rule_out.rule_id,
+                text=rule_out.text,
+                condition=rule_out.condition,
+                action=rule_out.action,
+                dependencies=rule_out.dependencies,
+                category=category,
+            )
+            rules.append(rule)
+
+        return LogicMap(
+            rules=rules,
+            root_rules=output.root_rules,
+        )
+
+    async def extract_with_retry(
+        self,
+        policy_text: str,
+        max_retries: int = 2,
+    ) -> LogicMap:
+        """
+        Extract with retry on validation failure.
+
+        Args:
+            policy_text: The policy document text
+            max_retries: Maximum retry attempts
+
+        Returns:
+            LogicMap with extracted rules
+
+        Raises:
+            ValueError: If extraction fails after all retries
+        """
+        last_error = None
+        for attempt in range(max_retries + 1):
+            try:
+                return await self.extract(policy_text)
+            except ValueError as e:
+                last_error = e
+                if attempt < max_retries:
+                    # Add hint for next attempt
+                    continue
+
+        raise ValueError(
+            f"Failed to extract valid Logic Map after {max_retries + 1} attempts: {last_error}"
+        )
+
+
+__all__ = ["LogicExtractor"]

synkro/generation/multiturn_responses.py

@@ -0,0 +1,177 @@
+"""Multi-turn response generation for complex conversations."""
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.types.core import Scenario, Trace, Message
+from synkro.prompts.multiturn_templates import (
+    MULTI_TURN_INITIAL_PROMPT,
+    MULTI_TURN_RESPONSE_PROMPT,
+)
+from synkro.prompts.templates import SYSTEM_PROMPT
+from synkro.generation.follow_ups import FollowUpGenerator
+
+
+class MultiTurnResponseGenerator:
+    """
+    Generates multi-turn conversations based on policy complexity.
+
+    Turn allocation:
+    - Simple (1-2 turns): Single query -> Straight answer
+    - Conditional (3 turns): Query -> Clarification -> Verdict
+    - Complex (5+ turns): Multiple rounds of exploration
+
+    Examples:
+        >>> gen = MultiTurnResponseGenerator()
+        >>> trace = await gen.generate_single(policy_text, scenario, target_turns=3)
+        >>> print(len([m for m in trace.messages if m.role == "assistant"]))
+        3
+    """
+
+    def __init__(self, llm: LLM | None = None, model: Model = OpenAI.GPT_4O_MINI):
+        """
+        Initialize the multi-turn response generator.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM
+        """
+        self.llm = llm or LLM(model=model)
+        self.follow_up_gen = FollowUpGenerator(llm=self.llm)
+
+    def _format_conversation(self, messages: list[Message]) -> str:
+        """Format conversation messages for prompt inclusion."""
+        formatted = []
+        for msg in messages:
+            role = msg.role.upper()
+            content = msg.content or "[No content]"
+            formatted.append(f"{role}: {content}")
+        return "\n\n".join(formatted)
+
+    async def _generate_initial_response(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        target_turns: int,
+    ) -> list[Message]:
+        """
+        Generate the initial exchange (system + user + assistant).
+
+        Args:
+            policy_text: The policy text
+            scenario: The scenario to respond to
+            target_turns: Total target turns for context
+
+        Returns:
+            List of initial messages [system, user, assistant]
+        """
+        prompt = MULTI_TURN_INITIAL_PROMPT.format(
+            target_turns=target_turns,
+            scenario=scenario.description,
+            context=scenario.context,
+            policy=policy_text,
+        )
+
+        response = await self.llm.generate(prompt)
+
+        return [
+            Message(role="system", content=SYSTEM_PROMPT),
+            Message(role="user", content=f"{scenario.description}\n\nContext: {scenario.context}"),
+            Message(role="assistant", content=response.strip()),
+        ]
+
+    async def _generate_response(
+        self,
+        policy_text: str,
+        messages: list[Message],
+        question: str,
+    ) -> str:
+        """
+        Generate a response to a follow-up question.
+
+        Args:
+            policy_text: The policy text
+            messages: Conversation history
+            question: The follow-up question to answer
+
+        Returns:
+            Assistant response text
+        """
+        conversation = self._format_conversation(messages)
+
+        prompt = MULTI_TURN_RESPONSE_PROMPT.format(
+            conversation=conversation,
+            question=question,
+            policy=policy_text,
+        )
+
+        response = await self.llm.generate(prompt)
+        return response.strip()
+
+    async def generate_single(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        target_turns: int,
+    ) -> Trace:
+        """
+        Generate a multi-turn trace for one scenario.
+
+        Args:
+            policy_text: The policy text
+            scenario: The scenario to generate for
+            target_turns: Number of user-assistant exchanges
+
+        Returns:
+            Trace with multi-turn conversation
+        """
+        # Generate initial exchange
+        messages = await self._generate_initial_response(
+            policy_text, scenario, target_turns
+        )
+
+        # Generate follow-up turns
+        for turn in range(1, target_turns):
+            # Generate follow-up question
+            follow_up = await self.follow_up_gen.generate(
+                policy_text=policy_text,
+                messages=messages,
+                turn_index=turn,
+            )
+
+            # Add user message with follow-up question
+            messages.append(Message(role="user", content=follow_up.question))
+
+            # Generate assistant response
+            response = await self._generate_response(
+                policy_text=policy_text,
+                messages=messages,
+                question=follow_up.question,
+            )
+
+            # Add assistant response
+            messages.append(Message(role="assistant", content=response))
+
+        return Trace(messages=messages, scenario=scenario)
+
+    async def generate(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+        target_turns: int,
+    ) -> list[Trace]:
+        """
+        Generate multi-turn traces for multiple scenarios.
+
+        Args:
+            policy_text: The policy text
+            scenarios: List of scenarios to generate for
+            target_turns: Number of turns per trace
+
+        Returns:
+            List of traces with multi-turn conversations
+        """
+        traces = []
+        for scenario in scenarios:
+            trace = await self.generate_single(policy_text, scenario, target_turns)
+            traces.append(trace)
+        return traces

synkro/generation/planner.py

@@ -0,0 +1,131 @@
+"""Planning for trace generation across categories."""
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.types.core import Plan, Category
+from synkro.prompts.templates import POLICY_PLANNING_PROMPT, POLICY_COMPLEXITY_PROMPT
+from synkro.schemas import PolicyPlan, PolicyComplexity
+
+
+class Planner:
+    """
+    Plans how to distribute trace generation across categories.
+
+    The planner analyzes the policy and creates an optimal distribution
+    of scenarios across different categories to ensure comprehensive
+    coverage.
+
+    Examples:
+        >>> planner = Planner()
+        >>> plan = await planner.plan(policy, target_traces=100)
+        >>> for cat in plan.categories:
+        ...     print(f"{cat.name}: {cat.traces} traces")
+    """
+
+    def __init__(self, llm: LLM | None = None, model: Model = OpenAI.GPT_4O):
+        """
+        Initialize the planner.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM
+        """
+        self.llm = llm or LLM(model=model)
+
+    async def analyze_complexity(self, policy_text: str) -> PolicyComplexity:
+        """
+        Analyze policy complexity to determine optimal conversation turns.
+
+        Uses the PolicyComplexity schema to assess:
+        - Variable count (rules, conditions, exceptions)
+        - Complexity level (simple, conditional, complex)
+        - Recommended turns (1-6)
+
+        Args:
+            policy_text: The policy text to analyze
+
+        Returns:
+            PolicyComplexity with recommended turns and complexity level
+        """
+        prompt = f"""{POLICY_COMPLEXITY_PROMPT}
+
+POLICY:
+{policy_text}
+
+Analyze the policy complexity and recommend conversation turns."""
+
+        try:
+            return await self.llm.generate_structured(prompt, PolicyComplexity)
+        except Exception:
+            # Default to simple single-turn
+            return PolicyComplexity(
+                variable_count=1,
+                complexity_level="simple",
+                recommended_turns=1,
+                reasoning="Default - unable to analyze complexity",
+            )
+
+    async def plan(self, policy_text: str, target_traces: int, analyze_turns: bool = True) -> Plan:
+        """
+        Create a generation plan for the policy.
+
+        Analyzes the policy and determines optimal category distribution
+        and conversation turn count.
+
+        Args:
+            policy_text: The policy text to analyze
+            target_traces: Target number of traces to generate
+            analyze_turns: Whether to analyze policy for turn recommendation
+
+        Returns:
+            Plan object with categories, reasoning, and turn recommendations
+        """
+        prompt = f"""{POLICY_PLANNING_PROMPT}
+
+POLICY:
+{policy_text}
+
+TARGET TRACES: {target_traces}
+
+Analyze the policy and create a plan with categories for generating training data."""
+
+        # Analyze complexity for turn recommendations
+        complexity = None
+        if analyze_turns:
+            complexity = await self.analyze_complexity(policy_text)
+
+        try:
+            # Use structured output for reliable planning
+            parsed = await self.llm.generate_structured(prompt, PolicyPlan)
+
+            # Convert to typed objects
+            categories = [
+                Category(
+                    name=c.name,
+                    description=c.description,
+                    count=c.traces,
+                )
+                for c in parsed.categories
+            ]
+
+            return Plan(
+                categories=categories,
+                reasoning=parsed.reasoning,
+                recommended_turns=complexity.recommended_turns if complexity else 1,
+                complexity_level=complexity.complexity_level if complexity else "simple",
+            )
+        except Exception:
+            # Fallback plan
+            third = target_traces // 3
+            remainder = target_traces - (third * 3)
+            return Plan(
+                categories=[
+                    Category(name="Happy Path", description="Clear success cases", count=third),
+                    Category(name="Edge Cases", description="Ambiguous situations", count=third),
+                    Category(name="Violations", description="Clear failure cases", count=third + remainder),
+                ],
+                reasoning="Default plan - unable to parse LLM response",
+                recommended_turns=complexity.recommended_turns if complexity else 1,
+                complexity_level=complexity.complexity_level if complexity else "simple",
+            )
+

synkro/generation/responses.py

@@ -0,0 +1,189 @@
+"""Response generation for scenarios."""
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.types.core import Scenario, Trace, Message
+from synkro.prompts.templates import BATCHED_RESPONSE_PROMPT, SYSTEM_PROMPT
+from synkro.schemas import SingleResponse
+from synkro.parsers import parse_batched_responses, extract_content
+from synkro.generation.multiturn_responses import MultiTurnResponseGenerator
+
+
+class ResponseGenerator:
+    """
+    Generates expert responses for scenarios.
+
+    Creates comprehensive, policy-grounded responses that demonstrate
+    deep domain understanding.
+
+    Examples:
+        >>> gen = ResponseGenerator()
+        >>> traces = await gen.generate(policy.text, scenarios)
+    """
+
+    def __init__(self, llm: LLM | None = None, model: Model = OpenAI.GPT_4O_MINI):
+        """
+        Initialize the response generator.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM
+        """
+        self.llm = llm or LLM(model=model)
+        self._multi_turn_gen: MultiTurnResponseGenerator | None = None
+
+    @property
+    def multi_turn_gen(self) -> MultiTurnResponseGenerator:
+        """Lazy initialization of multi-turn generator."""
+        if self._multi_turn_gen is None:
+            self._multi_turn_gen = MultiTurnResponseGenerator(llm=self.llm)
+        return self._multi_turn_gen
+
+    async def generate(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+        target_turns: int = 1,
+    ) -> list[Trace]:
+        """
+        Generate responses for scenarios.
+
+        Args:
+            policy_text: The policy text
+            scenarios: List of scenarios to respond to
+            target_turns: Number of conversation turns (1 for single-turn)
+
+        Returns:
+            List of traces with generated responses
+        """
+        traces = []
+
+        # Generate responses one at a time for better quality
+        for scenario in scenarios:
+            trace = await self._generate_single(policy_text, scenario, target_turns)
+            traces.append(trace)
+
+        return traces
+
+    async def _generate_single(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        target_turns: int = 1,
+    ) -> Trace:
+        """
+        Generate a single trace for one scenario.
+
+        Args:
+            policy_text: The policy text
+            scenario: The scenario to respond to
+            target_turns: Number of conversation turns (1 for single-turn)
+
+        Returns:
+            Trace with generated messages
+        """
+        # Delegate to multi-turn generator for multi-turn traces
+        if target_turns > 1:
+            return await self.multi_turn_gen.generate_single(
+                policy_text, scenario, target_turns
+            )
+
+        # Single-turn generation
+        prompt = f"""You are a domain expert generating a training example.
+
+Given the scenario and policy below, create a complete training example.
+
+The assistant response must:
+- Start with <reasoning> tags showing your thought process
+- Cite specific policy sections that apply
+- Give specific, actionable recommendations
+- Address all aspects of the scenario
+- Acknowledge edge cases and complications
+
+SCENARIO:
+{scenario.description}
+
+CONTEXT:
+{scenario.context}
+
+POLICY:
+{policy_text}
+
+Generate exactly 3 messages: system, user, and assistant."""
+
+        # Use structured output for reliable JSON
+        parsed = await self.llm.generate_structured(prompt, SingleResponse)
+        messages = [
+            Message(role=m.role, content=m.content) for m in parsed.messages
+        ]
+
+        return Trace(messages=messages, scenario=scenario)
+
+    async def generate_batch(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+        batch_size: int = 10,
+    ) -> list[Trace]:
+        """
+        Generate responses in batches.
+
+        More efficient than single generation for large numbers of scenarios.
+
+        Args:
+            policy_text: The policy text
+            scenarios: List of scenarios to respond to
+            batch_size: Number of scenarios per batch
+
+        Returns:
+            List of traces with generated responses
+        """
+        traces = []
+
+        for i in range(0, len(scenarios), batch_size):
+            batch = scenarios[i : i + batch_size]
+            batch_traces = await self._generate_batch(policy_text, batch)
+            traces.extend(batch_traces)
+
+        return traces
+
+    async def _generate_batch(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+    ) -> list[Trace]:
+        """Generate traces for a batch of scenarios."""
+        scenarios_text = "\n\n".join(
+            f"SCENARIO {i}:\n{s.description}\n\nCONTEXT:\n{s.context}"
+            for i, s in enumerate(scenarios)
+        )
+
+        prompt = f"""{BATCHED_RESPONSE_PROMPT}
+
+SYSTEM PROMPT TO USE:
+{SYSTEM_PROMPT}
+
+POLICY:
+{policy_text}
+
+SCENARIOS:
+{scenarios_text}"""
+
+        response = await self.llm.generate(prompt)
+        from synkro.schemas import ScenarioOutput
+
+        scenario_outputs = [
+            ScenarioOutput(scenario=s.description, context=s.context) for s in scenarios
+        ]
+        parsed = parse_batched_responses(response, len(scenarios), scenario_outputs)
+
+        traces = []
+        for i, p in enumerate(parsed):
+            scenario = scenarios[min(p["index"], len(scenarios) - 1)]
+            messages = [
+                Message(role=m.role, content=m.content) for m in p["messages"]
+            ]
+            traces.append(Trace(messages=messages, scenario=scenario))
+
+        return traces
+