synkro 0.4.12__py3-none-any.whl

synkro/generation/golden_scenarios.py

@@ -0,0 +1,276 @@
+"""Golden Scenario Generator - The Adversary.
+
+Generates typed scenarios (positive, negative, edge_case, irrelevant)
+with explicit rule targeting. This is Stage 2 of the Golden Trace pipeline.
+"""
+
+import asyncio
+from typing import Literal
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.schemas import GoldenScenariosArray
+from synkro.types.core import Category
+from synkro.types.logic_map import LogicMap, GoldenScenario, ScenarioType
+from synkro.prompts.golden_templates import (
+    GOLDEN_SCENARIO_PROMPT,
+    POSITIVE_SCENARIO_INSTRUCTIONS,
+    NEGATIVE_SCENARIO_INSTRUCTIONS,
+    EDGE_CASE_SCENARIO_INSTRUCTIONS,
+    IRRELEVANT_SCENARIO_INSTRUCTIONS,
+)
+
+
+# Default scenario type distribution
+DEFAULT_DISTRIBUTION = {
+    ScenarioType.POSITIVE: 0.35,    # 35% happy path
+    ScenarioType.NEGATIVE: 0.30,    # 30% violations
+    ScenarioType.EDGE_CASE: 0.25,   # 25% edge cases
+    ScenarioType.IRRELEVANT: 0.10,  # 10% out of scope
+}
+
+
+TYPE_INSTRUCTIONS = {
+    ScenarioType.POSITIVE: POSITIVE_SCENARIO_INSTRUCTIONS,
+    ScenarioType.NEGATIVE: NEGATIVE_SCENARIO_INSTRUCTIONS,
+    ScenarioType.EDGE_CASE: EDGE_CASE_SCENARIO_INSTRUCTIONS,
+    ScenarioType.IRRELEVANT: IRRELEVANT_SCENARIO_INSTRUCTIONS,
+}
+
+
+class GoldenScenarioGenerator:
+    """
+    The Adversary - Generates typed scenarios with rule targeting.
+
+    Produces scenarios across four types:
+    - POSITIVE (35%): Happy path, all criteria met
+    - NEGATIVE (30%): Violation, exactly one criterion fails
+    - EDGE_CASE (25%): Boundary conditions, exact limits
+    - IRRELEVANT (10%): Outside policy scope
+
+    Each scenario includes:
+    - Target rule IDs it's designed to test
+    - Expected outcome based on the rules
+    - Scenario type for classification
+
+    Examples:
+        >>> generator = GoldenScenarioGenerator(llm=LLM(model=OpenAI.GPT_4O_MINI))
+        >>> scenarios = await generator.generate(
+        ...     policy_text="...",
+        ...     logic_map=logic_map,
+        ...     category=category,
+        ...     count=10,
+        ... )
+    """
+
+    def __init__(
+        self,
+        llm: LLM | None = None,
+        model: Model = OpenAI.GPT_4O_MINI,
+        distribution: dict[ScenarioType, float] | None = None,
+    ):
+        """
+        Initialize the Golden Scenario Generator.
+
+        Args:
+            llm: LLM client to use (creates one if not provided)
+            model: Model to use if creating LLM
+            distribution: Custom scenario type distribution (defaults to 35/30/25/10)
+        """
+        self.llm = llm or LLM(model=model, temperature=0.8)
+        self.distribution = distribution or DEFAULT_DISTRIBUTION
+
+    async def generate(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        category: Category,
+        count: int,
+    ) -> list[GoldenScenario]:
+        """
+        Generate scenarios for a category with balanced type distribution.
+
+        Args:
+            policy_text: The policy document text
+            logic_map: The extracted Logic Map (DAG of rules)
+            category: The category to generate scenarios for
+            count: Total number of scenarios to generate
+
+        Returns:
+            List of GoldenScenarios with type distribution
+        """
+        # Calculate counts per type based on distribution
+        type_counts = self._calculate_type_distribution(count)
+
+        # Generate scenarios for each type in parallel
+        tasks = []
+        for scenario_type, type_count in type_counts.items():
+            if type_count > 0:
+                task = self._generate_type(
+                    policy_text=policy_text,
+                    logic_map=logic_map,
+                    category=category,
+                    scenario_type=scenario_type,
+                    count=type_count,
+                )
+                tasks.append(task)
+
+        # Gather all results
+        results = await asyncio.gather(*tasks)
+
+        # Flatten and return
+        scenarios = []
+        for batch in results:
+            scenarios.extend(batch)
+
+        return scenarios
+
+    def _calculate_type_distribution(self, total: int) -> dict[ScenarioType, int]:
+        """Calculate how many scenarios of each type to generate."""
+        counts = {}
+        remaining = total
+
+        # For small counts, prioritize non-IRRELEVANT types
+        # IRRELEVANT should only appear when we have enough scenarios
+        priority_order = [
+            ScenarioType.POSITIVE,
+            ScenarioType.NEGATIVE,
+            ScenarioType.EDGE_CASE,
+            ScenarioType.IRRELEVANT,  # Last priority
+        ]
+
+        if total <= 3:
+            # For very small counts, assign one to each priority type until exhausted
+            for stype in priority_order:
+                if remaining > 0:
+                    counts[stype] = 1
+                    remaining -= 1
+                else:
+                    counts[stype] = 0
+        else:
+            # Normal distribution for larger counts
+            for i, (stype, ratio) in enumerate(self.distribution.items()):
+                if i == len(self.distribution) - 1:
+                    # Last type gets remaining to ensure total is exact
+                    counts[stype] = remaining
+                else:
+                    count = round(total * ratio)
+                    counts[stype] = count
+                    remaining -= count
+
+        return counts
+
+    async def _generate_type(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        category: Category,
+        scenario_type: ScenarioType,
+        count: int,
+    ) -> list[GoldenScenario]:
+        """Generate scenarios of a specific type."""
+        # Get type-specific instructions
+        type_instructions = TYPE_INSTRUCTIONS[scenario_type]
+
+        # Format Logic Map for prompt
+        logic_map_str = self._format_logic_map(logic_map)
+
+        # Build prompt
+        prompt = GOLDEN_SCENARIO_PROMPT.format(
+            scenario_type=scenario_type.value.upper(),
+            policy_text=policy_text,
+            logic_map=logic_map_str,
+            category=category.name,
+            count=count,
+            type_specific_instructions=type_instructions,
+        )
+
+        # Generate structured output
+        result = await self.llm.generate_structured(prompt, GoldenScenariosArray)
+
+        # Convert to domain models
+        scenarios = []
+        for s in result.scenarios:
+            scenario = GoldenScenario(
+                description=s.description,
+                context=s.context,
+                category=category.name,
+                scenario_type=ScenarioType(s.scenario_type),
+                target_rule_ids=s.target_rule_ids,
+                expected_outcome=s.expected_outcome,
+            )
+            scenarios.append(scenario)
+
+        # Enforce requested count (LLM may return more or fewer)
+        return scenarios[:count]
+
+    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("\nROOT RULES (Entry Points):")
+        lines.append(f"  {', '.join(logic_map.root_rules)}")
+
+        return "\n".join(lines)
+
+    async def generate_for_categories(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        categories: list[Category],
+    ) -> tuple[list[GoldenScenario], dict[str, int]]:
+        """
+        Generate scenarios for multiple categories with distribution tracking.
+
+        Args:
+            policy_text: The policy document text
+            logic_map: The extracted Logic Map
+            categories: List of categories with counts
+
+        Returns:
+            Tuple of (all scenarios, type distribution counts)
+        """
+        # Generate for each category in parallel
+        tasks = [
+            self.generate(policy_text, logic_map, cat, cat.count)
+            for cat in categories
+        ]
+        results = await asyncio.gather(*tasks)
+
+        # Flatten scenarios
+        all_scenarios = []
+        for batch in results:
+            all_scenarios.extend(batch)
+
+        # Calculate distribution
+        distribution = {
+            ScenarioType.POSITIVE.value: 0,
+            ScenarioType.NEGATIVE.value: 0,
+            ScenarioType.EDGE_CASE.value: 0,
+            ScenarioType.IRRELEVANT.value: 0,
+        }
+        for s in all_scenarios:
+            distribution[s.scenario_type.value] += 1
+
+        return all_scenarios, distribution
+
+    def get_distribution_summary(self, scenarios: list[GoldenScenario]) -> dict[str, int]:
+        """Get a summary of scenario type distribution."""
+        distribution = {
+            "positive": 0,
+            "negative": 0,
+            "edge_case": 0,
+            "irrelevant": 0,
+        }
+        for s in scenarios:
+            distribution[s.scenario_type.value] += 1
+        return distribution
+
+
+__all__ = ["GoldenScenarioGenerator", "DEFAULT_DISTRIBUTION"]

synkro/generation/golden_tool_responses.py

@@ -0,0 +1,416 @@
+"""Golden Tool Response Generator - The Thinker for Tool Calls.
+
+Generates tool call traces with grounded reasoning and rule citations.
+This is Stage 3 of the Golden Trace pipeline for TOOL_CALL datasets.
+"""
+
+import json
+import uuid
+import asyncio
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.types.core import Trace, Message, Scenario
+from synkro.types.tool import ToolDefinition, ToolCall, ToolFunction
+from synkro.types.logic_map import LogicMap, GoldenScenario
+from synkro.prompts.golden_templates import GOLDEN_TOOL_TRACE_PROMPT
+
+if TYPE_CHECKING:
+    from synkro.generation.tool_simulator import ToolSimulator
+
+
+# =============================================================================
+# Pydantic models for structured JSON output
+# =============================================================================
+
+class GoldenToolCallRequest(BaseModel):
+    """A tool call request with rule citation."""
+
+    name: str = Field(description="Name of the tool to call")
+    arguments: str = Field(description="Arguments as JSON string")
+    rule_id: str = Field(description="Rule ID that requires this tool call")
+    reasoning: str = Field(description="Why this tool is needed for the rule")
+
+
+class GoldenToolDecision(BaseModel):
+    """Structured output for tool calling decision with rule grounding."""
+
+    needs_tool: bool = Field(description="Whether a tool call is needed")
+    reasoning: str = Field(description="Rule-based explanation of decision")
+    rule_ids_evaluated: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs that were evaluated"
+    )
+    tool_calls: list[GoldenToolCallRequest] = Field(
+        default_factory=list,
+        description="Tool calls with rule citations"
+    )
+    direct_response: str | None = Field(
+        default=None,
+        description="Direct response if no tool needed"
+    )
+
+
+class GoldenToolSynthesis(BaseModel):
+    """Structured output for synthesizing tool results."""
+
+    response: str = Field(description="Natural response incorporating tool results")
+    rules_applied: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs applied in the response"
+    )
+    rules_excluded: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs explicitly excluded"
+    )
+
+
+# =============================================================================
+# Golden Tool Call Response Generator
+# =============================================================================
+
+class GoldenToolCallResponseGenerator:
+    """
+    The Thinker for Tool Calls - Generates tool traces with grounded reasoning.
+
+    Produces tool call traces with:
+    - Rule citations for tool selection decisions
+    - Explicit reasoning linking rules to tool usage
+    - DAG-compliant evaluation order
+    - Verification-ready metadata
+
+    Examples:
+        >>> generator = GoldenToolCallResponseGenerator(
+        ...     tools=[web_search_tool],
+        ...     llm=LLM(model=OpenAI.GPT_4O_MINI),
+        ...     simulator=tool_simulator,
+        ... )
+        >>> trace = await generator.generate_single(
+        ...     policy_text="...",
+        ...     logic_map=logic_map,
+        ...     scenario=scenario,
+        ... )
+    """
+
+    def __init__(
+        self,
+        tools: list[ToolDefinition],
+        llm: LLM | None = None,
+        simulator: "ToolSimulator | None" = None,
+        model: Model = OpenAI.GPT_4O_MINI,
+    ):
+        """
+        Initialize the Golden Tool Call Response Generator.
+
+        Args:
+            tools: List of available tool definitions
+            llm: LLM client to use (creates one if not provided)
+            simulator: Tool simulator for generating tool responses
+            model: Model to use if creating LLM
+        """
+        self.tools = tools
+        self.tools_by_name = {t.name: t for t in tools}
+        self.llm = llm or LLM(model=model, temperature=0.7)
+        self.simulator = simulator
+
+    def _get_tools_description(self) -> str:
+        """Get formatted description of all tools."""
+        descriptions = []
+        for tool in self.tools:
+            descriptions.append(tool.to_system_prompt())
+        return "\n\n".join(descriptions)
+
+    def _generate_call_id(self) -> str:
+        """Generate a unique tool call ID."""
+        return f"call_{uuid.uuid4().hex[:12]}"
+
+    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}")
+        return "\n".join(lines)
+
+    async def generate_single(
+        self,
+        policy_text: str,
+        logic_map: LogicMap,
+        scenario: GoldenScenario,
+        target_turns: int = 1,
+    ) -> Trace:
+        """
+        Generate a single tool call 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 (currently single-turn only)
+
+        Returns:
+            Trace with proper tool calling format and rule citations
+        """
+        # TODO: Implement multi-turn tool calling support
+        tools_desc = self._get_tools_description()
+        logic_map_str = self._format_logic_map(logic_map)
+
+        # Step 1: Get LLM decision on tool usage with rule grounding
+        decision = await self._get_tool_decision(
+            policy_text, logic_map_str, scenario, tools_desc
+        )
+
+        # Step 2: Build the message sequence
+        messages = await self._build_message_sequence(
+            policy_text, logic_map_str, scenario, tools_desc, decision
+        )
+
+        # Convert GoldenScenario to base Scenario
+        base_scenario = scenario.to_base_scenario()
+
+        return Trace(messages=messages, scenario=base_scenario)
+
+    async def _get_tool_decision(
+        self,
+        policy_text: str,
+        logic_map_str: str,
+        scenario: GoldenScenario,
+        tools_desc: str,
+    ) -> GoldenToolDecision:
+        """Get the LLM's rule-grounded decision on tool usage."""
+        prompt = f"""You are a customer support agent deciding whether to use tools.
+Your decisions must be GROUNDED in the Logic Map rules.
+
+AVAILABLE TOOLS:
+{tools_desc}
+
+LOGIC MAP (Rules to Apply):
+{logic_map_str}
+
+POLICY GUIDELINES:
+{policy_text}
+
+SCENARIO:
+Type: {scenario.scenario_type.value.upper()}
+Request: {scenario.description}
+Context: {scenario.context}
+Target Rules: {', '.join(scenario.target_rule_ids)}
+
+YOUR TASK:
+1. Evaluate which rules from the Logic Map apply to this scenario
+2. Determine if any rule requires information that a tool can provide
+3. If tools are needed, specify which rule requires each tool call
+4. If no tools needed, explain based on which rules why direct response is sufficient
+
+TOOL CALLING RULES:
+- Only call a tool if a SPECIFIC RULE requires information the tool can provide
+- Cite the Rule ID that necessitates each tool call
+- If the scenario is IRRELEVANT type, no tools should be needed
+- If information is already in the context, don't call a tool for it"""
+
+        return await self.llm.generate_structured(prompt, GoldenToolDecision)
+
+    async def _build_message_sequence(
+        self,
+        policy_text: str,
+        logic_map_str: str,
+        scenario: GoldenScenario,
+        tools_desc: str,
+        decision: GoldenToolDecision,
+    ) -> list[Message]:
+        """Build the full message sequence based on the tool decision."""
+        messages = []
+
+        # System message with tool descriptions
+        system_content = f"""You are a helpful customer support agent. You have access to the following tools:
+
+{tools_desc}
+
+Follow the policy guidelines to assist customers effectively."""
+
+        messages.append(Message(role="system", content=system_content))
+
+        # User message
+        messages.append(Message(role="user", content=scenario.description))
+
+        if decision.needs_tool and decision.tool_calls:
+            # Assistant message with tool_calls
+            tool_calls = []
+            for tc in decision.tool_calls:
+                call_id = self._generate_call_id()
+                tool_calls.append(ToolCall(
+                    id=call_id,
+                    type="function",
+                    function=ToolFunction(
+                        name=tc.name,
+                        arguments=tc.arguments
+                    )
+                ))
+
+            messages.append(Message(
+                role="assistant",
+                content=None,
+                tool_calls=tool_calls
+            ))
+
+            # Tool response messages
+            tool_results = []
+            for tc in tool_calls:
+                result = await self._simulate_tool_call(tc)
+                tool_results.append(result)
+
+                messages.append(Message(
+                    role="tool",
+                    content=result,
+                    tool_call_id=tc.id
+                ))
+
+            # Final assistant message synthesizing results
+            final_response = await self._synthesize_response(
+                scenario, tool_calls, tool_results, decision, policy_text, logic_map_str
+            )
+            messages.append(Message(role="assistant", content=final_response))
+
+        else:
+            # Direct response without tools
+            response = decision.direct_response or await self._generate_direct_response(
+                policy_text, logic_map_str, scenario
+            )
+            messages.append(Message(role="assistant", content=response))
+
+        return messages
+
+    async def _simulate_tool_call(self, tool_call: ToolCall) -> str:
+        """Simulate a tool response."""
+        if self.simulator:
+            return await self.simulator.simulate(tool_call)
+
+        # Fallback: generate a mock response based on tool definition
+        tool_name = tool_call.function.name
+        if tool_name in self.tools_by_name:
+            tool = self.tools_by_name[tool_name]
+            if tool.mock_responses:
+                import random
+                return random.choice(tool.mock_responses)
+
+        # Default mock response
+        args = json.loads(tool_call.function.arguments)
+        return json.dumps({
+            "status": "success",
+            "result": f"Simulated response for {tool_name}",
+            "query": args
+        })
+
+    async def _synthesize_response(
+        self,
+        scenario: GoldenScenario,
+        tool_calls: list[ToolCall],
+        tool_results: list[str],
+        decision: GoldenToolDecision,
+        policy_text: str,
+        logic_map_str: str,
+    ) -> str:
+        """Synthesize a natural response from tool results with rule grounding."""
+        # Build context of tool calls and results
+        tools_context = []
+        for tc, result in zip(tool_calls, tool_results):
+            tools_context.append(f"Tool: {tc.function.name}")
+            tools_context.append(f"Arguments: {tc.function.arguments}")
+            tools_context.append(f"Result: {result}")
+            tools_context.append("")
+
+        prompt = f"""Based on the tool results and rules, provide a helpful response.
+
+USER REQUEST:
+{scenario.description}
+
+SCENARIO TYPE: {scenario.scenario_type.value.upper()}
+TARGET RULES: {', '.join(scenario.target_rule_ids)}
+
+TOOL RESULTS:
+{chr(10).join(tools_context)}
+
+LOGIC MAP:
+{logic_map_str}
+
+RULES EVALUATED: {', '.join(decision.rule_ids_evaluated)}
+
+Synthesize the tool results into a natural, helpful response.
+- Apply the relevant rules from the Logic Map
+- Incorporate the information from the tool results
+- Don't expose raw JSON or technical details
+- Be conversational and helpful"""
+
+        synthesis = await self.llm.generate_structured(prompt, GoldenToolSynthesis)
+        return synthesis.response
+
+    async def _generate_direct_response(
+        self,
+        policy_text: str,
+        logic_map_str: str,
+        scenario: GoldenScenario,
+    ) -> str:
+        """Generate a direct response when no tools are needed."""
+        prompt = f"""Provide a helpful response based on the rules.
+
+USER REQUEST:
+{scenario.description}
+
+CONTEXT:
+{scenario.context}
+
+SCENARIO TYPE: {scenario.scenario_type.value.upper()}
+TARGET RULES: {', '.join(scenario.target_rule_ids)}
+
+LOGIC MAP:
+{logic_map_str}
+
+POLICY GUIDELINES:
+{policy_text}
+
+No tools are needed for this request. Provide a direct, helpful response
+applying the relevant rules from the Logic Map."""
+
+        synthesis = await self.llm.generate_structured(prompt, GoldenToolSynthesis)
+        return synthesis.response
+
+    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 tool calling format
+        """
+        tasks = [
+            self.generate_single(policy_text, logic_map, s, target_turns)
+            for s in scenarios
+        ]
+        return await asyncio.gather(*tasks)
+
+
+__all__ = [
+    "GoldenToolCallResponseGenerator",
+    "GoldenToolDecision",
+    "GoldenToolCallRequest",
+    "GoldenToolSynthesis",
+]