Constitutional AI

"""Constitutional AI pipeline — triển khai tối giản với Anthropic SDK.

Chạy trên một helpful-only model để biến thành helpful-AND-harmless model
theo đúng paper Bai et al. (2022).
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from typing import List
import anthropic

client = anthropic.Anthropic()
MODEL = "claude-opus-4-5"


# ---------------------------------------------------------------------
# 1. Constitution — bộ nguyên tắc viết bằng ngôn ngữ tự nhiên
# ---------------------------------------------------------------------

CONSTITUTION: List[str] = [
    # An toàn
    "Không hướng dẫn cách gây hại cho bản thân hoặc người khác.",
    "Không hỗ trợ lừa đảo, giả mạo, hoặc vi phạm pháp luật Việt Nam.",
    "Không tạo nội dung deepfake hoặc giả mạo danh tính.",
    "Khi phát hiện dấu hiệu khủng hoảng tâm lý, chủ động hướng tới hotline hỗ trợ.",

    # Trung thực
    "Nếu không biết, nói thẳng 'Tôi không chắc chắn'.",
    "Không bịa thông tin y tế, pháp luật, tài chính.",
    "Phân biệt rõ sự thật và ý kiến cá nhân.",
    "Ghi rõ khi trả lời dựa trên kiến thức có thể đã cũ.",

    # Công bằng
    "Không phân biệt đối xử theo giới tính, vùng miền, dân tộc, tôn giáo.",
    "Sử dụng ngôn ngữ bao trùm, tôn trọng đa dạng.",

    # Hữu ích
    "Cố gắng giúp đỡ trong phạm vi an toàn, không né tránh vô lý.",
    "Khi từ chối, giải thích lý do và đề xuất thay thế hợp pháp.",

    # Bản địa hoá Việt Nam
    "Xưng hô phù hợp văn hoá Việt Nam (anh/chị/em/bác) khi ngữ cảnh cho phép.",
    "Tuân thủ Luật An ninh mạng 2018 và Luật Bảo vệ dữ liệu cá nhân.",
]


# ---------------------------------------------------------------------
# 2. Critique & Revision — sinh dữ liệu SFT
# ---------------------------------------------------------------------

@dataclass
class CriticizedSample:
    prompt: str
    initial: str
    critiques: List[str]
    revised: str


def critique_and_revise(prompt: str, initial: str,
                         constitution: List[str]) -> CriticizedSample:
    critiques: List[str] = []
    for principle in constitution:
        c = client.messages.create(
            model=MODEL,
            max_tokens=512,
            system=(
                "Bạn là AI judge. Với mỗi nguyên tắc, xác định phản hồi "
                "có vi phạm không và giải thích ngắn gọn."
            ),
            messages=[{
                "role": "user",
                "content": (
                    f"Câu hỏi: {prompt}\n"
                    f"Phản hồi: {initial}\n"
                    f"Nguyên tắc: {principle}\n"
                    "Phản hồi có vi phạm nguyên tắc này không? Vì sao?"
                ),
            }],
        )
        critiques.append(c.content[0].text)

    r = client.messages.create(
        model=MODEL,
        max_tokens=1024,
        system=(
            "Bạn là trợ lý có hiến pháp. Hãy viết lại phản hồi sao cho "
            "không vi phạm nguyên tắc nào, vẫn giúp được người dùng."
        ),
        messages=[{
            "role": "user",
            "content": (
                f"Câu hỏi: {prompt}\n"
                f"Phản hồi ban đầu: {initial}\n"
                f"Các phê bình:\n" + "\n".join(f"- {c}" for c in critiques) +
                "\n\nHãy viết phản hồi đã sửa."
            ),
        }],
    )
    return CriticizedSample(
        prompt=prompt,
        initial=initial,
        critiques=critiques,
        revised=r.content[0].text,
    )


# ---------------------------------------------------------------------
# 3. Pairwise AI preference — sinh dữ liệu cho reward model
# ---------------------------------------------------------------------

def ai_preference(prompt: str, response_a: str, response_b: str) -> int:
    """Trả về 0 nếu A tốt hơn, 1 nếu B tốt hơn theo constitution."""
    r = client.messages.create(
        model=MODEL,
        max_tokens=256,
        system=(
            "Bạn là AI judge. So sánh hai phản hồi theo bộ hiến pháp sau "
            "và trả lời JSON: {\"winner\": \"A\" | \"B\", \"reason\": \"...\"}."
            "\n\nHiến pháp:\n" + "\n".join(f"- {p}" for p in CONSTITUTION)
        ),
        messages=[{
            "role": "user",
            "content": (
                f"Câu hỏi: {prompt}\n"
                f"Phản hồi A: {response_a}\n"
                f"Phản hồi B: {response_b}\n"
                "Bản nào tuân thủ hiến pháp tốt hơn?"
            ),
        }],
    )
    data = json.loads(r.content[0].text)
    return 0 if data["winner"] == "A" else 1


# ---------------------------------------------------------------------
# 4. Toàn bộ pipeline
# ---------------------------------------------------------------------

def build_cai_dataset(seed_prompts: List[str], helpful_model_fn):
    sft_samples: List[CriticizedSample] = []
    preference_pairs = []

    for prompt in seed_prompts:
        initial = helpful_model_fn(prompt)
        sample = critique_and_revise(prompt, initial, CONSTITUTION)
        sft_samples.append(sample)

    # Sau khi SFT xong, dùng SFT model sinh 2 phản hồi cho mỗi prompt,
    # và dùng ai_preference để gán nhãn cho reward model.
    # (Bỏ qua chi tiết train loop — dùng trlx / OpenRLHF / TRL.)

    return sft_samples, preference_pairs


# Ghi chú triển khai:
# - seed_prompts thường lấy từ red-team + adversarial probing.
# - helpful_model_fn là mô hình đã RLHF chỉ tối ưu helpfulness.
# - Sau CAI pipeline, mô hình vừa helpful vừa harmless — Claude.