# PROMPTPILOTPRO LESSON 1 – Build Your First RAG System
# No math required. Copy each cell. Run in order.

# Cell 1: Install libraries
!pip install langchain-openai langchain-community langchain-chroma langchain-text-splitters chromadb openai pypdf

# Cell 2: Add your OpenAI API key
import os
os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY_HERE"  # Replace with your actual key

# Cell 3: Load document (upload northwind_products.txt to Colab first)
from langchain_community.document_loaders import TextLoader
loader = TextLoader("/content/northwind_products.txt")
documents = loader.load()
print(f"Loaded {len(documents)} document(s)")

# Cell 4: Split document into chunks
from langchain_text_splitters import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = text_splitter.split_documents(documents)
print(f"Split into {len(chunks)} chunks")

# Cell 5: Create vector store (embed + store)
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(chunks, embeddings)
print("Vector store ready!")

# Cell 6: Retrieve AND generate an answer (full RAG)
query = "Which products are discontinued?"
docs = vectorstore.similarity_search(query)
context = "\n\n".join([d.page_content for d in docs])

from openai import OpenAI
client = OpenAI()
response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": "Answer the question based ONLY on the context provided. If the answer is not in the context, say you don't know."},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
    ]
)
print(response.choices[0].message.content)

# Cell 7: Try more questions (uncomment and edit)
# query = "What is the most expensive product?"
# query = "Who supplies the Tofu product?"
# query = "Which customers are in Germany?"