# Install required packages
!pip install langchain langchain-openai langchain-core python-dotenv -q

import warnings
warnings.filterwarnings("ignore")

# API key management
# Option 1: Google Colab Secrets (recommended)
# from google.colab import userdata
# OPENROUTER_API_KEY = userdata.get('OPENROUTER_API_KEY')

# Option 2: Direct variable (local dev only — never commit this)
OPENROUTER_API_KEY = "your-api-key-here"

from langchain_openai import ChatOpenAI

# Initialize the model via OpenRouter
llm = ChatOpenAI(
    api_key=OPENROUTER_API_KEY,
    base_url="https://openrouter.ai/api/v1",
    model="arcee-ai/trinity-large-preview:free"  # free 400B params model
)

print("Model initialized:", llm.model_name)

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

# Create messages
system_msg = SystemMessage(content="You are an ML expert. Always answer briefly.")
human_msg = HumanMessage(content="What is Dropout?")

print("SystemMessage:", system_msg)
print("HumanMessage :", human_msg)

# invoke() — waits for the complete response, returns an AIMessage
response = llm.invoke([system_msg, human_msg])

print("Returned type:", type(response))
print("Content      :", response.content)

# stream() — returns tokens in real-time, token by token
# Use chunk.content — the official native attribute (not chunk.text)
for chunk in llm.stream([system_msg, human_msg]):
    print(chunk.content, end="", flush=True)

# Wrong: instantiates first, then calls the method
ChatPromptTemplate().from_messages([...])

# Correct: calls directly on the class
ChatPromptTemplate.from_messages([...])

from langchain_core.prompts import ChatPromptTemplate

# Create a template with a dynamic variable {subject}
# Pass tuples ("role", "text") — not Message objects
template = ChatPromptTemplate.from_messages([
    ("system", "You are an ML expert. Always answer briefly."),
    ("human", "Hi, I want you to explain {subject} to me")
])

# invoke() with a variable dictionary → returns a ChatPromptValue
prompt_value = template.invoke({"subject": "Dropout"})

print("Returned type:", type(prompt_value))
print("Content      :", prompt_value)

# The ChatPromptValue is passed directly to the llm
for chunk in llm.stream(prompt_value):
    print(chunk.content, end="", flush=True)

from langchain_core.output_parsers import StrOutputParser

parser = StrOutputParser()

# invoke() — receives an AIMessage, returns a str
response = llm.invoke(prompt_value)
parsed = parser.invoke(response)

print("Type before parser:", type(response))  # AIMessage
print("Type after parser :", type(parsed))    # str
print("\nContent:", parsed)

# Works: but loses streaming (invoke waits for the complete response)
parser.invoke(llm.invoke(prompt_value))

# Does not work: StrOutputParser expects an AIMessage, not a generator
parser.stream(llm.stream(prompt_value))  # → ValidationError

chain = template | llm | parser
chain.stream({"subject": "Dropout"})  # end-to-end streaming

from pydantic import BaseModel, Field

# Define the Pydantic schema
# Field(description=...) guides the model on what to put in each field
class MLConcept(BaseModel):
    name: str = Field(..., description="The name of the ML concept")
    definition: str = Field(..., description="A clear and concise definition of the concept")
    example: str = Field(..., description="A concrete real-world application example")

# Create the structured_llm — new object, the original llm stays intact
structured_llm = llm.with_structured_output(MLConcept)

# invoke() only — no stream() with structured output
response = structured_llm.invoke(prompt_value)

print("Returned type:", type(response))  # MLConcept — no longer an AIMessage!
print()
print("name      :", response.name)
print("definition:", response.definition)
print("example   :", response.example)

chain = template | llm | parser
chain.stream({"subject": "Dropout"})  # end-to-end streaming