Sequential Chains: Chaining Multiple Steps

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

model = ChatOpenAI(model="gpt-4o", temperature=0)
parser = StrOutputParser()

# Step 1: Extract drug class
classify_prompt = ChatPromptTemplate.from_template(
    "What drug class is {drug}? Reply with just the class name."
)

# Step 2: Explain the class mechanism
explain_prompt = ChatPromptTemplate.from_template(
    "Explain the mechanism of {drug_class} drugs in 2 sentences."
)

# Sequential chain: output of step 1 becomes input for step 2
classify_chain = classify_prompt | model | parser
explain_chain = explain_prompt | model | parser

# Problem: classify_chain outputs a string but explain_chain expects {"drug_class": ...}
# Solution: wrap with RunnablePassthrough
from langchain_core.runnables import RunnablePassthrough

sequential_chain = (
    classify_chain
    | (lambda drug_class: {"drug_class": drug_class})  # Reshape output → dict
    | explain_chain
)

result = sequential_chain.invoke({"drug": "warfarin"})
print(result)  # "Anticoagulant drugs work by..."

# Pattern: progressively enrich the context dict through each step

chain = (
    # Start: {"drug": "warfarin"}
    RunnablePassthrough.assign(
        drug_class=(
            ChatPromptTemplate.from_template("What drug class is {drug}? One word.")
            | model | parser
        )
    )
    # Now: {"drug": "warfarin", "drug_class": "Anticoagulant"}
    | RunnablePassthrough.assign(
        mechanism=(
            ChatPromptTemplate.from_template("Explain {drug_class} mechanism in one sentence.")
            | model | parser
        )
    )
    # Now: {"drug": "warfarin", "drug_class": "...", "mechanism": "..."}
    | RunnablePassthrough.assign(
        monitoring=(
            ChatPromptTemplate.from_template(
                "For {drug} ({drug_class}): what lab values need monitoring?"
            )
            | model | parser
        )
    )
    # Now has: drug, drug_class, mechanism, monitoring
)

result = chain.invoke({"drug": "warfarin"})
# result is a dict with all four keys
print(result["drug"])
print(result["drug_class"])
print(result["mechanism"])
print(result["monitoring"])

from langchain_core.runnables import RunnablePassthrough, RunnableLambda

# Generate a full clinical drug monograph in stages

# Stage 1: Drug classification and basic info
basic_info_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a clinical pharmacist."),
    ("human", "For {drug}, provide: therapeutic class, mechanism of action in 2 sentences."),
])

# Stage 2: Dosing information
dosing_prompt = ChatPromptTemplate.from_template(
    "For {drug} (a {drug_class}):\n"
    "Provide standard adult dosing for the most common indication. "
    "Include dose, frequency, and route."
)

# Stage 3: Safety profile
safety_prompt = ChatPromptTemplate.from_template(
    "For {drug}:\n"
    "- Top 3 adverse effects\n"
    "- Major contraindications\n"
    "- Most important drug interaction"
)

# Stage 4: Compile final report
compile_prompt = ChatPromptTemplate.from_template(
    """Compile this information into a brief clinical monograph:

Drug: {drug}
Basic Info: {basic_info}
Dosing: {dosing}
Safety: {safety}

Format as a professional clinical summary."""
)


def extract_drug_class(text: str) -> str:
    """Extract the drug class from basic_info text."""
    # Simple extraction — in production use structured output
    lines = text.lower()
    for drug_class in ["anticoagulant", "biguanide", "beta-blocker", "ace inhibitor", "statin"]:
        if drug_class in lines:
            return drug_class.title()
    return "Unclassified"


clinical_report_chain = (
    # Stage 1
    RunnablePassthrough.assign(
        basic_info=basic_info_prompt | model | parser
    )
    # Extract class from basic_info
    | RunnablePassthrough.assign(
        drug_class=RunnableLambda(lambda d: extract_drug_class(d["basic_info"]))
    )
    # Stage 2: Dosing (has access to drug and drug_class)
    | RunnablePassthrough.assign(
        dosing=dosing_prompt | model | parser
    )
    # Stage 3: Safety
    | RunnablePassthrough.assign(
        safety=safety_prompt | model | parser
    )
    # Stage 4: Compile
    | compile_prompt | model | parser
)

report = clinical_report_chain.invoke({"drug": "warfarin"})
print(report)

from operator import itemgetter
from langchain_core.runnables import RunnableLambda

# itemgetter: pick specific keys from dict
select_keys = RunnableLambda(lambda d: {
    "drug": d["drug"],
    "final_answer": d["answer"],
})

# Or use itemgetter for a single key
get_answer = itemgetter("answer")

chain = (
    RunnablePassthrough.assign(
        intermediate=some_chain
    )
    | RunnablePassthrough.assign(
        answer=final_chain
    )
    | RunnableLambda(lambda d: {"drug": d["drug"], "answer": d["answer"]})
    # Drop "intermediate" — only keep what matters
)

from langchain_core.runnables import RunnableLambda

def safe_drug_lookup(inputs: dict) -> dict:
    """Safely retrieve drug info with error handling."""
    drug = inputs.get("drug", "").strip()
    if not drug:
        return {"drug": drug, "error": "Empty drug name"}
    
    # Validate drug name is reasonable
    if len(drug) > 100 or not drug.replace(" ", "").isalpha():
        return {"drug": drug, "error": "Invalid drug name format"}
    
    return {"drug": drug}

safe_chain = (
    RunnableLambda(safe_drug_lookup)
    | RunnableLambda(lambda d: d if "error" not in d else {"drug": d["drug"], "basic_info": f"Error: {d['error']}"})
    | RunnablePassthrough.assign(
        # Only runs basic_info chain if "basic_info" not already set (no error)
        basic_info=RunnableLambda(
            lambda d: (basic_info_prompt | model | parser).invoke(d)
            if "basic_info" not in d
            else d["basic_info"]
        )
    )
)

from langchain_core.runnables import RunnableParallel

# Sequential: when step B needs output from step A
sequential = (
    RunnablePassthrough.assign(drug_class=classify_chain)
    | RunnablePassthrough.assign(mechanism=explain_chain)  # Needs drug_class
)

# Parallel: when steps are independent
parallel = RunnableParallel(
    mechanism=mechanism_chain,        # Independent
    side_effects=side_effects_chain,  # Independent
    interactions=interactions_chain,  # Independent
)

# Hybrid: some parallel, some sequential
hybrid = (
    # Parallel stage 1
    RunnablePassthrough.assign(
        **{key: chain for key, chain in parallel_stage_1.items()}
    )
    # Sequential stage 2 (needs results from stage 1)
    | RunnablePassthrough.assign(
        final_summary=summary_chain
    )
)

# Bad: 3 separate LLM calls for info that could come from 1
bad_chain = (
    RunnablePassthrough.assign(drug_class=classify_chain)
    | RunnablePassthrough.assign(mechanism=explain_chain)
    | RunnablePassthrough.assign(monitoring=monitoring_chain)
)

# Better: one structured LLM call for all three
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser

class DrugBasics(BaseModel):
    drug_class: str
    mechanism: str = Field(description="One sentence mechanism")
    monitoring: list[str] = Field(description="Key monitoring parameters")

pydantic_parser = PydanticOutputParser(pydantic_object=DrugBasics)
combined_prompt = ChatPromptTemplate.from_messages([
    ("system", "Extract drug information. {format_instructions}"),
    ("human", "Drug: {drug}"),
]).partial(format_instructions=pydantic_parser.get_format_instructions())

efficient_chain = combined_prompt | model | pydantic_parser
basics = efficient_chain.invoke({"drug": "warfarin"})
# One LLM call → drug_class, mechanism, monitoring all at once