{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 3: Independent Lab — Build Your Own RAG Assistant\n",
    "\n",
    "## Overview\n",
    "\n",
    "Build a domain-specific RAG assistant by choosing one of four business tracks. You will design a knowledge base, implement chunking, write a grounding prompt, create a golden test set, and iterate from v1 → v2 with measured improvement.\n",
    "\n",
    "### Deliverables\n",
    "| # | Item | Requirement |\n",
    "|---|------|-------------|\n",
    "| 1 | Track selection | Clearly marked |\n",
    "| 2 | Knowledge base | 4+ documents |\n",
    "| 3 | Chunking strategy | Documented rationale |\n",
    "| 4 | System prompt v1 | Grounding rules |\n",
    "| 5 | Golden test set | 10+ questions |\n",
    "| 6 | v1 evaluation | Accuracy + scores |\n",
    "| 7 | System prompt v2 | Improved rules |\n",
    "| 8 | v2 evaluation | Comparison |\n",
    "| 9 | Error analysis | Patterns + insights |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U google-genai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Imports ──────────────────────────────────────────────\n",
    "import os, time, json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from datetime import datetime, timezone\n",
    "from typing import List, Optional, Literal\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "from google import genai\n",
    "from google.genai import types\n",
    "\n",
    "# ── API Key ───────────────────────────────────────────────\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    API_KEY = userdata.get(\"GEMINI_API_KEY\")\n",
    "except Exception:\n",
    "    API_KEY = None\n",
    "\n",
    "if not API_KEY:\n",
    "    import getpass\n",
    "    API_KEY = getpass.getpass(\"Enter your Gemini API key: \")\n",
    "\n",
    "client = genai.Client(api_key=API_KEY)\n",
    "\n",
    "MODEL_ID = \"gemini-2.5-flash-lite\"\n",
    "EMBEDDING_MODEL = \"gemini-embedding-001\"\n",
    "\n",
    "print(f\"API key loaded: {'yes' if API_KEY else 'no'}\")\n",
    "print(f\"Generation model: {MODEL_ID}\")\n",
    "print(f\"Embedding model:  {EMBEDDING_MODEL}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Infrastructure (same as Guided Lab) ──────────────────\n",
    "PROMPT_LOG = []\n",
    "\n",
    "def _now():\n",
    "    return datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z')\n",
    "\n",
    "def generate(prompt, temperature=0.7, max_tokens=1000, log=True, label=None):\n",
    "    \"\"\"Generate free-form text. Returns raw string.\"\"\"\n",
    "    t0 = time.time()\n",
    "    response = client.models.generate_content(\n",
    "        model=MODEL_ID,\n",
    "        contents=prompt,\n",
    "        config={\"temperature\": temperature, \"max_output_tokens\": max_tokens},\n",
    "    )\n",
    "    latency = time.time() - t0\n",
    "    text = response.text\n",
    "    if log:\n",
    "        PROMPT_LOG.append({\n",
    "            \"timestamp\": _now(), \"label\": label or \"generate\",\n",
    "            \"type\": \"free_form\",\n",
    "            \"prompt\": prompt[:300] + \"...\" if len(prompt) > 300 else prompt,\n",
    "            \"prompt_length\": len(prompt),\n",
    "            \"temperature\": temperature,\n",
    "            \"response\": text[:300] + \"...\" if len(text) > 300 else text,\n",
    "            \"response_length\": len(text),\n",
    "            \"latency_s\": round(latency, 2),\n",
    "        })\n",
    "    return text\n",
    "\n",
    "def generate_structured(prompt, schema_model, temperature=0.2, log=True, label=None):\n",
    "    \"\"\"Generate structured JSON output. Returns Pydantic model instance.\"\"\"\n",
    "    t0 = time.time()\n",
    "    response = client.models.generate_content(\n",
    "        model=MODEL_ID,\n",
    "        contents=prompt,\n",
    "        config={\n",
    "            \"temperature\": temperature,\n",
    "            \"response_mime_type\": \"application/json\",\n",
    "            \"response_schema\": schema_model,\n",
    "        },\n",
    "    )\n",
    "    latency = time.time() - t0\n",
    "    result = schema_model.model_validate_json(response.text)\n",
    "    if log:\n",
    "        PROMPT_LOG.append({\n",
    "            \"timestamp\": _now(), \"label\": label or \"structured\",\n",
    "            \"type\": \"structured\",\n",
    "            \"schema\": schema_model.__name__,\n",
    "            \"prompt\": prompt[:300] + \"...\" if len(prompt) > 300 else prompt,\n",
    "            \"prompt_length\": len(prompt),\n",
    "            \"temperature\": temperature,\n",
    "            \"response\": response.text[:300],\n",
    "            \"response_length\": len(response.text),\n",
    "            \"latency_s\": round(latency, 2),\n",
    "        })\n",
    "    return result\n",
    "\n",
    "def show_log():\n",
    "    \"\"\"Display prompt log as DataFrame.\"\"\"\n",
    "    if not PROMPT_LOG:\n",
    "        print(\"No API calls logged yet.\")\n",
    "        return None\n",
    "    return pd.DataFrame(PROMPT_LOG)\n",
    "\n",
    "def embed_texts(texts, task_type=\"RETRIEVAL_DOCUMENT\"):\n",
    "    \"\"\"Embed one or more texts using Gemini Embeddings API.\"\"\"\n",
    "    if isinstance(texts, str):\n",
    "        texts = [texts]\n",
    "    response = client.models.embed_content(\n",
    "        model=EMBEDDING_MODEL,\n",
    "        contents=texts,\n",
    "        config=types.EmbedContentConfig(task_type=task_type),\n",
    "    )\n",
    "    return [np.array(e.values) for e in response.embeddings]\n",
    "\n",
    "def cosine_similarity(a, b):\n",
    "    \"\"\"Compute cosine similarity between two vectors.\"\"\"\n",
    "    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))\n",
    "\n",
    "def search(query, doc_embeddings, documents, top_k=3):\n",
    "    \"\"\"Find top-k most similar documents to a query.\"\"\"\n",
    "    query_emb = embed_texts(query, task_type=\"RETRIEVAL_QUERY\")[0]\n",
    "    scores = []\n",
    "    for i, doc_emb in enumerate(doc_embeddings):\n",
    "        sim = cosine_similarity(query_emb, doc_emb)\n",
    "        scores.append((i, sim, documents[i]))\n",
    "    scores.sort(key=lambda x: x[1], reverse=True)\n",
    "    return scores[:top_k]\n",
    "\n",
    "def rag_query(question, documents, doc_embeddings, top_k=3, system_prompt=None):\n",
    "    \"\"\"Complete RAG pipeline: retrieve → augment → generate.\"\"\"\n",
    "    results = search(question, doc_embeddings, documents, top_k)\n",
    "    context = \"\\n\\n\".join(\n",
    "        f\"[Chunk {i+1}]\\n{doc}\" for i, (_, _, doc) in enumerate(results)\n",
    "    )\n",
    "    if system_prompt is None:\n",
    "        system_prompt = (\n",
    "            \"You are a helpful assistant. Answer based ONLY on the provided context. \"\n",
    "            \"If the answer is not in the context, say \\\"I don't have that information.\\\"\"\n",
    "        )\n",
    "    prompt = f\"\"\"{system_prompt}\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "\n",
    "Question: {question}\"\"\"\n",
    "    answer = generate(prompt, temperature=0.3, label=f\"rag_{question[:30]}\")\n",
    "    return answer, results\n",
    "\n",
    "def chunk_sentences(text, max_words=60, overlap_words=15):\n",
    "    \"\"\"Split text at sentence boundaries with overlap.\"\"\"\n",
    "    sentences = [s.strip() for s in text.replace('\\n', ' ').split('. ') if s.strip()]\n",
    "    chunks = []\n",
    "    current = []\n",
    "    current_len = 0\n",
    "    for sent in sentences:\n",
    "        sent_words = len(sent.split())\n",
    "        if current_len + sent_words > max_words and current:\n",
    "            chunks.append(\". \".join(current) + \".\")\n",
    "            overlap_sents = []\n",
    "            overlap_len = 0\n",
    "            for s in reversed(current):\n",
    "                if overlap_len + len(s.split()) <= overlap_words:\n",
    "                    overlap_sents.insert(0, s)\n",
    "                    overlap_len += len(s.split())\n",
    "                else:\n",
    "                    break\n",
    "            current = overlap_sents\n",
    "            current_len = overlap_len\n",
    "        current.append(sent)\n",
    "        current_len += sent_words\n",
    "    if current:\n",
    "        chunks.append(\". \".join(current) + \".\")\n",
    "    return chunks\n",
    "\n",
    "# ── Additional helpers for this lab ──────────────────────\n",
    "\n",
    "def evaluate_retrieval(golden_set, chunks, chunk_embeddings, system_prompt, top_k=3):\n",
    "    \"\"\"Run RAG on a golden set and measure keyword accuracy.\"\"\"\n",
    "    results = []\n",
    "    for qa in golden_set:\n",
    "        answer, retrieved = rag_query(\n",
    "            qa[\"question\"], chunks, chunk_embeddings,\n",
    "            top_k=top_k, system_prompt=system_prompt\n",
    "        )\n",
    "        answer_lower = answer.lower()\n",
    "        keyword_hit = any(kw.lower() in answer_lower for kw in qa[\"expected_keywords\"])\n",
    "        results.append({\n",
    "            \"id\": qa[\"id\"],\n",
    "            \"question\": qa[\"question\"],\n",
    "            \"difficulty\": qa[\"difficulty\"],\n",
    "            \"answer\": answer,\n",
    "            \"keyword_match\": keyword_hit,\n",
    "            \"top_score\": retrieved[0][1],\n",
    "        })\n",
    "    return pd.DataFrame(results)\n",
    "\n",
    "def compare_versions(v1_df, v2_df):\n",
    "    \"\"\"Compare two evaluation DataFrames.\"\"\"\n",
    "    v1_acc = v1_df[\"keyword_match\"].mean()\n",
    "    v2_acc = v2_df[\"keyword_match\"].mean()\n",
    "    print(f\"v1 accuracy: {v1_acc:.0%} ({v1_df['keyword_match'].sum()}/{len(v1_df)})\")\n",
    "    print(f\"v2 accuracy: {v2_acc:.0%} ({v2_df['keyword_match'].sum()}/{len(v2_df)})\")\n",
    "    print(f\"Improvement: {v2_acc - v1_acc:+.0%}\")\n",
    "    # Show per-question comparison\n",
    "    comparison = v1_df[[\"id\", \"difficulty\", \"keyword_match\"]].rename(\n",
    "        columns={\"keyword_match\": \"v1_correct\"}\n",
    "    )\n",
    "    comparison[\"v2_correct\"] = v2_df[\"keyword_match\"].values\n",
    "    comparison[\"changed\"] = comparison[\"v1_correct\"] != comparison[\"v2_correct\"]\n",
    "    print(\"\\nPer-question comparison:\")\n",
    "    print(comparison.to_string(index=False))\n",
    "    return comparison\n",
    "\n",
    "# ── Hybrid Retrieval (TF-IDF + Semantic) ─────────────────\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "\n",
    "def build_sparse_index(chunks):\n",
    "    \"\"\"Build a TF-IDF sparse index for keyword retrieval.\"\"\"\n",
    "    vectorizer = TfidfVectorizer(stop_words='english')\n",
    "    matrix = vectorizer.fit_transform(chunks)\n",
    "    return vectorizer, matrix\n",
    "\n",
    "def search_sparse(query, vectorizer, tfidf_matrix, chunks, top_k=5):\n",
    "    \"\"\"Sparse keyword search using TF-IDF cosine similarity.\"\"\"\n",
    "    from sklearn.metrics.pairwise import cosine_similarity as sklearn_cosine\n",
    "    query_vec = vectorizer.transform([query])\n",
    "    sims = sklearn_cosine(query_vec, tfidf_matrix).ravel()\n",
    "    top_indices = np.argsort(sims)[::-1][:top_k]\n",
    "    return [(int(i), float(sims[i]), chunks[i]) for i in top_indices if sims[i] > 0]\n",
    "\n",
    "def hybrid_search(query, doc_embeddings, chunks, vectorizer, tfidf_matrix,\n",
    "                  top_k=5, semantic_weight=0.6, keyword_weight=0.4):\n",
    "    \"\"\"\n",
    "    Combine semantic and keyword search results.\n",
    "    Merges by chunk index, deduplicates, and scores with weighted average.\n",
    "    \"\"\"\n",
    "    # Semantic search\n",
    "    sem_results = search(query, doc_embeddings, chunks, top_k=top_k * 2)\n",
    "    # Sparse search\n",
    "    sparse_results = search_sparse(query, vectorizer, tfidf_matrix, chunks, top_k=top_k * 2)\n",
    "    \n",
    "    # Merge by chunk index\n",
    "    scores = {}\n",
    "    for idx, score, text in sem_results:\n",
    "        scores[idx] = {\"semantic\": score, \"sparse\": 0.0, \"text\": text}\n",
    "    for idx, score, text in sparse_results:\n",
    "        if idx in scores:\n",
    "            scores[idx][\"sparse\"] = score\n",
    "        else:\n",
    "            scores[idx] = {\"semantic\": 0.0, \"sparse\": score, \"text\": text}\n",
    "    \n",
    "    # Weighted combination\n",
    "    combined = []\n",
    "    for idx, s in scores.items():\n",
    "        combo = semantic_weight * s[\"semantic\"] + keyword_weight * s[\"sparse\"]\n",
    "        combined.append((idx, combo, s[\"semantic\"], s[\"sparse\"], s[\"text\"]))\n",
    "    \n",
    "    combined.sort(key=lambda x: x[1], reverse=True)\n",
    "    return combined[:top_k]\n",
    "\n",
    "def search_with_filter(query, doc_embeddings, chunks, chunk_metadata, top_k=3, filter_field=None, filter_value=None):\n",
    "    \"\"\"Semantic search with optional metadata filtering.\"\"\"\n",
    "    query_emb = embed_texts(query, task_type=\"RETRIEVAL_QUERY\")[0]\n",
    "    scores = []\n",
    "    for i, doc_emb in enumerate(doc_embeddings):\n",
    "        # Apply metadata filter\n",
    "        if filter_field and filter_value:\n",
    "            if chunk_metadata[i].get(filter_field) != filter_value:\n",
    "                continue\n",
    "        sim = cosine_similarity(query_emb, doc_emb)\n",
    "        scores.append((i, sim, chunks[i]))\n",
    "    scores.sort(key=lambda x: x[1], reverse=True)\n",
    "    return scores[:top_k]\n",
    "\n",
    "# ── LLM Re-ranking (optional) ────────────────────────────\n",
    "\n",
    "RERANK_PROMPT = \"\"\"You are re-ranking retrieved excerpts for a question.\n",
    "Return the {k} most relevant excerpt numbers as a JSON list.\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Excerpts:\n",
    "{excerpts}\n",
    "\n",
    "Return ONLY a JSON list of excerpt numbers like: [2, 0, 3]\n",
    "The numbers refer to the excerpt indices shown above.\"\"\"\n",
    "\n",
    "def rerank_with_llm(question, candidates, k=3):\n",
    "    \"\"\"Use LLM to re-rank candidates by relevance. Returns re-ordered list.\"\"\"\n",
    "    excerpts = \"\\n\\n\".join(\n",
    "        f\"[Excerpt {i}]: {text[:200]}...\"\n",
    "        for i, (_, _, text) in enumerate(candidates)\n",
    "    )\n",
    "    prompt = RERANK_PROMPT.format(question=question, excerpts=excerpts, k=k)\n",
    "    response = generate(prompt, temperature=0.1, label=\"rerank\")\n",
    "    try:\n",
    "        import json as _json\n",
    "        indices = _json.loads(response.strip())\n",
    "        reranked = [candidates[i] for i in indices if i < len(candidates)]\n",
    "        return reranked[:k]\n",
    "    except Exception:\n",
    "        return candidates[:k]  # Fallback to original order\n",
    "\n",
    "print(\"All infrastructure ready (includes hybrid retrieval, metadata filtering, and LLM re-ranking).\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 1: Choose Your Track\n",
    "\n",
    "Select one of four tracks. Set `SELECTED_TRACK` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── TODO: Set your track ─────────────────────────────────\n",
    "SELECTED_TRACK = \"\"  # Set to \"A\", \"B\", \"C\", or \"D\"\n",
    "assert SELECTED_TRACK in (\"A\", \"B\", \"C\", \"D\"), \"Please set SELECTED_TRACK to A, B, C, or D\"\n",
    "print(f\"Selected track: {SELECTED_TRACK}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Track Data ───────────────────────────────────────────\n",
    "\n",
    "TRACK_DATA = {\n",
    "    \"A\": {\n",
    "        \"name\": \"HR Policy Assistant\",\n",
    "        \"documents\": [\n",
    "            # Doc 0: Remote Work Policy\n",
    "            \"\"\"Remote Work Policy (Effective January 2025)\n",
    "\n",
    "Eligibility: All full-time employees who have completed their 90-day probation \n",
    "period are eligible for remote work. Contractors and interns must work on-site.\n",
    "\n",
    "Schedule: Employees may work remotely up to 3 days per week. Core collaboration \n",
    "days are Tuesday and Thursday — attendance is mandatory and in-person. Managers \n",
    "may require additional on-site days during project-critical phases.\n",
    "\n",
    "Equipment: The company provides a one-time EUR 500 stipend for home office \n",
    "equipment (desk, chair, monitor). Laptops are issued by IT. Internet costs are \n",
    "reimbursed up to EUR 30/month with receipt submission.\n",
    "\n",
    "Performance: Remote workers are evaluated on output, not hours. However, \n",
    "employees must be available during core hours (10:00-15:00 CET) for meetings \n",
    "and collaboration.\"\"\",\n",
    "\n",
    "            # Doc 1: Leave Policy\n",
    "            \"\"\"Annual Leave and Absence Policy (Effective January 2025)\n",
    "\n",
    "Annual Leave: Full-time employees receive 30 days of paid annual leave per year, \n",
    "plus public holidays. Part-time employees receive leave pro-rata.\n",
    "\n",
    "Carry-Over: Up to 5 unused days may be carried into the next calendar year. \n",
    "Carried-over days must be used by March 31 or they expire.\n",
    "\n",
    "Sick Leave: Employees receive continued pay for up to 6 weeks of illness per \n",
    "year. A doctor's note is required from the third consecutive sick day.\n",
    "\n",
    "Parental Leave: Birth parents receive 14 weeks of fully paid parental leave. \n",
    "Non-birth parents receive 4 weeks of fully paid leave. Additional unpaid leave \n",
    "of up to 6 months is available upon request.\n",
    "\n",
    "Special Leave: 2 days for marriage, 2 days for bereavement of immediate family, \n",
    "1 day for moving house.\"\"\",\n",
    "\n",
    "            # Doc 2: Benefits Overview\n",
    "            \"\"\"Employee Benefits Overview (2025)\n",
    "\n",
    "Health Insurance: The company covers 50% of statutory health insurance \n",
    "contributions. Supplementary private dental insurance is available at group \n",
    "rates through our partner Allianz.\n",
    "\n",
    "Pension: The company matches employee pension contributions up to 4% of gross \n",
    "salary through our company pension scheme (bAV).\n",
    "\n",
    "Professional Development: Each employee has an annual learning budget of \n",
    "EUR 1,500 for courses, conferences, and certifications. Unused budget does \n",
    "not carry over.\n",
    "\n",
    "Wellness: EUR 50/month gym or fitness subsidy (receipt required). Annual \n",
    "company health check-up offered in Q2.\n",
    "\n",
    "Other Benefits: Public transport subsidy (Deutschlandticket), free drinks and \n",
    "snacks in the office, quarterly team events with EUR 50/person budget.\"\"\",\n",
    "\n",
    "            # Doc 3: Work Hours Policy\n",
    "            \"\"\"Working Hours Policy (Effective January 2025)\n",
    "\n",
    "Standard Hours: Full-time employees work 40 hours per week. Core hours are \n",
    "10:00 to 15:00 CET — all employees must be available during this window.\n",
    "\n",
    "Flextime: Outside of core hours, employees may arrange their schedule \n",
    "flexibly. Start times between 07:00 and 10:00 are permitted.\n",
    "\n",
    "Overtime: Overtime must be pre-approved by the direct manager. Approved \n",
    "overtime is compensated as time-off-in-lieu (TOIL) at a 1:1 ratio. \n",
    "Overtime exceeding 10 hours per month requires HR approval.\n",
    "\n",
    "On-Call: Engineers on the on-call rotation receive EUR 200/week flat \n",
    "compensation plus EUR 50 per incident responded to outside business hours.\n",
    "\n",
    "Time Tracking: All employees must log hours in the Personio system by \n",
    "end of each Friday. Failure to log for 2+ consecutive weeks will trigger \n",
    "a reminder from HR.\"\"\",\n",
    "        ],\n",
    "    },\n",
    "    \"B\": {\n",
    "        \"name\": \"Product Documentation\",\n",
    "        \"documents\": [\n",
    "            # Doc 0: Authentication Guide\n",
    "            \"\"\"DataSync Pro — Authentication Guide (v3.2)\n",
    "\n",
    "API Keys: Every API request must include a valid API key in the \n",
    "X-API-Key header. Keys are generated in the Dashboard under Settings > \n",
    "API Keys. Each organisation can have up to 10 active keys.\n",
    "\n",
    "OAuth 2.0: For user-level access, DataSync Pro supports the Authorization \n",
    "Code flow. Redirect URI must be registered in advance. Access tokens expire \n",
    "after 1 hour; refresh tokens are valid for 30 days.\n",
    "\n",
    "Token Refresh: POST to /oauth/token with grant_type=refresh_token and your \n",
    "refresh_token. The response includes a new access_token and a new refresh_token \n",
    "(rotate both).\n",
    "\n",
    "Rate Limits: API key requests are limited to 100 requests/minute and \n",
    "10,000 requests/day. OAuth tokens are limited to 200 requests/minute. \n",
    "Exceeding limits returns HTTP 429 with a Retry-After header.\"\"\",\n",
    "\n",
    "            # Doc 1: REST API Reference\n",
    "            \"\"\"DataSync Pro — REST API Reference (v3.2)\n",
    "\n",
    "Base URL: https://api.datasyncpro.com/v3\n",
    "\n",
    "GET /connections — List all active connections. Returns array of connection \n",
    "objects. Supports ?status=active|paused|error filtering.\n",
    "\n",
    "POST /connections — Create a new connection. Required body: source_type, \n",
    "destination_type, schedule. Returns connection_id.\n",
    "\n",
    "GET /connections/{id}/runs — List sync runs for a connection. Returns \n",
    "timestamps, row counts, and status. Supports ?limit and ?offset pagination.\n",
    "\n",
    "POST /sync/trigger — Manually trigger a sync. Required: connection_id. \n",
    "Optional: full_sync (boolean, default false for incremental).\n",
    "\n",
    "GET /health — System health check. Returns status, version, and uptime. \n",
    "No authentication required.\n",
    "\n",
    "Error Responses: 400 (bad request), 401 (invalid key), 403 (insufficient \n",
    "permissions), 404 (resource not found), 429 (rate limited), 500 (server error).\"\"\",\n",
    "\n",
    "            # Doc 2: SDK Quickstart\n",
    "            \"\"\"DataSync Pro — SDK Quickstart Guide\n",
    "\n",
    "Python SDK:\n",
    "  pip install datasyncpro\n",
    "  from datasyncpro import Client\n",
    "  client = Client(api_key=\"your-key\")\n",
    "  connections = client.connections.list()\n",
    "\n",
    "JavaScript SDK:\n",
    "  npm install @datasyncpro/sdk\n",
    "  import { DataSyncClient } from '@datasyncpro/sdk';\n",
    "  const client = new DataSyncClient({ apiKey: 'your-key' });\n",
    "  const connections = await client.connections.list();\n",
    "\n",
    "Java SDK:\n",
    "  Add Maven dependency: com.datasyncpro:sdk:3.2.0\n",
    "  DataSyncClient client = new DataSyncClient(\"your-key\");\n",
    "  List<Connection> connections = client.connections().list();\n",
    "\n",
    "All SDKs support automatic retry with exponential backoff for 429 and 5xx \n",
    "errors. Default: 3 retries with 1s/2s/4s delays. Configure via \n",
    "client.config.maxRetries.\"\"\",\n",
    "\n",
    "            # Doc 3: Common Errors\n",
    "            \"\"\"DataSync Pro — Troubleshooting Common Errors\n",
    "\n",
    "HTTP 401 Unauthorized:\n",
    "- API key is missing, expired, or revoked. Regenerate in Dashboard.\n",
    "- OAuth token has expired. Use the refresh token to obtain a new one.\n",
    "\n",
    "HTTP 403 Forbidden:\n",
    "- Your API key lacks permission for this endpoint. Check key scopes in Dashboard.\n",
    "- Organisation-level endpoints require admin role.\n",
    "\n",
    "HTTP 429 Too Many Requests:\n",
    "- You've exceeded rate limits. Check the Retry-After header and wait.\n",
    "- Implement exponential backoff: wait 1s, 2s, 4s between retries.\n",
    "- Consider batching requests to reduce call volume.\n",
    "\n",
    "HTTP 500 Internal Server Error:\n",
    "- Temporary server issue. Retry after 30 seconds.\n",
    "- If persistent (>5 minutes), check status.datasyncpro.com.\n",
    "- Contact support@datasyncpro.com with your request_id from the response header.\n",
    "\n",
    "Sync Failures:\n",
    "- \"Schema mismatch\": Source schema changed. Re-map fields in Dashboard.\n",
    "- \"Connection timeout\": Destination unreachable. Check firewall rules.\n",
    "- \"Row limit exceeded\": Free plan limited to 100K rows per sync.\"\"\",\n",
    "        ],\n",
    "    },\n",
    "    \"C\": {\n",
    "        \"name\": \"Customer Support\",\n",
    "        \"documents\": [\n",
    "            # Doc 0: Account Access\n",
    "            \"\"\"CloudBase — Account Access Guide\n",
    "\n",
    "Password Reset: Click \"Forgot Password\" on the login page. A reset link \n",
    "is sent to your registered email (valid for 24 hours). If you don't receive \n",
    "it, check spam/junk folders or contact support.\n",
    "\n",
    "Two-Factor Authentication (2FA): Enable 2FA in Settings > Security. We \n",
    "support authenticator apps (Google Authenticator, Authy) and SMS. Recovery \n",
    "codes are provided during setup — store them securely.\n",
    "\n",
    "Locked Accounts: After 5 failed login attempts, accounts are locked for \n",
    "30 minutes. Contact support to unlock immediately if urgent.\n",
    "\n",
    "Single Sign-On (SSO): Enterprise plans support SAML 2.0 SSO. Configuration \n",
    "requires: Identity Provider metadata URL, attribute mapping, and admin \n",
    "approval. Setup takes 1-2 business days.\"\"\",\n",
    "\n",
    "            # Doc 1: Billing Guide\n",
    "            \"\"\"CloudBase — Billing & Subscription Guide\n",
    "\n",
    "Plans: Free (5 users, 1GB), Team (EUR 12/user/month, 50GB), Enterprise \n",
    "(custom pricing, unlimited storage). All prices exclude VAT.\n",
    "\n",
    "Upgrades: Upgrades take effect immediately. You are charged a prorated \n",
    "amount for the remainder of the billing cycle.\n",
    "\n",
    "Downgrades: Downgrades take effect at the start of the next billing cycle. \n",
    "No partial refunds are issued for the current cycle.\n",
    "\n",
    "Refund Policy: Full refunds are available within 14 days of initial \n",
    "purchase or upgrade. After 14 days, no refunds are issued. To request a \n",
    "refund, email billing@cloudbase.io with your account ID.\n",
    "\n",
    "Payment Methods: Credit card (Visa, Mastercard, Amex), SEPA direct debit \n",
    "(EU only), bank transfer (annual plans only, Enterprise tier).\n",
    "\n",
    "Invoices: Monthly invoices are sent on the 1st of each month to the billing \n",
    "email. Past invoices are available in Settings > Billing > Invoice History.\"\"\",\n",
    "\n",
    "            # Doc 2: Troubleshooting\n",
    "            \"\"\"CloudBase — Troubleshooting Guide\n",
    "\n",
    "Slow Performance:\n",
    "- Clear browser cache and cookies. Try incognito/private mode.\n",
    "- Check status.cloudbase.io for ongoing incidents.\n",
    "- Large dashboards (>50 widgets) may load slowly. Consider splitting them.\n",
    "\n",
    "Sync Errors:\n",
    "- \"Sync conflict\": Two users edited the same record. Manually resolve by \n",
    "  choosing the correct version in the conflict resolution panel.\n",
    "- \"Connection lost\": Check internet connectivity. CloudBase auto-retries \n",
    "  every 30 seconds for up to 5 minutes.\n",
    "\n",
    "Browser Compatibility:\n",
    "- Supported: Chrome 90+, Firefox 88+, Edge 90+, Safari 15+.\n",
    "- NOT supported: Internet Explorer (any version).\n",
    "- Mobile: Use the CloudBase mobile app for iOS 15+ and Android 12+.\n",
    "\n",
    "Export Issues:\n",
    "- CSV exports are limited to 100,000 rows. Use the API for larger datasets.\n",
    "- PDF reports timeout after 60 seconds. Reduce the date range or number of \n",
    "  charts to speed up generation.\"\"\",\n",
    "\n",
    "            # Doc 3: Usage Limits\n",
    "            \"\"\"CloudBase — Usage Limits & Fair Use Policy\n",
    "\n",
    "Storage Quotas:\n",
    "- Free: 1 GB total. Team: 50 GB total (shared across team). \n",
    "  Enterprise: Unlimited (fair use applies).\n",
    "- File upload limit: 100 MB per file on all plans.\n",
    "\n",
    "API Limits:\n",
    "- Free: 1,000 API calls/day. Team: 50,000 API calls/day. \n",
    "  Enterprise: 500,000 API calls/day (higher limits available on request).\n",
    "- Rate limit: 60 requests/minute per API key.\n",
    "\n",
    "Team Size:\n",
    "- Free: Up to 5 users. Team: Up to 100 users. Enterprise: Unlimited.\n",
    "- Adding users beyond plan limits requires an upgrade.\n",
    "\n",
    "Fair Use Policy: \"Unlimited\" features (Enterprise storage, API calls) are \n",
    "subject to fair use. Sustained usage exceeding 10x the median for your plan \n",
    "tier may trigger a review. We will contact you before taking any action.\n",
    "\n",
    "Data Retention: Deleted data is retained in backups for 30 days (Team) or \n",
    "90 days (Enterprise). Free plan data is not backed up.\"\"\",\n",
    "        ],\n",
    "    },\n",
    "    \"D\": {\n",
    "        \"name\": \"Research Analyst\",\n",
    "        \"documents\": [\n",
    "            # Doc 0: Market Overview\n",
    "            \"\"\"Cloud Analytics Market — Overview (Q4 2024 Report)\n",
    "\n",
    "Total Addressable Market (TAM): The global cloud analytics market was valued \n",
    "at USD 65 billion in 2024 and is projected to reach USD 130 billion by 2028, \n",
    "growing at a CAGR of 19%.\n",
    "\n",
    "Key Segments: Business Intelligence (35% of market), Data Integration (25%), \n",
    "Advanced Analytics & ML (25%), Data Governance (15%).\n",
    "\n",
    "Growth Drivers: Increasing cloud adoption (78% of enterprises now multi-cloud), \n",
    "rising data volumes (estimated 120 ZB generated globally in 2024), and \n",
    "regulatory compliance requirements (GDPR, CCPA, AI Act).\n",
    "\n",
    "Market Maturity: The BI segment is mature with consolidation expected. \n",
    "Advanced Analytics & ML is the fastest-growing segment at 28% CAGR, driven \n",
    "by generative AI adoption. Data Governance is emerging as organisations \n",
    "prepare for the EU AI Act (effective August 2025).\"\"\",\n",
    "\n",
    "            # Doc 1: Competitive Landscape\n",
    "            \"\"\"Cloud Analytics Market — Competitive Landscape (Q4 2024)\n",
    "\n",
    "Top 5 Players by Market Share:\n",
    "1. Snowflake (18%): Dominant in data warehousing. Strengths: performance, \n",
    "   ecosystem. Weakness: premium pricing, vendor lock-in concerns.\n",
    "2. Databricks (15%): Leader in ML/AI workloads. Strengths: open-source \n",
    "   (Apache Spark), unified analytics. Weakness: complexity for non-technical users.\n",
    "3. Microsoft Fabric (14%): Fastest-growing. Strengths: Office 365 integration, \n",
    "   enterprise relationships. Weakness: relatively new, fewer third-party integrations.\n",
    "4. Google BigQuery (12%): Strong in serverless analytics. Strengths: pricing \n",
    "   model, Gemini AI integration. Weakness: smaller partner ecosystem.\n",
    "5. AWS Redshift (11%): Established player. Strengths: AWS ecosystem, \n",
    "   broad feature set. Weakness: aging architecture, migration complexity.\n",
    "\n",
    "Rest of Market (30%): Includes Tableau (Salesforce), Qlik, Palantir, dbt Labs, \n",
    "and 200+ niche players. Consolidation is accelerating — 23 acquisitions in 2024.\"\"\",\n",
    "\n",
    "            # Doc 2: Customer Trends\n",
    "            \"\"\"Cloud Analytics Market — Customer Trends (Q4 2024)\n",
    "\n",
    "Adoption Patterns: 67% of enterprises now use 2+ analytics platforms \n",
    "(up from 45% in 2022). Multi-tool strategies are driven by best-of-breed \n",
    "preferences and avoiding vendor lock-in.\n",
    "\n",
    "Churn Drivers (in order of impact):\n",
    "1. Total cost of ownership surprises (cited by 34% of churned customers)\n",
    "2. Poor data integration with existing tools (28%)\n",
    "3. Insufficient self-service capabilities for business users (22%)\n",
    "4. Performance issues at scale (16%)\n",
    "\n",
    "NPS Benchmarks: Industry average NPS is 32. Leaders: Databricks (52), \n",
    "Snowflake (45). Laggards: AWS Redshift (18), legacy on-premise tools (5).\n",
    "\n",
    "Decision Criteria (ranked by importance): 1) Total cost of ownership, \n",
    "2) Ease of integration, 3) Performance at scale, 4) AI/ML capabilities, \n",
    "5) Vendor support quality.\n",
    "\n",
    "Buying Process: Average sales cycle is 4.5 months for enterprise deals. \n",
    "Technical evaluation involves 3-4 stakeholders. CFO sign-off required \n",
    "for deals exceeding USD 100K ARR.\"\"\",\n",
    "\n",
    "            # Doc 3: Regional Insights\n",
    "            \"\"\"Cloud Analytics Market — Regional Insights (Q4 2024)\n",
    "\n",
    "North America (45% of global market): Most mature market. Cloud-first \n",
    "mandates are standard in 80% of enterprises. Average spend: USD 2.1M per \n",
    "enterprise. Key trend: Consolidation of analytics tools to reduce costs.\n",
    "\n",
    "Europe (30%): Growing at 22% CAGR, faster than global average. GDPR has \n",
    "made data governance a top priority. EU AI Act (effective August 2025) is \n",
    "driving demand for explainability and audit tools. Average spend: EUR 1.4M.\n",
    "Germany and UK are largest markets (combined 55% of European spend).\n",
    "\n",
    "Asia-Pacific (20%): Fastest-growing region at 26% CAGR. India and Southeast \n",
    "Asia are key growth markets. Government digitalisation programs are a major \n",
    "driver. Average spend is lower (USD 0.6M) but growing rapidly.\n",
    "\n",
    "Rest of World (5%): Nascent markets in Middle East (UAE, Saudi Arabia) and \n",
    "Latin America (Brazil, Mexico). Primarily adopting SaaS models. Cloud \n",
    "infrastructure availability remains a constraint in some regions.\n",
    "\n",
    "Pricing Differences: Enterprise pricing is 15-20% lower in APAC and \n",
    "emerging markets compared to North America. EU pricing is comparable \n",
    "to North America.\"\"\",\n",
    "        ],\n",
    "    },\n",
    "}\n",
    "\n",
    "# Select the correct track data\n",
    "track = TRACK_DATA[SELECTED_TRACK]\n",
    "documents = track[\"documents\"]\n",
    "print(f\"Track: {track['name']}\")\n",
    "print(f\"Documents loaded: {len(documents)}\")\n",
    "for i, doc in enumerate(documents):\n",
    "    print(f\"  Doc {i}: {doc.split(chr(10))[0][:60]}... ({len(doc.split())} words)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 2: Design Your Chunking Strategy\n",
    "\n",
    "Choose your chunk size and overlap. Start with the defaults and adjust after evaluating."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── TODO: Configure your chunking parameters ─────────────\n",
    "CHUNK_MAX_WORDS = 80    # Adjust: try 60-120\n",
    "CHUNK_OVERLAP_WORDS = 20  # Adjust: try 10-25\n",
    "\n",
    "# Chunk all documents WITH metadata\n",
    "all_chunks = []\n",
    "chunk_metadata = []  # Parallel list of metadata dicts\n",
    "\n",
    "for doc_idx, doc in enumerate(documents):\n",
    "    chunks = chunk_sentences(doc, max_words=CHUNK_MAX_WORDS, overlap_words=CHUNK_OVERLAP_WORDS)\n",
    "    doc_title = doc.strip().split('\\n')[0]  # First line as title\n",
    "    for chunk_idx, chunk_text in enumerate(chunks):\n",
    "        all_chunks.append(chunk_text)\n",
    "        chunk_metadata.append({\n",
    "            \"chunk_id\": f\"DOC{doc_idx}::C{chunk_idx+1}\",\n",
    "            \"doc_index\": doc_idx,\n",
    "            \"doc_title\": doc_title,\n",
    "        })\n",
    "\n",
    "print(f\"Chunking: max_words={CHUNK_MAX_WORDS}, overlap={CHUNK_OVERLAP_WORDS}\")\n",
    "print(f\"Total chunks: {len(all_chunks)}\\n\")\n",
    "for i, (chunk, meta) in enumerate(zip(all_chunks, chunk_metadata)):\n",
    "    print(f\"Chunk {i:2d} [{meta['chunk_id']}] ({len(chunk.split()):3d} words): {chunk[:70]}...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed all chunks (dense index)\n",
    "chunk_embeddings = embed_texts(all_chunks, task_type=\"RETRIEVAL_DOCUMENT\")\n",
    "print(f\"Dense index: {len(chunk_embeddings)} chunks (dimension: {len(chunk_embeddings[0])})\")\n",
    "\n",
    "# Build sparse index (TF-IDF for keyword search)\n",
    "tfidf_vectorizer, tfidf_matrix = build_sparse_index(all_chunks)\n",
    "print(f\"Sparse index: {tfidf_matrix.shape[0]} chunks, {tfidf_matrix.shape[1]} vocabulary terms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "711ghdli54o",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 2b: Test Hybrid Retrieval\n",
    "\n",
    "Compare semantic-only vs. hybrid (semantic + keyword) retrieval. Hybrid search catches queries where exact terms matter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0t6zsnfrgfvs",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare semantic-only vs hybrid retrieval\n",
    "test_query = all_chunks[0].split()[:5]  # Use first few words from doc as test\n",
    "test_query = \" \".join(test_query)\n",
    "\n",
    "print(f\"Test query: \\\"{test_query}\\\"\\n\")\n",
    "\n",
    "# Semantic only\n",
    "print(\"Semantic-only results:\")\n",
    "sem_results = search(test_query, chunk_embeddings, all_chunks, top_k=3)\n",
    "for idx, score, text in sem_results:\n",
    "    print(f\"  [{score:.3f}] Chunk {idx}: {text[:70]}...\")\n",
    "\n",
    "# Hybrid\n",
    "print(\"\\nHybrid results (0.6 semantic + 0.4 keyword):\")\n",
    "hyb_results = hybrid_search(\n",
    "    test_query, chunk_embeddings, all_chunks,\n",
    "    tfidf_vectorizer, tfidf_matrix, top_k=3\n",
    ")\n",
    "for idx, combo, sem, sparse, text in hyb_results:\n",
    "    print(f\"  [{combo:.3f}] Chunk {idx} (sem={sem:.3f}, kw={sparse:.3f}): {text[:60]}...\")\n",
    "\n",
    "print(\"\\nHybrid search boosts chunks that match both meaning AND exact keywords.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cpm71sbv5k",
   "metadata": {},
   "source": [
    "### Step 2c: Metadata Filtering\n",
    "\n",
    "In production, you often want to restrict search to specific document categories (e.g., only search HR policies, not engineering docs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2p3z7w8v8yr",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate metadata filtering\n",
    "print(\"All documents in index:\")\n",
    "seen = set()\n",
    "for meta in chunk_metadata:\n",
    "    if meta[\"doc_title\"] not in seen:\n",
    "        print(f\"  Doc {meta['doc_index']}: {meta['doc_title'][:60]}\")\n",
    "        seen.add(meta[\"doc_title\"])\n",
    "\n",
    "# Search with filter (restrict to first document only)\n",
    "print(f\"\\nFiltered search (doc_index=0 only):\")\n",
    "filtered = search_with_filter(\n",
    "    \"TODO: enter a question here\",\n",
    "    chunk_embeddings, all_chunks, chunk_metadata,\n",
    "    top_k=3, filter_field=\"doc_index\", filter_value=0\n",
    ")\n",
    "for idx, score, text in filtered:\n",
    "    meta = chunk_metadata[idx]\n",
    "    print(f\"  [{score:.3f}] {meta['chunk_id']}: {text[:70]}...\")\n",
    "\n",
    "print(\"\\nMetadata filtering ensures the retriever only searches relevant document subsets.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 3: Write Your System Prompt (v1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── TODO: Write your grounding system prompt ─────────────\n",
    "\n",
    "SYSTEM_PROMPT_V1 = \"\"\"You are a [ROLE] assistant for [DOMAIN].\n",
    "\n",
    "Rules:\n",
    "- Answer ONLY based on the provided context excerpts\n",
    "- If the answer is not in the context, say: \"I don't have that information in the provided documents.\"\n",
    "- Cite which chunk(s) support your answer using [Chunk N] notation\n",
    "- Keep answers concise (under 80 words)\n",
    "- Do not make up information or combine context with your own knowledge\n",
    "\n",
    "Scope:\n",
    "- You answer questions about [TOPICS]\n",
    "- For questions outside this scope, say: \"That's outside my area of expertise.\"\n",
    "\n",
    "TODO: Customize the placeholders above for your chosen track!\n",
    "\"\"\"\n",
    "\n",
    "print(\"System prompt v1:\")\n",
    "print(SYSTEM_PROMPT_V1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 4: Create Your Golden Test Set\n",
    "\n",
    "Create 10+ questions BEFORE running the pipeline. Mix easy, medium, hard, and refusal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── TODO: Create your golden test set ────────────────────\n",
    "# Include at least: 5-6 easy, 2-3 medium/edge, 2-3 refusal\n",
    "\n",
    "GOLDEN_SET = [\n",
    "    # Easy questions (direct answers in a single chunk)\n",
    "    {\n",
    "        \"id\": \"Q01\",\n",
    "        \"question\": \"TODO: Write your first easy question\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"easy\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q02\",\n",
    "        \"question\": \"TODO: Write your second easy question\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"easy\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q03\",\n",
    "        \"question\": \"TODO\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"easy\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q04\",\n",
    "        \"question\": \"TODO\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"easy\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q05\",\n",
    "        \"question\": \"TODO\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"easy\",\n",
    "    },\n",
    "    # Medium questions (may require combining chunks)\n",
    "    {\n",
    "        \"id\": \"Q06\",\n",
    "        \"question\": \"TODO: Write a question that spans multiple documents\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"medium\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q07\",\n",
    "        \"question\": \"TODO\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"medium\",\n",
    "    },\n",
    "    # Edge case questions\n",
    "    {\n",
    "        \"id\": \"Q08\",\n",
    "        \"question\": \"TODO: Write an ambiguous or boundary question\",\n",
    "        \"expected_keywords\": [\"TODO\"],\n",
    "        \"difficulty\": \"hard\",\n",
    "    },\n",
    "    # Refusal questions (no answer in knowledge base)\n",
    "    {\n",
    "        \"id\": \"Q09\",\n",
    "        \"question\": \"TODO: Write a question with NO answer in your documents\",\n",
    "        \"expected_keywords\": [\"don't have\", \"not in\", \"no information\", \"outside\"],\n",
    "        \"difficulty\": \"refusal\",\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"Q10\",\n",
    "        \"question\": \"TODO: Another out-of-scope question\",\n",
    "        \"expected_keywords\": [\"don't have\", \"not in\", \"no information\", \"outside\"],\n",
    "        \"difficulty\": \"refusal\",\n",
    "    },\n",
    "]\n",
    "\n",
    "print(f\"Golden set: {len(GOLDEN_SET)} questions\")\n",
    "for q in GOLDEN_SET:\n",
    "    print(f\"  [{q['difficulty']:7s}] {q['id']}: {q['question']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 5: Evaluate v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run v1 evaluation\n",
    "# (Make sure you've filled in your SYSTEM_PROMPT_V1 and GOLDEN_SET above!)\n",
    "\n",
    "v1_results = evaluate_retrieval(GOLDEN_SET, all_chunks, chunk_embeddings, SYSTEM_PROMPT_V1)\n",
    "\n",
    "v1_acc = v1_results[\"keyword_match\"].mean()\n",
    "print(f\"\\nv1 Keyword Accuracy: {v1_acc:.0%} ({v1_results['keyword_match'].sum()}/{len(v1_results)})\\n\")\n",
    "print(v1_results[[\"id\", \"difficulty\", \"keyword_match\", \"top_score\"]].to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM-as-Judge scoring for v1\n",
    "class RAGScore(BaseModel):\n",
    "    groundedness: int = Field(description=\"1-5: Is the answer supported by the context?\")\n",
    "    relevance: int = Field(description=\"1-5: Does the answer address the question?\")\n",
    "    explanation: str = Field(description=\"Brief explanation of the scores\")\n",
    "\n",
    "v1_scores = []\n",
    "for _, row in v1_results.iterrows():\n",
    "    eval_prompt = f\"\"\"Rate this RAG system answer.\n",
    "\n",
    "Question: {row['question']}\n",
    "Answer: {row['answer']}\n",
    "\n",
    "Score groundedness (1-5): Is the answer based on facts from context, not made up?\n",
    "Score relevance (1-5): Does the answer address the question asked?\"\"\"\n",
    "    score = generate_structured(eval_prompt, RAGScore, label=f\"eval_v1_{row['id']}\")\n",
    "    v1_scores.append({\"id\": row[\"id\"], \"groundedness\": score.groundedness,\n",
    "                       \"relevance\": score.relevance, \"explanation\": score.explanation})\n",
    "\n",
    "v1_scores_df = pd.DataFrame(v1_scores)\n",
    "print(\"v1 LLM-as-Judge Scores:\")\n",
    "print(v1_scores_df[[\"id\", \"groundedness\", \"relevance\"]].to_string(index=False))\n",
    "print(f\"\\nAvg Groundedness: {v1_scores_df['groundedness'].mean():.1f}/5\")\n",
    "print(f\"Avg Relevance:    {v1_scores_df['relevance'].mean():.1f}/5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 6: Improve Your Prompt → v2\n",
    "\n",
    "Based on v1 errors, improve your system prompt. Common improvements:\n",
    "- Add specific rules for categories or edge cases\n",
    "- Add examples of good answers\n",
    "- Strengthen refusal instructions\n",
    "- Add format requirements (citations, length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── TODO: Write your improved system prompt ──────────────\n",
    "\n",
    "SYSTEM_PROMPT_V2 = \"\"\"TODO: Write your improved system prompt here.\n",
    "\n",
    "Start from v1 and add:\n",
    "1. Specific rules based on errors you observed\n",
    "2. Examples of good answers for tricky questions\n",
    "3. Stronger refusal instructions\n",
    "4. Any additional scope or format rules\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "print(\"System prompt v2:\")\n",
    "print(SYSTEM_PROMPT_V2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 7: Evaluate v2 and Compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run v2 evaluation\n",
    "v2_results = evaluate_retrieval(GOLDEN_SET, all_chunks, chunk_embeddings, SYSTEM_PROMPT_V2)\n",
    "\n",
    "v2_acc = v2_results[\"keyword_match\"].mean()\n",
    "print(f\"v2 Keyword Accuracy: {v2_acc:.0%} ({v2_results['keyword_match'].sum()}/{len(v2_results)})\\n\")\n",
    "\n",
    "# Compare v1 vs v2\n",
    "print(\"=\" * 60)\n",
    "comparison = compare_versions(v1_results, v2_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2b4c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect answers for GOLDEN_SET with (1) semantic search then (2) hybrid search\n",
    "\n",
    "# 1) Semantic-only (uses rag_query -> search)\n",
    "semantic_inspect = evaluate_retrieval(\n",
    "    GOLDEN_SET, all_chunks, chunk_embeddings, SYSTEM_PROMPT_V2, top_k=3\n",
    ")[[\"id\", \"difficulty\", \"question\", \"answer\", \"top_score\", \"keyword_match\"]]\n",
    "\n",
    "# 2) Hybrid (TF-IDF + semantic) -> build context -> generate\n",
    "hybrid_rows = []\n",
    "for qa in GOLDEN_SET:\n",
    "    hyb = hybrid_search(\n",
    "        qa[\"question\"],\n",
    "        chunk_embeddings, all_chunks,\n",
    "        tfidf_vectorizer, tfidf_matrix,\n",
    "        top_k=3, semantic_weight=0.6, keyword_weight=0.4\n",
    "    )\n",
    "\n",
    "    context = \"\\n\\n\".join(\n",
    "        f\"[Chunk {j+1}]\\n{text}\" for j, (_, _, _, _, text) in enumerate(hyb)\n",
    "    )\n",
    "\n",
    "    prompt = f\"\"\"{SYSTEM_PROMPT_V2}\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "\n",
    "Question: {qa[\"question\"]}\"\"\"\n",
    "\n",
    "    ans = generate(prompt, temperature=0.3, label=f\"hybrid_{qa['id']}\")\n",
    "    hybrid_rows.append({\n",
    "        \"id\": qa[\"id\"],\n",
    "        \"difficulty\": qa[\"difficulty\"],\n",
    "        \"question\": qa[\"question\"],\n",
    "        \"answer\": ans,\n",
    "        \"hybrid_top_score\": hyb[0][1] if hyb else None,\n",
    "    })\n",
    "\n",
    "hybrid_inspect = pd.DataFrame(hybrid_rows)\n",
    "\n",
    "# Optional: side-by-side comparison\n",
    "comparison = semantic_inspect.merge(hybrid_inspect, on=[\"id\", \"difficulty\", \"question\"], suffixes=(\"_semantic\", \"_hybrid\"))\n",
    "display(comparison[[\"id\",\"difficulty\",\"question\",\"answer_semantic\",\"answer_hybrid\",\"top_score\",\"hybrid_top_score\",\"keyword_match\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 8: Error Analysis\n",
    "\n",
    "Reflect on your results. Fill in the sections below.\n",
    "\n",
    "### Top 3 Error Patterns\n",
    "\n",
    "**Pattern 1:** [Name]\n",
    "- **What happens:** [Description]\n",
    "- **Example:** [Specific question that failed]\n",
    "- **Root cause:** [Why the model gets this wrong]\n",
    "\n",
    "**Pattern 2:** [Name]\n",
    "- **What happens:** [Description]\n",
    "- **Example:** [Specific question that failed]\n",
    "- **Root cause:** [Why the model gets this wrong]\n",
    "\n",
    "**Pattern 3:** [Name]\n",
    "- **What happens:** [Description]\n",
    "- **Example:** [Specific question that failed]\n",
    "- **Root cause:** [Why the model gets this wrong]\n",
    "\n",
    "### Changes from v1 → v2\n",
    "\n",
    "| Change | Why | Impact |\n",
    "|--------|-----|--------|\n",
    "| [What you modified] | [Why you expected improvement] | [Actual effect on accuracy] |\n",
    "| [Second change] | [Reasoning] | [Result] |\n",
    "\n",
    "### Remaining Risks\n",
    "\n",
    "- **Risk 1:** [What could still go wrong]\n",
    "- **Mitigation:** [How you would handle this in production]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 9: Export Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export all results\n",
    "# Golden set\n",
    "with open(\"day3_lab2_golden_set.json\", \"w\") as f:\n",
    "    json.dump(GOLDEN_SET, f, indent=2)\n",
    "print(\"Exported golden set to day3_lab2_golden_set.json\")\n",
    "\n",
    "# v1 and v2 results\n",
    "v1_export = v1_results.to_dict(orient=\"records\")\n",
    "v2_export = v2_results.to_dict(orient=\"records\")\n",
    "with open(\"day3_lab2_results.json\", \"w\") as f:\n",
    "    json.dump({\"v1\": v1_export, \"v2\": v2_export}, f, indent=2, default=str)\n",
    "print(\"Exported results to day3_lab2_results.json\")\n",
    "\n",
    "# Prompt log\n",
    "if PROMPT_LOG:\n",
    "    log_df = pd.DataFrame(PROMPT_LOG)\n",
    "    log_df.to_csv(\"day3_lab2_prompt_log.csv\", index=False)\n",
    "    print(f\"Exported {len(PROMPT_LOG)} API calls to day3_lab2_prompt_log.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Checklist\n",
    "\n",
    "- [ ] Chose a track and set `SELECTED_TRACK`\n",
    "- [ ] Reviewed and understood the knowledge base documents\n",
    "- [ ] Configured chunking strategy with rationale\n",
    "- [ ] Chunks include metadata (chunk_id, doc_title)\n",
    "- [ ] Built both dense (embedding) and sparse (TF-IDF) indices\n",
    "- [ ] Tested hybrid retrieval vs semantic-only\n",
    "- [ ] Tested metadata filtering\n",
    "- [ ] Wrote system prompt v1 with grounding rules\n",
    "- [ ] Created golden test set with 10+ questions (easy + medium + refusal)\n",
    "- [ ] Ran v1 evaluation (keyword accuracy + LLM-as-judge)\n",
    "- [ ] Improved prompt to v2 based on v1 errors\n",
    "- [ ] Ran v2 evaluation and compared versions\n",
    "- [ ] Completed error analysis (patterns, changes, risks)\n",
    "- [ ] Exported all results\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "→ Proceed to the [Day 3 Assignment](03-04_assignment_3.qmd)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
