"""runnable_passthrough_demo.py — pass data through a chain
Run with venv active:  python runnable_passthrough_demo.py
"""
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableParallel

load_dotenv()

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

explain_prompt = ChatPromptTemplate.from_template(
    "What does the HTML <{tag}> tag do? One short sentence."
)
explain_chain = explain_prompt | model | parser

tip_prompt = ChatPromptTemplate.from_template(
    "Tag: <{tag}>. Explanation: {explanation}\nGive one beginner tip. One short sentence."
)

tags = ["a", "img"]

# 1. RunnablePassthrough() — passes input through unchanged
print("=== RunnablePassthrough() ===")
passthrough = RunnablePassthrough()
for tag in tags:
    data = passthrough.invoke({"tag": tag})
    print(f"in:  {{'tag': '{tag}'}}")
    print(f"out: {data}\n")

# 2. RunnablePassthrough.assign() — keep input, add new keys
assign_chain = (
    RunnablePassthrough.assign(explanation=explain_chain)
    | tip_prompt
    | model
    | parser
)

print("=== RunnablePassthrough.assign() ===")
for tag in tags:
    tip = assign_chain.invoke({"tag": tag})
    print(f"Tag: <{tag}>")
    print(f"Tip: {tip}\n")

# 3. With RunnableParallel — keep original dict alongside an LLM branch
parallel_chain = RunnableParallel(
    original=RunnablePassthrough(),
    purpose=explain_chain,
)

print("=== Passthrough + RunnableParallel ===")
for tag in tags:
    result = parallel_chain.invoke({"tag": tag})
    print(f"Tag: <{tag}>")
    print(f"  original: {result['original']}")
    print(f"  purpose:  {result['purpose']}\n")