"""fallbacks_retries_streaming_demo.py — resilience and streaming in LCEL
Run with venv active:  python fallbacks_retries_streaming_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()

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

# --- 1. Streaming — tokens arrive one chunk at a time ---
print("--- Streaming ---")
print("Answer: ", end="", flush=True)
for chunk in explain_chain.stream({"tag": "a"}):
    print(chunk, end="", flush=True)
print("\n")

# --- 2. Retry — LangChain retries on transient API errors ---
print("--- With retry ---")
retry_chain = explain_chain.with_retry(stop_after_attempt=3)
answer = retry_chain.invoke({"tag": "img"})
print(f"Answer: {answer}\n")

# --- 3. Fallback — backup chain runs if the primary fails ---
print("--- Fallback ---")
broken_primary = explain_prompt | ChatOpenAI(model="gpt-does-not-exist") | parser
reliable_chain = broken_primary.with_fallbacks([explain_chain])
answer = reliable_chain.invoke({"tag": "button"})
print(f"Answer: {answer}")