# file: 01_basic_summarizer.py

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

OPENROUTER_API_URL = "https://openrouter.ai/api/v1"

def main():
    """
    Main function to load environment variables, create a summarization chain,
    and process a sample text, as described in Chapter 2.
    """
    # Load environment variables from .env file
    load_dotenv()

    # 1. Define the LLM we want to use
    # We select a fast and cost-effective model from the Gemini family and
    # configure the client to use OpenRouter's API.
    '''chat_model = ChatOpenAI(
        model="google/gemini-2.0-flash-001",
        temperature=0.3,
        openai_api_key=os.getenv("OPENROUTER_API_KEY"),
        openai_api_base=OPENROUTER_API_URL # Using the constant defined above
    )'''

    chat_model = ChatOpenAI(
        model="qwen/qwen3-coder:free", # Swapped from Gemini to Qwen's Coder
        temperature=0.3,
        openai_api_key=os.getenv("OPENROUTER_API_KEY"),
        openai_api_base=OPENROUTER_API_URL
    )

    # 2. Create our prompt template
    # This guides the model on what to do with the input.
    template = """
    You are a helpful assistant who summarizes text.
    Summarize the following text in 1 sentence:

    Text: "{text_to_summarize}"
    """
    prompt = ChatPromptTemplate.from_template(template)

    # 3. Create the output parser
    # This will convert the model's chat message output into a string.
    output_parser = StrOutputParser()

    # 4. Create the "chain" by piping the components together
    # This uses the LangChain Expression Language (LCEL) to define the flow of data.
    summarization_chain = prompt | chat_model | output_parser

    # 5. Invoke the chain with a sample text
    # We pass the input text from the book to test our chain.
    sample_text = """
    AI engineering represents a significant paradigm shift from traditional machine
    learning. While traditional ML focused heavily on model training and optimization
    in environments like Jupyter Notebooks, modern AI engineering treats pre-trained
    Large Language Models (LLMs) as components, akin to a CPU. The primary challenge
    is no longer just building an accurate model, but architecting a robust system
    around it. This requires a focus on three pillars: observability, to understand
    the system's internal state; reliability, to ensure predictable and structured
    outputs; and scalability, to manage cost and performance effectively.
    """
    
    print("--- Summarizing Text ---")
    summary = summarization_chain.invoke({"text_to_summarize": sample_text})
    print(summary)
    print("------------------------")

if __name__ == "__main__":
    main()