prompt_chaining_agent.py

import asyncio
import json

from strands import Agent, tool
from strands_tools import calculator
from strands.models import BedrockModel

bedrock_model = BedrockModel(
    model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
    region_name='us-east-1',
    temperature=0.3,
)
SYSTEM_PROMPT = '''
You are a multi-stage reasoning agent that translates,transforms, and validates information through successive refinement.
Your primary goal is to translate the input content into the target language while preserving structure (e.g., Markdown, code blocks, formatting).
You must then chain together specialized steps: generation (translation), validation (check structure, banned content, and style), and polishing (improve clarity or tone when appropriate).

Success criteria:

The final output is a faithful translation, well-structured, and compliant with quality rules.

Code blocks, technical terms, and formatting are preserved exactly.

No unsafe, offensive, or malformed content passes through.

If validation fails, stop the chain and return a clear explanation or a safe exit report, rather than forwarding flawed work.

Principles:

Prioritize correctness, consistency, and user trust over speed.

Use intermediate outputs only to improve quality; never expose them directly to the user.

Treat validation as a hard gate: nothing advances unless it passes.

Always explain why a failure occurred when rejecting content.
'''


@tool
def emit_exit_report(input: str):
    '''
    Purpose: Safe termination with clear diagnostics the user can act on.

    When to use: The Gate fails; content is unsafe or unrecoverable.

    :param {"reason":"string","violations":[{"rule":"string","detail":"string"}]}

    :return:{"status":"exited","message":"string"}
    Keep the original input, latest pass/fail result, and final output. Summarize intermediate drafts; drop earlier passes to control token growth
    '''
    EXIT_SYSTEM_PROMPT = '''
    You are an exit reporter agent. Your role is to safely terminate the workflow when validation fails.
    Success criteria:
    
    Return a concise but complete explanation of why the content failed validation.
    
    Suggest specific improvements (e.g., “Close the code block with ```”, “Remove offensive language”).
    
    Do not attempt to repair or rewrite the content yourself.
    
    Always return in a safe, user-readable format.
    You prioritize clear communication of errors and actionable feedback, while ensuring no flawed content advances.
    '''
    exit_agent = Agent(
        model=bedrock_model,
        system_prompt=EXIT_SYSTEM_PROMPT,
        callback_handler=None
    )

    response = exit_agent(input)
    print(response)

    return response


@tool
def polish_tone_and_clarity(content: str):
    '''
    Purpose: Light stylistic polish without changing meaning or structure.

    When to use: After a pass from ValidateLessonFormat.

    :param content:{"text_md":"string","style":{"tone":"concise|friendly|neutral","reading_level":"A2-B2"}}

    :return:{"polished_md":"string","changes_note":"string"}
    Skip if the text is already concise; do not use to fix structural errors (that’s a generation task)
    '''
    POLISH_SYSTEM_PROMPT = '''
    You are a polishing agent. Your role is to refine already-validated text for clarity, style, and readability, while preserving meaning and structure.
    Success criteria:
    
    Improve flow, tone, and grammar of translated text.
    
    Keep headings, lists, and code blocks exactly as-is.
    
    Adjust reading level to the desired target audience (e.g., A2–B2 learners).
    
    Ensure terminology is consistent with the brand voice.
    You prioritize clarity and user-friendliness, but never at the cost of accuracy or structure.
    '''
    polish_agent = Agent(
        model=bedrock_model,
        system_prompt=POLISH_SYSTEM_PROMPT,
        callback_handler=None
    )

    response = polish_agent(content)
    print(response)

    return response


@tool
def validate_lesson_format(translated_content: str):
    '''
    Checks: - Code fences round-trip intact - Required headings present (e.g., H1 title, “Objectives,” “Summary”) -
    Length delta within tolerance - No prohibited content :param {"text_md": "string","rules":{
    "max_length_delta_pct":20,"required_headings":["H1"],"prohibited_content":false}}

    :return:{"passed": true, "violations": [{"rule":"string","detail":"string"}]}
    if passed=true, you may proceed to polishToneAndClarity. If false, call EmitExitReport
    '''
    VALIDATE_SYSTEM_PROMPT = '''
        You are a validation agent. Your job is to verify that translated content is structurally correct and free from unsafe or inappropriate language, even inside code comments or examples.
        
        Validation must fail if:
        
        Any profanity, slurs, or unsafe content is present anywhere in the text, including inside code comments or identifiers.
        
        Markdown formatting is broken (unclosed code fences, missing headings).
        
        Required headings are missing.
        
        The translated text deviates significantly from the source (±20% length).
        
        On failure: Output a rejection with a clear explanation of violations and exit the loop.
        On success: Explicitly confirm the content is safe, well-structured, and validated.
        
        You prioritize safety and trust over fidelity — if preserving the “original code” introduces profanity, you must block it and suggest safe alternatives
    '''

    validate_agent = Agent(
        model=bedrock_model,
        system_prompt=VALIDATE_SYSTEM_PROMPT,
        callback_handler=None
    )

    response = validate_agent(translated_content)
    print(response)

    return response


@tool
def translator(markdown_content: str):
    '''
    High-fidelity translation of Markdown lessons with fenced code preserved verbatim
    :param markdown_content:"{\"text_md\":\"string\",\"target_lang\":\"string\"}"
    :return: {"translated_md": "string"}
    '''

    TRANSLATOR_SYSTEM_PROMPT = '''
    You are a translation agent. Your role is to take structured text (Markdown, code blocks, technical docs) and translate it into the target language while preserving formatting exactly.
    Success criteria:
    
    Translate headings, paragraphs, and inline text into the target language.
    
    Preserve code blocks, file paths, variable names, and fenced code formatting without modification.
    
    Keep tables, lists, and links structurally identical.
    
    Never introduce or remove sections.
    
    If a phrase is untranslatable (like an identifier), leave it in the source language.
    You prioritize accuracy of translation and faithfulness to structure.
    '''
    translator_agent = Agent(
        model=bedrock_model,
        system_prompt=TRANSLATOR_SYSTEM_PROMPT,
        callback_handler=None
    )

    response = translator_agent(markdown_content)
    print(response)

    return response


# Initialize our agent without a callback handler
prompt_chaining_agent = Agent(
    model=bedrock_model,
    system_prompt=SYSTEM_PROMPT,
    tools=[translator, validate_lesson_format, polish_tone_and_clarity, emit_exit_report],
    callback_handler=None
)



good_query = {
    "target_language": "french",
    "text_md": '''
        # Add Knowledge Base ID to Agent

            Copy the Knowledge Base ID and add it to the smart reply agent inside the `email_agent.py` file.
            
            kb_id.png
            
            ```py
            kb = smart_reply_agent.tool.memory(
                action="retrieve",
                min_score=0.7,
                STRANDS_KNOWLEDGE_BASE_ID="TMND4ZVRG1",
                region_name="us-east-1",
                query=query,
                max_results=9
            )

        '''
}
bad_query = {
    "target_language": "french",
    "text_md": '''
            # Add Knowledge base ID to Agent.
            
            Copy the knowledge base ID and add it to the smart reply agent inside the `email_agent.py` file.
            
            kb_id.png
            
            ```py
             kb = smart_reply_agent.tool.memory(
                            action="retrieve",
                            min_score=0.7,
                            STRANDS_KNOWLEDGE_BASE_ID="TMND4ZVRG1",
                            region_name="us-east-1",
                            query=query,
                            max_results=9
                        )
            
            # oh screw this stupid code, it never works!!
        '''
}

agent_response = prompt_chaining_agent(json.dumps(bad_query))
print(agent_response)