from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

# Prompt template
prompt = ChatPromptTemplate.from_template("tell me a joke about {topic}")
# Main model
model = OllamaLLM(model="llama3.2", base_url="http://studio1.dorceus.local:11434" )
# Define fallback as a RunnableLambda
fallback_runnable = RunnableLambda(lambda x: "Oops! Something went wrong.")
# Add fallback to model
model_with_fallback = model.with_fallbacks([fallback_runnable])
# Compose full chain
chain = prompt | model_with_fallback
# Run the chain
response = chain.invoke({"topic": "bears"})
print(response)