"""route_between_llms_demo.py — RunnableBranch picks one LLM chain
Run with venv active:  python route_between_llms_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 RunnableBranch

load_dotenv()

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

simple_prompt = ChatPromptTemplate.from_template(
    "In one short sentence, what does the HTML <{tag}> tag do?"
)
simple_chain = simple_prompt | mini | parser

detail_prompt = ChatPromptTemplate.from_template(
    "Explain what the HTML <{tag}> tag does. "
    "Name one common attribute and show a short example."
)
detail_chain = detail_prompt | pro | parser

COMPLEX_TAGS = {"form", "table", "select", "textarea"}

router = RunnableBranch(
    (lambda x: x["tag"] in COMPLEX_TAGS, detail_chain),
    simple_chain,
)

for tag in ["a", "form", "p", "table"]:
    answer = router.invoke({"tag": tag})
    print(f"tag={tag}")
    print(f"  {answer}\n")