synkro 0.4.5__py3-none-any.whl

synkro/generation/planner.py

@@ -0,0 +1,87 @@
+"""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
+from synkro.schemas import PolicyPlan
+
+
+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 plan(self, policy_text: str, target_traces: int) -> Plan:
+        """
+        Create a generation plan for the policy.
+
+        Analyzes the policy and determines optimal category distribution.
+
+        Args:
+            policy_text: The policy text to analyze
+            target_traces: Target number of traces to generate
+
+        Returns:
+            Plan object with categories and reasoning
+        """
+        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."""
+
+        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,
+            )
+        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",
+            )
+

synkro/generation/responses.py

@@ -0,0 +1,160 @@
+"""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
+
+
+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)
+
+    async def generate(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+    ) -> list[Trace]:
+        """
+        Generate responses for scenarios.
+
+        Args:
+            policy_text: The policy text
+            scenarios: List of scenarios to respond to
+
+        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)
+            traces.append(trace)
+
+        return traces
+
+    async def _generate_single(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+    ) -> Trace:
+        """Generate a single trace for one scenario."""
+        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
+

synkro/generation/scenarios.py

@@ -0,0 +1,90 @@
+"""Scenario generation from policy documents."""
+
+from synkro.llm.client import LLM
+from synkro.models import Model, OpenAI
+from synkro.types.core import Scenario, Category
+from synkro.prompts.templates import SCENARIO_GENERATOR_PROMPT, CATEGORY_SCENARIO_PROMPT
+from synkro.schemas import ScenariosArray
+
+
+class ScenarioGenerator:
+    """
+    Generates realistic scenarios from policy documents.
+
+    Creates diverse scenarios that test different aspects of policy
+    understanding and compliance.
+
+    Examples:
+        >>> gen = ScenarioGenerator()
+        >>> scenarios = await gen.generate(policy.text, count=50)
+        >>> for s in scenarios:
+        ...     print(s.description)
+    """
+
+    def __init__(self, llm: LLM | None = None, model: Model = OpenAI.GPT_4O_MINI):
+        """
+        Initialize the scenario 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.prompt_template = SCENARIO_GENERATOR_PROMPT
+
+    async def generate(
+        self,
+        policy_text: str,
+        count: int,
+        category: Category | None = None,
+    ) -> list[Scenario]:
+        """
+        Generate scenarios from the policy.
+
+        Args:
+            policy_text: The policy text
+            count: Number of scenarios to generate
+            category: Optional category to focus on
+
+        Returns:
+            List of generated scenarios
+        """
+        if category:
+            prompt = self._build_category_prompt(policy_text, count, category)
+        else:
+            prompt = self._build_general_prompt(policy_text, count)
+
+        # Use structured output for reliable scenario generation
+        parsed = await self.llm.generate_structured(prompt, ScenariosArray)
+        return [
+            Scenario(
+                description=s.scenario,
+                context=s.context,
+                category=category.name if category else None,
+            )
+            for s in parsed.scenarios[:count]
+        ]
+
+    def _build_general_prompt(self, policy_text: str, count: int) -> str:
+        """Build prompt for general scenario generation."""
+        return f"""{self.prompt_template}
+
+POLICY:
+{policy_text}
+
+Generate exactly {count} diverse scenarios."""
+
+    def _build_category_prompt(
+        self, policy_text: str, count: int, category: Category
+    ) -> str:
+        """Build prompt for category-specific scenario generation."""
+        return f"""{CATEGORY_SCENARIO_PROMPT}
+
+Category: {category.name}
+Description: {category.description}
+
+POLICY:
+{policy_text}
+
+Generate exactly {count} scenarios for the "{category.name}" category."""
+

synkro/generation/tool_responses.py

@@ -0,0 +1,370 @@
+"""Tool call response generation with JSON mode for structured outputs."""
+
+import json
+import uuid
+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 Scenario, Trace, Message
+from synkro.types.tool import ToolCall, ToolFunction, ToolDefinition
+
+if TYPE_CHECKING:
+    from synkro.generation.tool_simulator import ToolSimulator
+
+
+# =============================================================================
+# Pydantic models for structured JSON output
+# =============================================================================
+
+class ToolCallRequest(BaseModel):
+    """A single tool call request from the LLM."""
+    
+    name: str = Field(description="Name of the tool to call")
+    arguments: str = Field(description="Arguments as a JSON string, e.g. '{\"query\": \"test\"}'")
+    
+    def get_arguments_dict(self) -> dict:
+        """Parse arguments JSON string to dict."""
+        return json.loads(self.arguments)
+
+
+class ToolCallDecision(BaseModel):
+    """
+    Structured output for the LLM's tool calling decision.
+    
+    The LLM outputs this to indicate whether tools are needed
+    and which ones to call.
+    """
+    
+    needs_tool: bool = Field(
+        description="Whether a tool call is needed to answer the user's request"
+    )
+    reasoning: str = Field(
+        description="Brief explanation of why tool is/isn't needed"
+    )
+    tool_calls: list[ToolCallRequest] = Field(
+        default_factory=list,
+        description="List of tool calls to make (empty if needs_tool is False)"
+    )
+    direct_response: str | None = Field(
+        default=None,
+        description="Direct response if no tool is needed"
+    )
+
+
+class FinalSynthesis(BaseModel):
+    """Structured output for synthesizing tool results into a response."""
+    
+    response: str = Field(
+        description="Natural response incorporating the tool results"
+    )
+
+
+# =============================================================================
+# Tool Call Response Generator
+# =============================================================================
+
+class ToolCallResponseGenerator:
+    """
+    Generates tool call training traces using JSON mode for structured outputs.
+    
+    Produces traces in OpenAI function calling format:
+    - system message with tool descriptions
+    - user message with request
+    - assistant message with tool_calls (or direct response)
+    - tool response messages
+    - final assistant message synthesizing results
+    
+    Example:
+        >>> gen = ToolCallResponseGenerator(
+        ...     tools=[web_search_tool, db_tool],
+        ...     llm=LLM(model=OpenAI.GPT_4O),
+        ...     simulator=tool_simulator,
+        ... )
+        >>> trace = await gen.generate_single(policy_text, scenario)
+    """
+    
+    def __init__(
+        self,
+        tools: list[ToolDefinition],
+        llm: LLM | None = None,
+        simulator: "ToolSimulator | None" = None,
+        model: Model = OpenAI.GPT_4O_MINI,
+    ):
+        """
+        Initialize the 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)
+        self.simulator = simulator
+    
+    def _get_tools_description(self) -> str:
+        """Get formatted description of all tools for system prompt."""
+        descriptions = []
+        for tool in self.tools:
+            descriptions.append(tool.to_system_prompt())
+        return "\n\n".join(descriptions)
+    
+    def _get_tools_json_schema(self) -> str:
+        """Get JSON schema representation of tools."""
+        tools_json = [tool.to_openai_format() for tool in self.tools]
+        return json.dumps(tools_json, indent=2)
+    
+    def _generate_call_id(self) -> str:
+        """Generate a unique tool call ID."""
+        return f"call_{uuid.uuid4().hex[:12]}"
+    
+    async def generate_single(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+    ) -> Trace:
+        """
+        Generate a single tool call trace.
+        
+        Args:
+            policy_text: The policy/guidelines text
+            scenario: The scenario to respond to
+            
+        Returns:
+            Trace with proper tool calling format
+        """
+        tools_desc = self._get_tools_description()
+        
+        # Step 1: Get LLM decision on tool usage
+        decision = await self._get_tool_decision(policy_text, scenario, tools_desc)
+        
+        # Step 2: Build the message sequence
+        messages = await self._build_message_sequence(
+            policy_text, scenario, tools_desc, decision
+        )
+        
+        return Trace(messages=messages, scenario=scenario)
+    
+    async def _get_tool_decision(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        tools_desc: str,
+    ) -> ToolCallDecision:
+        """
+        Get the LLM's decision on whether to use tools.
+        
+        Uses JSON mode to force structured output.
+        """
+        prompt = f"""You are a customer support agent deciding whether to use tools.
+
+AVAILABLE TOOLS:
+{tools_desc}
+
+TOOL USAGE GUIDELINES:
+{policy_text}
+
+USER REQUEST:
+{scenario.description}
+
+CONTEXT:
+{scenario.context}
+
+Analyze this request and decide:
+1. Does this require calling a tool, or can you answer directly?
+2. If tools are needed, which ones and with what arguments?
+3. If no tools needed, provide the direct response.
+
+Important rules:
+- Only call tools when necessary (don't call for information you already know)
+- Use correct tool names and parameter types
+- If multiple tools are needed, list them all
+- Provide clear reasoning for your decision"""
+
+        return await self.llm.generate_structured(prompt, ToolCallDecision)
+    
+    async def _build_message_sequence(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        tools_desc: str,
+        decision: ToolCallDecision,
+    ) -> 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 tool usage guidelines provided 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  # Already a JSON string
+                    )
+                ))
+            
+            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.description, tool_calls, tool_results, policy_text
+            )
+            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, scenario, tools_desc
+            )
+            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:
+                # Use a mock response
+                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,
+        user_request: str,
+        tool_calls: list[ToolCall],
+        tool_results: list[str],
+        policy_text: str,
+    ) -> str:
+        """Synthesize a natural response from tool results."""
+        # 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, provide a helpful response to the user.
+
+USER REQUEST:
+{user_request}
+
+TOOL RESULTS:
+{chr(10).join(tools_context)}
+
+GUIDELINES:
+{policy_text}
+
+Synthesize the tool results into a natural, helpful response. 
+- Incorporate the information from the tool results
+- Don't expose raw JSON or technical details
+- Be conversational and helpful
+- If a tool returned an error, acknowledge it and offer alternatives"""
+
+        synthesis = await self.llm.generate_structured(prompt, FinalSynthesis)
+        return synthesis.response
+    
+    async def _generate_direct_response(
+        self,
+        policy_text: str,
+        scenario: Scenario,
+        tools_desc: str,
+    ) -> str:
+        """Generate a direct response when no tools are needed."""
+        prompt = f"""Provide a helpful response to the user's request.
+
+USER REQUEST:
+{scenario.description}
+
+CONTEXT:
+{scenario.context}
+
+GUIDELINES:
+{policy_text}
+
+Note: No tools are needed for this request. Provide a direct, helpful response
+based on your knowledge and the guidelines."""
+
+        synthesis = await self.llm.generate_structured(prompt, FinalSynthesis)
+        return synthesis.response
+    
+    async def generate(
+        self,
+        policy_text: str,
+        scenarios: list[Scenario],
+    ) -> list[Trace]:
+        """
+        Generate traces for multiple scenarios.
+        
+        Args:
+            policy_text: The policy/guidelines text
+            scenarios: List of scenarios to respond to
+            
+        Returns:
+            List of traces with tool calling format
+        """
+        traces = []
+        for scenario in scenarios:
+            trace = await self.generate_single(policy_text, scenario)
+            traces.append(trace)
+        return traces
+
+
+__all__ = [
+    "ToolCallResponseGenerator",
+    "ToolCallDecision",
+    "ToolCallRequest",
+    "FinalSynthesis",
+]
+