"""multi_llm_workflows_demo.py — different models for different steps
Run with venv active:  python multi_llm_workflows_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

load_dotenv()

parser = StrOutputParser()

# Fast model for a quick extraction step
fast_model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# Stronger model for a richer explanation step
smart_model = ChatOpenAI(model="gpt-4o", temperature=0)

# Step 1: pick tag name (fast model)
pick_prompt = ChatPromptTemplate.from_template(
    "Name one HTML tag used for {feature}. Reply with only the tag name, no brackets."
)
pick_chain = pick_prompt | fast_model | parser

# Step 2: beginner-friendly explanation (smart model)
explain_prompt = ChatPromptTemplate.from_template(
    "Explain what the HTML <{tag}> tag does for a complete beginner. "
    "Two short sentences with one practical example."
)

# Multi-LLM workflow: mini picks the tag, gpt-4o writes the explanation
chain = pick_chain | explain_prompt | smart_model | parser

features = ["links", "images"]

for feature in features:
    explanation = chain.invoke({"feature": feature})
    print(f"Feature: {feature}")
    print(f"Explanation: {explanation}\n")