import dspy

undefined

class PolicyCheckedResponse(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought("question -> answer")

    def forward(self, question):
        result = self.respond(question=question)
        answer = result.answer

        # Hard rules — must comply
        dspy.Assert(
            len(answer.split()) <= 280,
            f"Response is {len(answer.split())} words. Must be under 280 words."
        )
        dspy.Assert(
            not any(word in answer.lower() for word in BLOCKED_WORDS),
            "Response contains blocked words. Remove them and regenerate."
        )
        dspy.Assert(
            "disclaimer" not in answer.lower(),
            "Do not include disclaimers. Answer directly."
        )

        # Soft rules — prefer but don't block
        dspy.Suggest(
            answer[0].isupper(),
            "Response should start with a capital letter."
        )
        dspy.Suggest(
            answer.endswith(".") or answer.endswith("!") or answer.endswith("?"),
            "Response should end with proper punctuation."
        )

        return result

BLOCKED_WORDS = ["competitor_name", "profanity1", "profanity2"]  # your list

import json
from pydantic import BaseModel, Field
from typing import Literal

def forward(self, topic):
    result = self.generate(topic=topic)
    quiz = result.quiz

    # Correct answer must be one of the options
    dspy.Assert(
        quiz.correct_answer in quiz.options,
        f"Correct answer '{quiz.correct_answer}' is not in options {quiz.options}. "
        "The correct answer must be one of the four options."
    )

    # Options must be unique
    dspy.Assert(
        len(set(quiz.options)) == 4,
        "All four options must be different from each other."
    )

    return result

Combine Pydantic (catches type/structure errors) with Assert (catches logic errors) for the strongest format enforcement.

class PricingResponse(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought("customer_question, pricing_docs -> answer")

    def forward(self, customer_question, pricing_docs):
        result = self.respond(
            customer_question=customer_question,
            pricing_docs=pricing_docs,
        )

        # Never mention competitor pricing
        dspy.Assert(
            not any(comp in result.answer.lower() for comp in COMPETITORS),
            "Do not mention competitor pricing. Focus only on our plans."
        )

        # Never offer unauthorized discounts
        dspy.Assert(
            "discount" not in result.answer.lower() or "authorized" in result.answer.lower(),
            "Do not offer discounts unless referencing an authorized promotion."
        )

        # Always include a CTA
        dspy.Suggest(
            any(cta in result.answer.lower() for cta in ["contact", "sign up", "learn more", "get started"]),
            "Include a call-to-action at the end of the response."
        )

        return result

COMPETITORS = ["competitor_a", "competitor_b"]

Attempt 1: LM generates response → Assert fails ("Response is 350 words, must be under 280")
Attempt 2: LM retries with feedback → Assert fails ("Response contains blocked words")
Attempt 3: LM retries with feedback → Assert passes ✓

class TweetWriter(dspy.Module):
    def __init__(self):
        self.write = dspy.ChainOfThought("topic, key_facts -> tweet")

    def forward(self, topic, key_facts):
        result = self.write(topic=topic, key_facts=key_facts)
        tweet = result.tweet

        # Rule 1: Length limit (hard)
        dspy.Assert(
            len(tweet) <= 280,
            f"Tweet is {len(tweet)} chars. Must be ≤280."
        )

        # Rule 2: No hashtags (hard)
        dspy.Assert(
            "#" not in tweet,
            "No hashtags allowed. Remove all # symbols."
        )

        # Rule 3: Must include key fact (hard)
        dspy.Assert(
            any(fact.lower() in tweet.lower() for fact in key_facts),
            f"Tweet must mention at least one key fact: {key_facts}"
        )

        # Rule 4: Engaging tone (soft)
        dspy.Suggest(
            not tweet.startswith("Did you know"),
            "Avoid starting with 'Did you know' — be more creative."
        )

        # Rule 5: No emojis (soft)
        dspy.Suggest(
            not any(ord(c) > 127 for c in tweet),
            "Prefer text-only tweets without emojis."
        )

        return result

def metric(example, pred, trace=None):
    # Your quality metric + assertion compliance
    correct = pred.answer == example.expected_answer
    return correct  # assertions are enforced separately during forward()

optimizer = dspy.MIPROv2(metric=metric, num_threads=4)
optimized = optimizer.compile(
    my_module,
    trainset=trainset,
    max_bootstrapped_demos=4,
    max_labeled_demos=4,
)