synkro 0.4.12__py3-none-any.whl

synkro/schemas.py

@@ -0,0 +1,472 @@
+"""Pydantic schemas for structured LLM outputs and validation."""
+
+from typing import Literal
+from pydantic import BaseModel, Field
+
+
+# =============================================================================
+# SCENARIO SCHEMAS
+# =============================================================================
+
+
+class ScenarioOutput(BaseModel):
+    """Output schema for scenario generation."""
+
+    scenario: str = Field(description="Detailed scenario description")
+    context: str = Field(description="Relevant background information")
+
+
+class ScenariosArray(BaseModel):
+    """Array of generated scenarios."""
+
+    scenarios: list[ScenarioOutput]
+
+
+# =============================================================================
+# POLICY ANALYSIS SCHEMAS
+# =============================================================================
+
+
+class PolicyComplexity(BaseModel):
+    """Policy complexity analysis for auto-detecting optimal turns."""
+
+    variable_count: int = Field(
+        description="Number of variables/conditions in the policy (rules, exceptions, conditions)"
+    )
+    complexity_level: Literal["simple", "conditional", "complex"] = Field(
+        description="Overall complexity: simple (1 var), conditional (2-3 vars), complex (4+ vars)"
+    )
+    recommended_turns: int = Field(
+        ge=1, le=6, description="Recommended conversation turns based on complexity"
+    )
+    reasoning: str = Field(description="Brief explanation of the complexity assessment")
+
+
+class PlanCategory(BaseModel):
+    """A category in the generation plan."""
+
+    name: str = Field(description='Short category name (e.g., "Consent Violations", "Edge Cases")')
+    description: str = Field(description="What this category tests")
+    traces: int = Field(ge=1, description="Number of traces to generate for this category")
+
+
+class PolicyPlan(BaseModel):
+    """LLM-generated plan for dataset generation."""
+
+    categories: list[PlanCategory] = Field(
+        min_length=2, max_length=10, description="Scenario categories with trace allocations"
+    )
+    reasoning: str = Field(
+        description="Explanation of why these categories were chosen based on policy content"
+    )
+
+
+# =============================================================================
+# CHAT MESSAGE SCHEMAS
+# =============================================================================
+
+
+class ChatMessage(BaseModel):
+    """A single chat message in OpenAI format."""
+
+    role: Literal["system", "user", "assistant"] = Field(description="Message role")
+    content: str = Field(description="Message content")
+
+
+class ConversationOutput(BaseModel):
+    """Output from response generation - a complete conversation."""
+
+    index: int = Field(description="Scenario index (0-based)")
+    messages: list[ChatMessage] = Field(
+        description="Full conversation with system, user, and assistant messages"
+    )
+
+
+class BatchedConversations(BaseModel):
+    """Batch of generated conversations."""
+
+    conversations: list[ConversationOutput]
+
+
+# =============================================================================
+# GRADING SCHEMAS
+# =============================================================================
+
+
+class GradeOutput(BaseModel):
+    """Grading result for a single response."""
+
+    index: int = Field(description="Scenario index (0-based)")
+    passed: bool = Field(
+        alias="pass", description="Is the response FULLY correct, policy-compliant, and format-valid?"
+    )
+    policy_violations: list[str] = Field(
+        default_factory=list,
+        description="Specific policy rules that were violated or misinterpreted",
+    )
+    missing_citations: list[str] = Field(
+        default_factory=list,
+        description="Policy sections that should have been cited but were not",
+    )
+    incomplete_reasoning: list[str] = Field(
+        default_factory=list, description="Logical gaps or missing steps in the chain of thought"
+    )
+    vague_recommendations: list[str] = Field(
+        default_factory=list,
+        description="Recommendations that need to be more specific or actionable",
+    )
+    feedback: str = Field(description="Summary of how to fix the issues")
+
+    class Config:
+        populate_by_name = True
+
+
+class BatchedGrades(BaseModel):
+    """Batch of grading results."""
+
+    grades: list[GradeOutput]
+
+
+# =============================================================================
+# SINGLE-ITEM SCHEMAS (for parallel generation)
+# =============================================================================
+
+
+class SingleResponse(BaseModel):
+    """Single response output for parallel generation."""
+
+    messages: list[ChatMessage] = Field(
+        min_length=3, max_length=3, description="Exactly 3 messages: system, user, assistant"
+    )
+
+
+class MultiTurnResponse(BaseModel):
+    """Multi-turn response output for complexity-driven generation."""
+
+    messages: list[ChatMessage] = Field(
+        min_length=3,
+        description="Conversation messages (variable length based on turn count)"
+    )
+    turn_count: int = Field(
+        ge=1, le=10,
+        description="Number of user-assistant exchanges in this conversation"
+    )
+
+
+class SingleGrade(BaseModel):
+    """Single grade output for parallel generation."""
+
+    passed: bool = Field(
+        alias="pass", description="Is the response FULLY correct, policy-compliant, and format-valid?"
+    )
+    policy_violations: list[str] = Field(
+        default_factory=list, description="Specific policy rules that were violated"
+    )
+    missing_citations: list[str] = Field(
+        default_factory=list, description="Policy sections that should have been cited"
+    )
+    incomplete_reasoning: list[str] = Field(
+        default_factory=list, description="Logical gaps or missing reasoning steps"
+    )
+    vague_recommendations: list[str] = Field(
+        default_factory=list, description="Recommendations that need to be more specific"
+    )
+    feedback: str = Field(description='Summary of issues or "Correct" if passing')
+
+    class Config:
+        populate_by_name = True
+
+
+# =============================================================================
+# MULTI-TURN SCHEMAS
+# =============================================================================
+
+
+class FollowUpQuestion(BaseModel):
+    """A follow-up question for multi-turn conversations."""
+
+    index: int = Field(description="Scenario index")
+    question: str = Field(description="Follow-up question from the user")
+    question_type: Literal["clarification", "edge_case", "what_if", "specificity", "challenge"] = (
+        Field(description="Type of follow-up")
+    )
+
+
+class TurnGrade(BaseModel):
+    """Grade for a single turn in a multi-turn conversation."""
+
+    turn_index: int = Field(description="Which turn (0-based, only assistant turns)")
+    passed: bool = Field(alias="pass", description="Does this turn pass all criteria?")
+    policy_violations: list[str] = Field(
+        default_factory=list, description="Policy violations in this turn"
+    )
+    missing_citations: list[str] = Field(
+        default_factory=list, description="Missing citations in this turn"
+    )
+    incomplete_reasoning: list[str] = Field(
+        default_factory=list, description="Reasoning gaps in this turn"
+    )
+    vague_recommendations: list[str] = Field(
+        default_factory=list, description="Vague recommendations in this turn"
+    )
+    feedback: str = Field(description="Specific feedback for this turn")
+
+    class Config:
+        populate_by_name = True
+
+
+class ConversationGrade(BaseModel):
+    """Full grading for a multi-turn conversation."""
+
+    index: int = Field(description="Scenario index")
+    overall_pass: bool = Field(description="Does the ENTIRE conversation pass?")
+    turn_grades: list[TurnGrade] = Field(description="Grade for each assistant turn")
+    coherence_pass: bool = Field(
+        description="Is the conversation coherent with no contradictions?"
+    )
+    coherence_issues: list[str] = Field(
+        default_factory=list, description="Any contradictions or incoherence across turns"
+    )
+    progressive_depth: bool = Field(
+        description="Does each turn build on previous context appropriately?"
+    )
+    overall_feedback: str = Field(
+        description="Summary of what needs to be fixed across the conversation"
+    )
+
+
+# =============================================================================
+# AGENTIC SCHEMAS
+# =============================================================================
+
+
+class ToolCall(BaseModel):
+    """A tool call in an agentic trace."""
+
+    tool_name: str = Field(description="Name of the tool to call")
+    arguments: dict[str, str] = Field(description="Arguments to pass to the tool")
+
+
+class AgenticStep(BaseModel):
+    """A single step in an agentic trace."""
+
+    reasoning: str = Field(description="Reasoning before tool call")
+    tool_name: str = Field(description="Tool to call")
+    tool_args: dict = Field(description="Tool arguments")
+
+
+class AgenticTrace(BaseModel):
+    """Complete agentic trace with tool usage."""
+
+    index: int = Field(description="Scenario index")
+    steps: list[AgenticStep] = Field(description="Steps of tool usage")
+    final_answer: str = Field(description="Final comprehensive answer")
+
+
+# =============================================================================
+# TOOL CALL GRADING SCHEMAS
+# =============================================================================
+
+
+class ToolCallGrade(BaseModel):
+    """Grading result for a tool call trace.
+    
+    Evaluates tool usage on four criteria:
+    - Tool Selection: Did they use the right tool?
+    - Parameter Accuracy: Were the parameters correct?
+    - Response Synthesis: Did they use tool results correctly?
+    - Timing: Did they call tools at the right time?
+    """
+    
+    passed: bool = Field(
+        alias="pass",
+        description="Does the trace pass ALL criteria?"
+    )
+    
+    # Criterion 1: Tool Selection
+    tool_selection_correct: bool = Field(
+        description="Did the assistant choose the appropriate tool for the task?"
+    )
+    tool_selection_issues: list[str] = Field(
+        default_factory=list,
+        description="Specific issues with tool selection (wrong tool, missing tool, unnecessary tool)"
+    )
+    
+    # Criterion 2: Parameter Accuracy
+    parameters_valid: bool = Field(
+        description="Were the tool parameters correct (types, values, required fields)?"
+    )
+    parameter_issues: list[str] = Field(
+        default_factory=list,
+        description="Specific issues with parameters (wrong type, invalid value, missing required)"
+    )
+    
+    # Criterion 3: Response Synthesis
+    synthesis_accurate: bool = Field(
+        description="Did the assistant correctly use tool results without hallucination?"
+    )
+    synthesis_issues: list[str] = Field(
+        default_factory=list,
+        description="Specific issues with synthesis (hallucinated data, ignored results, misinterpreted)"
+    )
+    
+    # Criterion 4: Timing
+    timing_appropriate: bool = Field(
+        description="Did the assistant call tools at the right moment?"
+    )
+    timing_issues: list[str] = Field(
+        default_factory=list,
+        description="Specific issues with timing (premature call, delayed call, should have called earlier)"
+    )
+    
+    # Overall feedback
+    feedback: str = Field(
+        description="Summary of issues or 'Correct' if passing"
+    )
+    
+    class Config:
+        populate_by_name = True
+    
+    def get_all_issues(self) -> list[str]:
+        """Get all issues combined."""
+        return (
+            self.tool_selection_issues
+            + self.parameter_issues
+            + self.synthesis_issues
+            + self.timing_issues
+        )
+
+
+# =============================================================================
+# GOLDEN TRACE SCHEMAS
+# =============================================================================
+
+
+class RuleExtraction(BaseModel):
+    """A single rule extracted from the policy."""
+
+    rule_id: str = Field(description="Unique identifier (e.g., 'R001')")
+    text: str = Field(description="Exact rule text from the policy")
+    condition: str = Field(description="The 'if' part - when this rule applies")
+    action: str = Field(description="The 'then' part - what happens")
+    dependencies: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs that must be evaluated before this rule"
+    )
+    category: Literal["constraint", "permission", "procedure", "exception"] = Field(
+        description="Type of rule"
+    )
+
+
+class LogicMapOutput(BaseModel):
+    """Output schema for logic extraction - the complete DAG of rules."""
+
+    rules: list[RuleExtraction] = Field(
+        description="All rules extracted from the policy"
+    )
+    root_rules: list[str] = Field(
+        description="Rule IDs with no dependencies (entry points)"
+    )
+    reasoning: str = Field(
+        description="Explanation of rule extraction and relationship identification"
+    )
+
+
+class RefinedLogicMapOutput(BaseModel):
+    """Output schema for Logic Map refinement based on user feedback."""
+
+    rules: list[RuleExtraction] = Field(
+        description="All rules in the refined Logic Map (modified and unmodified)"
+    )
+    root_rules: list[str] = Field(
+        description="Rule IDs with no dependencies (entry points)"
+    )
+    changes_summary: str = Field(
+        description="Brief summary of changes made (e.g., 'Added R009, removed R005')"
+    )
+    reasoning: str = Field(
+        description="Explanation of how user feedback was interpreted and applied"
+    )
+
+
+class GoldenScenarioOutput(BaseModel):
+    """Output schema for a single golden scenario."""
+
+    description: str = Field(description="The user's request or question")
+    context: str = Field(default="", description="Additional context")
+    scenario_type: Literal["positive", "negative", "edge_case", "irrelevant"] = Field(
+        description="Type of scenario"
+    )
+    target_rule_ids: list[str] = Field(
+        description="Rule IDs this scenario tests"
+    )
+    expected_outcome: str = Field(
+        description="Expected behavior based on rules"
+    )
+
+
+class GoldenScenariosArray(BaseModel):
+    """Array of generated golden scenarios."""
+
+    scenarios: list[GoldenScenarioOutput]
+
+
+class ReasoningStepOutput(BaseModel):
+    """A single step in the Chain-of-Thought reasoning."""
+
+    rule_id: str = Field(description="The rule being evaluated")
+    rule_text: str = Field(description="The text of the rule")
+    applies: bool = Field(description="Whether this rule applies")
+    reasoning: str = Field(description="Why the rule does/doesn't apply")
+    exclusions: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs excluded because this rule applies"
+    )
+
+
+class GoldenTraceOutput(BaseModel):
+    """Output schema for a golden trace with grounded reasoning."""
+
+    messages: list[ChatMessage] = Field(
+        description="The conversation messages"
+    )
+    reasoning_chain: list[ReasoningStepOutput] = Field(
+        description="Step-by-step reasoning with rule citations"
+    )
+    rules_applied: list[str] = Field(
+        description="Rule IDs that were applied in the response"
+    )
+    rules_excluded: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs that were explicitly excluded and why"
+    )
+
+
+class VerificationOutput(BaseModel):
+    """Output schema for trace verification against Logic Map."""
+
+    passed: bool = Field(description="Whether the trace passed verification")
+    issues: list[str] = Field(
+        default_factory=list,
+        description="List of issues found"
+    )
+    skipped_rules: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs that should have been applied but weren't"
+    )
+    hallucinated_rules: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs cited that don't exist or don't apply"
+    )
+    contradictions: list[str] = Field(
+        default_factory=list,
+        description="Logical contradictions found"
+    )
+    rules_verified: list[str] = Field(
+        default_factory=list,
+        description="Rule IDs correctly applied"
+    )
+    feedback: str = Field(
+        description="Summary of verification result"
+    )
+

synkro/types/__init__.py

@@ -0,0 +1,41 @@
+"""Type definitions for Synkro.
+
+Usage:
+    from synkro.types import DatasetType, Message, Trace
+    from synkro.types import ToolDefinition, ToolCall, ToolFunction
+"""
+
+from synkro.types.core import (
+    Role,
+    Message,
+    Scenario,
+    Trace,
+    GradeResult,
+    Plan,
+    Category,
+)
+from synkro.types.dataset_type import DatasetType
+from synkro.types.tool import (
+    ToolDefinition,
+    ToolCall,
+    ToolFunction,
+    ToolResult,
+)
+
+__all__ = [
+    # Dataset type
+    "DatasetType",
+    # Core types
+    "Role",
+    "Message",
+    "Scenario",
+    "Trace",
+    "GradeResult",
+    "Plan",
+    "Category",
+    # Tool types
+    "ToolDefinition",
+    "ToolCall",
+    "ToolFunction",
+    "ToolResult",
+]

synkro/types/core.py

@@ -0,0 +1,126 @@
+"""Core Pydantic models for Synkro."""
+
+from typing import Literal, Any
+from pydantic import BaseModel, Field
+
+
+Role = Literal["system", "user", "assistant", "tool"]
+
+
+class Message(BaseModel):
+    """
+    A single message in a conversation.
+    
+    Supports both regular chat messages and tool-calling messages.
+    
+    Examples:
+        >>> # Regular message
+        >>> Message(role="user", content="Hello")
+        
+        >>> # Assistant with tool call (tool_calls is list of dicts or ToolCall objects)
+        >>> Message(role="assistant", content=None, tool_calls=[...])
+        
+        >>> # Tool response
+        >>> Message(role="tool", content="Result", tool_call_id="call_123")
+    """
+
+    role: Role
+    content: str | None = None
+    tool_calls: list[Any] | None = Field(
+        default=None,
+        description="Tool calls made by the assistant (list of ToolCall or dicts)"
+    )
+    tool_call_id: str | None = Field(
+        default=None,
+        description="ID of the tool call this message responds to (for tool role)"
+    )
+    
+    def model_post_init(self, __context) -> None:
+        """Validate message structure based on role."""
+        # For backwards compatibility, ensure content is string for non-tool roles
+        if self.role in ("system", "user") and self.content is None:
+            self.content = ""
+
+
+class Scenario(BaseModel):
+    """A test scenario for trace generation."""
+
+    description: str = Field(description="The scenario description")
+    context: str = Field(description="Additional context and background")
+    category: str | None = Field(default=None, description="Category this scenario belongs to")
+
+
+class GradeResult(BaseModel):
+    """Result of grading a trace."""
+
+    passed: bool = Field(description="Whether the trace passes quality checks")
+    issues: list[str] = Field(default_factory=list, description="List of issues found")
+    feedback: str = Field(default="", description="Summary feedback for improvement")
+
+
+class Trace(BaseModel):
+    """A complete training trace with messages and metadata."""
+
+    messages: list[Message] = Field(description="The conversation messages")
+    scenario: Scenario = Field(description="The scenario this trace was generated from")
+    grade: GradeResult | None = Field(default=None, description="Grading result if graded")
+
+    # Golden Trace metadata (for verification)
+    reasoning_chain: list[Any] | None = Field(default=None, description="Chain-of-thought reasoning steps with rule citations")
+    rules_applied: list[str] | None = Field(default=None, description="Rule IDs that were applied in the response")
+    rules_excluded: list[str] | None = Field(default=None, description="Rule IDs that were explicitly excluded")
+
+    @property
+    def system_message(self) -> str | None:
+        """Get the system message content."""
+        for m in self.messages:
+            if m.role == "system":
+                return m.content
+        return None
+
+    @property
+    def user_message(self) -> str:
+        """Get the first user message content."""
+        for m in self.messages:
+            if m.role == "user":
+                return m.content or ""
+        return ""
+
+    @property
+    def assistant_message(self) -> str:
+        """Get the last assistant message content."""
+        for m in reversed(self.messages):
+            if m.role == "assistant":
+                return m.content or ""
+        return ""
+    
+    @property
+    def has_tool_calls(self) -> bool:
+        """Check if this trace contains any tool calls."""
+        for m in self.messages:
+            if m.tool_calls:
+                return True
+        return False
+
+
+class Category(BaseModel):
+    """A category for organizing scenarios."""
+
+    name: str = Field(description="Category name")
+    description: str = Field(description="What this category tests")
+    count: int = Field(description="Number of traces to generate for this category")
+
+
+class Plan(BaseModel):
+    """A generation plan with categories and complexity analysis."""
+
+    categories: list[Category] = Field(description="Categories with trace allocations")
+    reasoning: str = Field(description="Explanation of why these categories were chosen")
+    recommended_turns: int = Field(
+        default=1,
+        description="Recommended conversation turns based on policy complexity"
+    )
+    complexity_level: Literal["simple", "conditional", "complex"] = Field(
+        default="simple",
+        description="Policy complexity level: simple (1-2 turns), conditional (3 turns), complex (5+ turns)"
+    )

synkro/types/dataset_type.py

@@ -0,0 +1,30 @@
+"""Dataset type enum for steering generation pipeline."""
+
+from enum import Enum
+
+
+class DatasetType(str, Enum):
+    """
+    Type of dataset to generate.
+
+    The dataset type determines:
+    - Prompts used for scenario and response generation
+    - Grading criteria
+    - Output format and schema
+
+    Examples:
+        >>> from synkro import DatasetType
+        >>> synkro.generate(policy, dataset_type=DatasetType.QA)
+        >>> synkro.generate(policy, dataset_type=DatasetType.SFT)
+        >>> synkro.generate(policy, dataset_type=DatasetType.TOOL_CALL, tools=[...])
+    """
+
+    QA = "qa"
+    """Question-Answer pairs: {question, answer, context}"""
+
+    SFT = "sft"
+    """Supervised Fine-Tuning: {messages: [system, user, assistant]}"""
+
+    TOOL_CALL = "tool_call"
+    """Tool Calling: {messages: [..., {tool_calls: [...]}, {role: tool}, ...]}"""
+