Title: Lab 1 — Instrument a Mini RAG Function with Local MLflow Tracing
Course: AIINFRA 301: Production RAG & LLMOps — Observability and Evaluation > Section: Week 01 — Introduction to Production RAG Architecture and the LLMOps Lifecycle

🛠️ Lab 1 — Instrument a Mini RAG Function with Local MLflow Tracing (50 pts)

Goal: Stand up a free, local observability backbone and use it to capture a full trace — inputs, outputs, latency, and token counts — of a minimal "retrieve → generate" RAG function. By the end of this lab you will have seen, not just read about, what a production trace looks like. Steps:

1. Set up your environment. Create a new folder for this lab and a virtual environment (python -m venv venv, then activate it). Install MLflow: pip install mlflow.
2. Launch the local MLflow tracing UI. Run mlflow ui in a terminal and confirm you can open the local URL it prints (typically http://127.0.0.1:5000) in your browser. Leave this running in its own terminal tab.
3. Build a tiny knowledge base. In a new Python file, create a Python list of 5–8 short text strings (a few sentences each) about any topic you like (e.g., your own notes from AIINFRA 300). This stands in for your "corpus" — no external vector database required.
4. Write a bare-bones retrieval function. Write a function retrieve(query, corpus) that does simple keyword overlap scoring (e.g., count shared words between the query and each corpus entry) and returns the top 1–2 matching strings. This keeps the lab free and dependency-light — no embedding API key needed.
5. Write a bare-bones generation function. Write a function generate(query, retrieved_chunks) that builds a prompt combining the query and retrieved chunks. If you have free/local access to a model (e.g., via Ollama or a free-tier API from AIINFRA 300), call it; otherwise, simulate generation with a simple string template — the point of this lab is the trace, not model quality.
6. Instrument both functions with MLflow tracing. Add the @mlflow.trace decorator (or mlflow.start_span) around retrieve and generate. Set mlflow.set_experiment("week01-rag-trace") at the top of your script. Follow the pattern in the MLflow GenAI Tracing Quickstart resource linked below.
7. Run your pipeline 3–5 times with different sample queries so you generate multiple traces to compare.
8. Inspect your traces in the MLflow UI. Open the Traces tab and confirm you can see: the full call tree (retrieve → generate), the inputs/outputs at each step, and the latency per step. Take a screenshot.
9. (Optional stretch) Compare with Langfuse. If you want to see a second free observability backbone this course will build on later, follow the Langfuse self-hosted Docker Compose guide linked below to spin up Langfuse locally, and send a couple of traces there too. Note one difference you noticed between the MLflow and Langfuse trace views.
10. Write a short reflection (150–250 words). Answer: Which of the four production RAG layers from the lecture does this lab actually exercise? What did seeing the trace teach you that reading about "observability" didn't?

Deliverables: Submit your Python script(s), a screenshot of your MLflow trace view showing at least 3 traces, and your written reflection — either as an uploaded file (.py + .png/.pdf) or pasted directly into the text entry box.