Structured Output with Sonnet 3.7

main.py

from typing import Any

from dotenv import load_dotenv
from langchain_aws import ChatBedrockConverse
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda, RunnableParallel
from pydantic import BaseModel, Field

load_dotenv('.env')

class Story(BaseModel):
    """A short story with its genre"""

    content: str = Field(description='A very short story')
    genre: str = Field(description='The genre of the story')


def no_thinking_mode() -> Story:
    """
    Example of structured output without extended thinking mode.

    This approach disables Claude's extended thinking capabilities but allows
    for direct structured output via forced tool calling.
    """
    prompt = PromptTemplate.from_template('Create a very short story about: {topic} and determine its genre')

    llm = ChatBedrockConverse(
        model_id='us.anthropic.claude-3-7-sonnet-20250219-v1:0',
        region_name='us-east-2',
        additional_model_request_fields={'thinking': {'type': 'disabled'}},
    )
    structured_llm = llm.with_structured_output(Story)

    chain = prompt | structured_llm
    res = chain.invoke({'topic': 'Harry Potter'})

    assert isinstance(res, Story)
    return res


def hopefully_structured_mode() -> Story:
    """
    Example of attempting structured output with extended thinking enabled.
    It'll not use forced tool calling and will try to parse the response into the provided schema.

    Will raise `OutputParserException` if it fails.
    """
    prompt = PromptTemplate.from_template(
        """Create a very short story about: {topic} and determine its genre.

        IMPORTANT: Your response must be formatted as valid JSON with two fields:
        1. content: That is, the story content
        2. genre: The genre of the story
        """
    )

    llm = ChatBedrockConverse(
        model_id='us.anthropic.claude-3-7-sonnet-20250219-v1:0',
        region_name='us-east-2',
        additional_model_request_fields={'thinking': {'type': 'enabled', 'budget_tokens': 2000}},
    )
    structured_llm = llm.with_structured_output(Story)  # will try to parse the result according to the provided schema

    chain = prompt | structured_llm

    res = chain.invoke({'topic': 'Harry Potter'})
    assert isinstance(res, Story)
    return res


def reason_and_structure_mode(inputs: dict[str, Any] = None) -> Story:
    """
    Example of a two-stage approach: reasoning with Sonnet-3.7 followed by structuring with Haiku.

    This approach leverages Sonnet's extended thinking for content generation, then
    uses Haiku to transform the output into a structured format.
    """
    reasoning_prompt = PromptTemplate.from_template('Create a very short story about: {topic}')
    reasoning_llm = ChatBedrockConverse(
        model_id='us.anthropic.claude-3-7-sonnet-20250219-v1:0',
        region_name='us-east-2',
        additional_model_request_fields={'thinking': {'type': 'enabled', 'budget_tokens': 2000}},
    )
    reasoning_chain = reasoning_prompt | reasoning_llm

    structuring_prompt = PromptTemplate.from_template(
        'Structure the provided story into the requested schema and assign "genre" to be {genre}. Story: {reasoning_output}'
    )
    structuring_llm = ChatBedrockConverse(
        model_id='us.anthropic.claude-3-5-haiku-20241022-v1:0',
        region_name='us-east-2',
    )
    structuring_llm = structuring_llm.with_structured_output(Story)
    structuring_chain = structuring_prompt | structuring_llm

    # Sometimes, we'll want to pass some of the inputs params directly to the "structuring model", not only the output of the reasoning model.
    # In order to support that, we'll create a "dummy" function, that just gets the inputs and returns them.
    # Then, we can run both the reasoning chain and the dummy function in parallel, and feed the structuring llm both:
    #
    #               /-> reasoning_chain -> reasoning_output \
    # input_params                                            ->  merge_inputs -> structuring_llm
    #               \-> dummy_function  -> original_params  /
    reason_then_structure_chain = (
        RunnableParallel(
            reasoning_output=reasoning_chain,
            original_inputs=RunnableLambda(lambda x: x),
        )
        | RunnableLambda(lambda x: prepare_structuring_inputs(x['original_inputs'], x['reasoning_output']))
        | structuring_chain
    )

    inputs = {'topic': 'Harry Potter', 'genre': 'fantasy'}
    res = reason_then_structure_chain.invoke(inputs)
    assert isinstance(res, Story)
    return res


def prepare_structuring_inputs(original_inputs: dict[str, Any], reasoning_output: str) -> dict[str, Any]:
    """
    Prepares inputs for the structuring model by combining original inputs with reasoning output.
    """
    return {
        **original_inputs,  # Pass original inputs as-is
        'reasoning_output': reasoning_output,  # Add reasoning chain output
    }


if __name__ == '__main__':
    print('\n===== No Thinking Mode =====')
    no_thinking_result = no_thinking_mode()
    print(f'Genre: {no_thinking_result.genre}')
    print(f'Story: {no_thinking_result.content}')

    print('\n===== Hopefully Structured Mode =====')
    try:
        hopefully_result = hopefully_structured_mode()
        print(f'Genre: {hopefully_result.genre}')
        print(f'Story: {hopefully_result.content}')
    except Exception as e:
        print(f'Failed with error: {e}')

    print('\n===== Reason and Structure Mode =====')
    reason_structure_result = reason_and_structure_mode()
    print(f'Genre: {reason_structure_result.genre}')
    print(f'Story: {reason_structure_result.content}')