diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..73a1b48 --- /dev/null +++ b/.env.example @@ -0,0 +1,3 @@ +GROQ_API_KEY=your_groq_api_key_here +GROQ_MODEL=llama-3.3-70b-versatile +HF_API_TOKEN=your_hf_api_token_here \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1b2a90b --- /dev/null +++ b/.gitignore @@ -0,0 +1,35 @@ +# Python +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ + +# Virtual environments +.venv/ +venv/ + +# Local environment variables +.env + +# Generated databases +*.db +*.sqlite +*.sqlite3 + + +# OS/editor files +.DS_Store +.vscode/ +.idea/ + +# Build artifacts +dist/ +build/ +*.egg-info/ +# Local inspection exports +data/*.xlsx +data/*_preview.csv + +# Local RAG vector index +data/chroma_db/ diff --git a/NOTES.md b/NOTES.md new file mode 100644 index 0000000..2b78618 --- /dev/null +++ b/NOTES.md @@ -0,0 +1,358 @@ +# NOTES + +## 1. CSV problems, fixes, and what I left + +The CSV was usable, but it needed cleaning before using it for search and RAG. + +Problems I found: + +- `date_added` had many leading spaces and needed parsing. +- Some text fields had messy whitespace such as repeated spaces, newlines, tabs, and non-breaking spaces. +- Several optional fields were missing, especially `director`, `cast`, and `country`. +- `country`, `cast`, `director`, and `listed_in` were comma-separated multi-value fields. +- Some list fields had trailing commas, which created empty values after splitting. +- `duration` was stored as text like `90 min`, `1 Season`, or `3 Seasons`. +- I found one duplicate-content row where all cleaned fields matched another row except `show_id`. + +Fixes made: + +- Cleaned whitespace in text fields. +- Parsed `date_added` into `YYYY-MM-DD`. +- Parsed `release_year` as an integer. +- Split `duration` into `duration_value` and `duration_unit`. +- Converted missing user-facing metadata to `"Unknown"` for fields like director, cast, country, and rating. +- Normalized rating abbreviations: `NR` to `Not Rated` and `UR` to `Unrated`. +- Split multi-value fields into separate child tables. +- Removed empty list values and duplicate list values while preserving order. +- Dropped unusable rows and exact duplicate-content rows. + +What I left: + +- I did not drop rows just because optional metadata was missing. A title can still be useful without director, cast, country, rating, or date_added. +- I did not manually edit the CSV. All cleaning happens in code. +- I did not add external data such as IMDb ratings, posters, languages, or current Netflix availability because that would make the assignment less reproducible. +- I did not do fuzzy duplicate matching. I only removed exact duplicate cleaned content except `show_id`. + +--- + +## 2. Schema decisions + +I used SQLite because the dataset is small, local, and easy to rebuild. + +The main table is `titles`. It stores one row per title with: + +- `show_id` +- `type` +- `title` +- `release_year` +- `rating` +- `duration_value` +- `duration_unit` +- `date_added` +- `description` + +I used `show_id` as the primary key because title names are not guaranteed to be unique. + +For multi-value fields, I created child tables: + +- `title_countries` +- `title_genres` +- `title_cast` +- `title_directors` + +I did this because fields like `country`, `cast`, `director`, and `listed_in` can contain multiple values in one CSV cell. Storing them separately makes filtering and stats cleaner. For example, a title with `United States, India` should count under both countries. + +Each child table has a `position` column to preserve the original order from the CSV, especially for cast and directors. + +--- + +## 3. RAG document design + +I used one RAG document per Netflix title. I did not chunk further because each title row is already short. + +For each title, I embedded a text block containing: + +- title +- type +- release year +- rating +- duration +- countries +- genres +- directors +- cast +- description + +I included metadata, not only description, because users may ask questions like: + +```text +Indian comedy movie +TV-MA Japanese show +movie with Vijay +light heart movie of India after 2018 +``` + +These questions depend on country, genre, rating, type, year, and cast, not only the description. + +Each Chroma document stores metadata like: + +```python +{ + "show_id": "...", + "title": "...", + "type": "Movie", + "release_year": 2019, + "rating": "TV-14", + "countries": ["India"], + "genres": ["Comedies", "Dramas"] +} +``` + +I store `countries` and `genres` as lists so ChromaDB can filter them using `$contains`. + +--- + +## 4. Moving from local SentenceTransformer to Hugging Face API + +Originally, the RAG system used local SentenceTransformer embeddings. + +Old flow: + +```text +Text +→ local SentenceTransformer model +→ embedding vector +→ ChromaDB +``` + +This worked locally, but it was heavy for deployment because `sentence-transformers` pulls large ML dependencies like PyTorch. + +I changed the embedding system to use Hugging Face Inference API. + +New flow: + +```text +Text +→ Hugging Face Inference API +→ embedding vector +→ ChromaDB +``` + +The embedding model is still: + +```text +sentence-transformers/all-MiniLM-L6-v2 +``` + +but the model is not loaded inside the backend anymore. + +This made deployment lighter because the backend only needs the lightweight Hugging Face client, not the full local embedding model stack. + +I also kept manual L2 normalization because the old local SentenceTransformer flow used normalized embeddings. This keeps retrieval behavior consistent. + +--- + +## 5. Chroma indexing flow + +When I run: + +```bash +python -m rag.index +``` + +the system: + +1. Loads cleaned Netflix titles from SQLite. +2. Converts each title into a RAG document. +3. Sends document text to Hugging Face Inference API. +4. Receives an embedding vector. +5. Normalizes the vector. +6. Stores the id, document text, metadata, and embedding in ChromaDB. + +The Chroma index is stored locally in: + +```text +data/chroma_db/ +``` + +This folder is generated and should not be committed. + +I also store embedding metadata such as model name and dimension so the retriever can detect mismatches if the embedding model changes. + +--- + +## 6. AI Intelligence Layer + +I added an AI Intelligence Layer before Chroma retrieval. + +Old `/ask` flow: + +```text +User question +→ embed raw question +→ search whole Chroma vector database +→ Groq final answer +``` + +New `/ask` flow: + +```text +User question +→ LLM query planner +→ semantic_query + Chroma where filter +→ filtered Chroma vector search +→ Groq final answer +``` + +The planner returns a direct Chroma-searchable plan. + +Example user query: + +```text +I want light heart movie of India after 2018 +``` + +Planner output: + +```python +SearchPlan( + semantic_query="lighthearted feel-good entertaining comedy drama", + where={ + "$and": [ + {"countries": {"$contains": "India"}}, + {"type": "Movie"}, + {"release_year": {"$gt": 2018}} + ] + } +) +``` + +This makes the retriever smarter before vector search happens. + +The exact filters decide which titles are eligible. The semantic query decides which eligible titles are most relevant. + +--- + +## 7. Why the AI layer was needed + +Plain vector search is semantic, not strict. + +If the user asks: + +```text +Indian comedy movies +``` + +plain vector search may return non-Indian titles if their descriptions are semantically similar. + +The AI Intelligence Layer solves this by creating a direct ChromaDB search plan: + +```text +semantic_query = funny entertaining comedy movies +where = countries contains India, genres contains Comedies, type Movie +``` + +So ChromaDB searches only inside the correct filtered subset. + +This is better than sending the raw query directly to vector search. + +--- + +## 8. Example `/ask` flow + +Example query: + +```text +Suggest Bollywood comedy movies after 2018 +``` + +Planner creates: + +```python +SearchPlan( + semantic_query="funny lighthearted entertaining comedy movies", + where={ + "$and": [ + {"countries": {"$contains": "India"}}, + {"genres": {"$contains": "Comedies"}}, + {"type": "Movie"}, + {"release_year": {"$gt": 2018}} + ] + } +) +``` + +Then the retriever does: + +```text +Embed semantic_query with Hugging Face +Search ChromaDB using query embedding + where filter +Return matching titles +``` + +Then Groq receives only retrieved catalogue titles and writes the final grounded answer. + +--- + +## 9. RAG and LLM guardrails + +The final answer prompt tells Groq: + +- Use only provided catalogue sources. +- Do not use outside knowledge. +- Do not invent titles, actors, countries, ratings, years, or plot details. +- Return JSON with `answer` and `used_show_ids`. +- Only cite `show_id`s actually used in the answer. + +The API then returns only the sources listed in `used_show_ids`. + +This prevents the API from showing every retrieved title as a source if the answer only used some of them. + +--- + +## 10. What would break at 100,000 titles? + +At 100,000 titles, the current design would still work conceptually, but these parts would need improvement: + + +- Rebuilding the full Chroma index every time would be slow. I would make indexing incremental and only re-embed changed rows. +- Local Chroma storage may not be ideal for production. I would consider a managed vector store or a more controlled indexing pipeline. +- LLM latency and cost would matter more. I would cache common answers and keep retrieved context smaller. +- SQLite is good for this assignment and dataset size, but at larger production scale I would consider PostgreSQL because it handles concurrent access, indexing, backups, and production operations better. + +--- + +## 11. AI usage + +I used AI tools while building this project. + +I used AI for: + +- planning the ingestion package structure +- drafting first versions of helper functions +- discussing cleaning choices like `"Unknown"` vs `NULL` +- comparing schema options for multi-value fields +- drafting parts of the FastAPI and RAG scaffolding +- reviewing prompts and source-grounding behavior +- debugging deployment issues +- designing the AI Intelligence Layer before Chroma retrieval + +I changed AI-generated output in multiple places instead of accepting it blindly. + +One example was the RAG prompt and source handling. The first version returned all retrieved titles as sources, even when the answer only used some of them. I changed the prompt to require a JSON response with `answer` and `used_show_ids`, added guardrails against outside knowledge and invented metadata, and updated the API to return only the titles listed in `used_show_ids`. + +Another example was the AI Intelligence Layer. I first considered rule-based normalization and tool-based approaches, but I kept the final implementation simpler: the LLM directly generates a Chroma-searchable plan, and Python only validates and executes it. + +I also rejected suggestions that felt unnecessary for this assignment, such as MinMax scaling numeric fields, adding many boolean quality-flag columns, using a full framework like LangChain, or adding fuzzy duplicate matching without time to evaluate it. + +I ran the code, checked outputs, and adjusted decisions based on actual behavior. I can explain the final path from CSV to SQLite, from SQLite to Chroma, from Chroma retrieval to Groq answer, and from raw query to AI-planned filtered retrieval. + +--- + +## 12. What I would improve next + +- **TMDB enrichment:** If business needs richer data, I would use TMDB to fill missing fields like unknown director, cast, language, poster, and external ratings with source/confidence tracking. + +- **Fallback retrieval:** If strict filters return no results, I would show a clear “no exact match found” message and then return the closest relaxed alternatives. + +- **RAG evaluation:** I would create a small evaluation set of sample questions and expected `show_id`s to measure whether retrieval quality is improving. diff --git a/README.md b/README.md index 6074be7..70495ef 100644 --- a/README.md +++ b/README.md @@ -1,171 +1,454 @@ -# Intern Assignment — Netflix Catalog Q&A System +# Netflix Catalogue Q&A System -Welcome. This is a take-home assignment to help us understand how you think about -data, backend, and AI/LLM problems. We're not looking for a perfect product — -we're looking for the **reasoning behind your choices**. +This project is a backend system for searching and asking questions about a Netflix catalogue dataset. -**Time expectation:** ~5–7 hours. Stop at 8 hours regardless of where you are. -We'd rather see a smaller scope done thoughtfully than a large scope half-done. +It has four main parts: -**You can use AI tools** (ChatGPT, Cursor, Copilot, Claude, etc.). We expect you -to. But you must understand every line you submit and be able to explain it. -There is a live walkthrough after submission where we will ask you to defend -specific lines, modify code on the spot, and debug something you didn't write. -If you can't, the assignment counts for very little. +1. **Data ingestion** — cleans the raw Netflix CSV and stores it in SQLite. +2. **Structured API** — provides FastAPI endpoints for title search, title lookup, and catalogue stats. +3. **RAG Q&A** — retrieves relevant catalogue titles from ChromaDB and uses Groq to answer questions with source `show_id`s. +4. **AI Intelligence Layer** — converts a raw user query into a Chroma-searchable plan before vector retrieval. + +The source dataset is: + +```text +data/netflix_titles.csv +``` + +Generated files such as `data/netflix.db`, `data/chroma_db/`, `.env`, and `.venv/` are not committed because they can be rebuilt locally. + +--- + +## Tech Stack + +- Python 3.10+ +- SQLite +- FastAPI +- ChromaDB +- Hugging Face Inference API +- Groq +- pytest +- ruff +- Optional demo UI: Streamlit + +--- + +## High-Level Architecture + +```text +Netflix CSV + ↓ +Ingestion pipeline + ↓ +Clean SQLite database + ↓ +RAG document builder + ↓ +Hugging Face embedding API + ↓ +ChromaDB vector index + ↓ +/ask API endpoint + ↓ +AI Intelligence Layer + ↓ +Filtered Chroma vector search + ↓ +Groq final answer +``` + +--- + +## AI Intelligence Layer + +The `/ask` endpoint does not send the raw user question directly to ChromaDB. + +Instead, the system first sends the question to an LLM-based query planner. The planner returns a direct Chroma search plan containing: + +1. `semantic_query` — a rewritten query used for vector search. +2. `where` — a ChromaDB metadata filter used for exact constraints. + +Example user query: + +```text +I want light heart movie of India after 2018 +``` + +Planner output: + +```json +{ + "semantic_query": "lighthearted feel-good entertaining comedy drama", + "where": { + "$and": [ + {"countries": {"$contains": "India"}}, + {"type": "Movie"}, + {"release_year": {"$gt": 2018}} + ] + } +} +``` + +This improves retrieval because: + +```text +Metadata filters decide eligibility. +Semantic search decides relevance. +``` + +For example, if the user asks for Indian movies, ChromaDB only searches inside Indian titles instead of searching the entire vector database. + +--- + +## Hugging Face Embeddings + +Originally, the project used local SentenceTransformer embeddings. That worked locally, but it added heavy deployment dependencies such as PyTorch. + +The current version uses Hugging Face Inference API for embeddings. + +Old approach: + +```text +Text → local SentenceTransformer → embedding +``` + +New approach: + +```text +Text → Hugging Face Inference API → embedding +``` + +The embedding model is still: + +```text +sentence-transformers/all-MiniLM-L6-v2 +``` + +but the model is not loaded locally anymore. + +Benefits: + +- No local PyTorch model loading +- Smaller dependency size +- Easier Render deployment +- Same embedding model family +- Normalized embeddings preserved + +--- + +## Setup + +Clone the repository: + +```bash +git clone +cd intern-assignment +``` + +Create and activate a virtual environment: + +```bash +python3 -m venv .venv +source .venv/bin/activate +``` + +Install backend dependencies: + +```bash +pip install -r requirements.txt +``` + +Create a local environment file: + +```bash +cp .env.example .env +``` + +Add your API keys in `.env`: + +```env +GROQ_API_KEY=your_groq_api_key_here +GROQ_MODEL=llama-3.3-70b-versatile +HF_API_TOKEN=your_hf_api_token_here +``` + +> [!IMPORTANT] +> You need a Hugging Face token with read access. The project uses Hugging Face Inference API for embeddings instead of loading SentenceTransformer locally. + +> [!WARNING] +> If you change the embedding model or metadata structure, rebuild the Chroma index with `python -m rag.index`. + +Do not commit `.env`. + +--- + +## How to Run + +### 1. Run ingestion + +```bash +python ingest.py +``` + +This reads the raw CSV and creates: + +```text +data/netflix.db +``` + +The script prints a summary report with rows read, rows loaded, rows dropped, missing values handled, and cleaning fixes applied. --- -## What you're building +### 2. Build the RAG index -A small system that lets a user query the Netflix catalogue in two ways: -1. **Structured search** — filter shows by country, year, type, rating. -2. **Natural-language Q&A** — ask questions like *"Suggest me a comedy movie from - India with a strong female lead"* and get an answer grounded in the catalogue. +```bash +python -m rag.index +``` + +This creates the local Chroma vector index: + +```text +data/chroma_db/ +``` -You will also write up your reasoning. The write-up matters as much as the code. +The index stores one embedded document per Netflix title. + +Each Chroma record stores: + +- document id / `show_id` +- embedded document text +- metadata such as title, type, rating, release year, countries, and genres +- embedding vector generated by Hugging Face --- -## Provided data +### 3. Start the API -`data/netflix_titles.csv` — ~6,200 rows of Netflix titles as of 2019. Real, -public, and **deliberately not cleaned for you**. Treat it as you'd treat any -data you receive in the real world. +```bash +uvicorn api.main:app --reload +``` -Columns: `show_id, type, title, director, cast, country, date_added, -release_year, rating, duration, listed_in, description` +The API runs at: -You will discover problems with this data as you work. Document what you find. +```text +http://127.0.0.1:8000 +``` + +Interactive docs are available at: + +```text +http://127.0.0.1:8000/docs +``` --- -## The four parts +## API Endpoints + +### Health check + +```bash +curl http://127.0.0.1:8000/health +``` + +### List titles + +```bash +curl "http://127.0.0.1:8000/titles?country=India&type=Movie&page=1&page_size=5" +``` + +Supported filters: + +- `country` +- `release_year` +- `type` +- `rating` +- `page` +- `page_size` + +### Get one title by ID + +```bash +curl http://127.0.0.1:8000/titles/81075235 +``` + +### Stats + +```bash +curl http://127.0.0.1:8000/stats +``` -### Part 1 — Data ingestion (scripting + data) +Returns total titles, count by type, and top 10 countries. -Write a Python script `ingest.py` that: -- Loads `data/netflix_titles.csv` -- Cleans/normalises the data (your judgement call — explain in NOTES.md) -- Writes the cleaned data into a SQLite database `data/netflix.db` -- Prints a clear summary report when it finishes (rows loaded, rows dropped/fixed, - any anomalies you found) +### Ask a question -You should be able to run `python ingest.py` from a fresh checkout and have it -work end-to-end. +```bash +curl -X POST http://127.0.0.1:8000/ask \ + -H "Content-Type: application/json" \ + -d '{"question":"I want light heart movie of India after 2018"}' +``` -### Part 2 — Backend API (backend) +Example response: + +```json +{ + "answer": "A good match is 15-Aug, 90 ML, Kaake Da Viyah or House Arrest. They are all Indian movies released after 2018 and listed under Comedies, so they fit your request for a light heart movie.", + "sources": [ + {"show_id": "81033429", "title": "15-Aug"}, + {"show_id": "81147278", "title": "90 ML"}, + {"show_id": "81155859", "title": "Kaake Da Viyah"}, + {"show_id": "81076114", "title": "House Arrest"} + ] +} +``` -Build a small HTTP API using **FastAPI** (preferred) or Flask, with these -endpoints: +--- -- `GET /titles` — list titles, with query params for filtering by `country`, - `release_year`, `type` (Movie / TV Show), and `rating`. Support pagination. -- `GET /titles/{show_id}` — fetch a single title by id. -- `GET /stats` — return small summary stats (total titles, count by type, - count by country — top 10). +## `/ask` Flow + +```text +User question + ↓ +AI Intelligence Layer + ↓ +LLM creates: +- semantic_query +- Chroma where filter + ↓ +Hugging Face embeds semantic_query + ↓ +ChromaDB filtered vector search + ↓ +Retrieved catalogue titles + ↓ +Groq final answer + ↓ +Answer + source show_ids +``` -Run it locally with one command. Document that command in your README. +Example: -### Part 3 — RAG-based Q&A (AI / LLM) +```text +User: +Suggest Bollywood comedy movies after 2018 -Add an endpoint: -- `POST /ask` — body: `{ "question": "..." }`. Returns an answer grounded in - the catalogue, plus the `show_id`s of the titles your answer was based on. +AI planner: +semantic_query = "funny lighthearted entertaining comedy movies" +where = countries contains India, genres contains Comedies, type Movie, release_year > 2018 +``` -Build a basic RAG pipeline: -1. Embed the catalogue (you decide what text to embed — that's a real design - choice, document it). -2. On a question, retrieve the most relevant titles. -3. Pass them to an LLM with a prompt and return the answer + sources. +--- -You can use **any** LLM — OpenAI, Groq (free tier), Gemini (free tier), Ollama -locally, anything. You can use any embedding model (sentence-transformers -locally is free and fine). You can use any vector store (FAISS, Chroma, or even -just numpy + cosine similarity — for 6k rows you don't need a real vector DB). +## Optional Streamlit Demo -The answer must cite which titles it used. Hallucinated answers without sources -will be penalised heavily. +The Streamlit demo is only a thin UI over the FastAPI backend. It does not replace the API. -### Part 4 — NOTES.md (the most important file) +Install demo dependency: -A single markdown file in the project root, **maximum 2 pages**, answering: +```bash +pip install -r requirements-demo.txt +``` -1. **What problems did you find in the CSV?** What did you fix, what did you - leave, and why? -2. **Schema decisions.** How did you store multi-value fields like `country`, - `cast`, `listed_in`? Why? -3. **RAG choices.** What text did you embed for each title? What chunk size / - model did you pick? Why? -4. **What would break at 100,000 titles instead of 6,000?** Be concrete. -5. **AI usage — be honest.** What did you ask AI tools to write? Where did you - override or rewrite their output, and why? Where did you accept their output - without fully understanding it? -6. **What would you do with another 4 hours?** +Start the backend first: -We read this **first**, before looking at your code. Don't bullshit it. Honest -"I don't know why I picked this" beats a confident-sounding lie. We will ask -follow-up questions on every answer in the live round. +```bash +python ingest.py +python -m rag.index +uvicorn api.main:app --reload +``` + +Then run: + +```bash +streamlit run demo/streamlit_app.py +``` + +The Streamlit sidebar lets you set the API base URL. For local use, keep: + +```text +http://127.0.0.1:8000 +``` + +For deployed use, paste the Render backend URL. --- -## Coding standards (read STANDARDS.md before you start) +## Render Deployment -You must follow `STANDARDS.md` — it lists the conventions we use so we can -read your code quickly. It's short. Skipping it makes your submission look -like vibe-coded output we can't review. +For Render backend deployment, use this build command: + +```bash +pip install -r requirements.txt && python ingest.py && python -m rag.index +``` + +Use this start command: + +```bash +python -m uvicorn api.main:app --host 0.0.0.0 --port ${PORT:-10000} +``` + +Set these environment variables on Render: + +```env +GROQ_API_KEY=your_groq_api_key_here +GROQ_MODEL=llama-3.3-70b-versatile +HF_API_TOKEN=your_hf_api_token_here +PYTHON_VERSION=3.11.11 +``` + +`HF_API_TOKEN` is required during build because `python -m rag.index` generates embeddings using Hugging Face. + +The first `/ask` request after a restart can be slower because the retriever and API clients are loaded lazily. --- -## What to submit +## How to Test + +Run the tests: + +```bash +python -m pytest -q +``` -A single zip file (or a public git repo link) containing: +Run linting: +```bash +ruff check . ``` -your_submission/ -├── README.md # how to install + run, in your own words -├── NOTES.md # the write-up (Part 4) -├── STANDARDS.md # (optional — copy ours, or write your own and tell us why) -├── requirements.txt # or pyproject.toml -├── .gitignore -├── ingest.py -├── api/ # your FastAPI / Flask app -│ └── ... -├── rag/ # your RAG code -│ └── ... -├── tests/ # at least a few real tests, see STANDARDS.md -│ └── ... -└── data/ - └── netflix_titles.csv # leave the original here for reproducibility + +To fix import sorting or formatting issues: + +```bash +ruff check --select I --fix . +ruff format . ``` -If you use git: **do not squash your commits**. We want to see the messy -history — failed attempts, reverts, "trying X" commits. That's the most -honest signal of how you actually worked. A repo with two commits ("initial", -"final") tells us you only committed the AI's output. +The test suite covers ingestion row dropping, metadata parsing, duplicate-content handling, filtered title listing, missing title errors, embedding utilities, and `/ask` source output. --- -## What happens next +## Generated Files -1. You submit the assignment. -2. We read NOTES.md and skim the code. -3. **60–90 minute live walkthrough**, screen-share. We will: - - Ask you to walk through one part of your code we pick at random - - Ask "why" on specific lines - - Ask you to add a small feature live, in 15 minutes (no new prompting to - ChatGPT during this — you can google syntax, that's it) - - Hand you a 30-line Python script we wrote that's broken, and ask you to - debug it in 10 minutes +These files are generated locally and should not be committed: -The live round is where the assignment is actually graded. The code you -submit is the artifact we will discuss; the discussion is what we evaluate. +```text +.env +.venv/ +data/netflix.db +data/chroma_db/ +__pycache__/ +.pytest_cache/ +.ruff_cache/ +``` --- -## Questions before you start +## Known Limitations + +- The dataset is a static Netflix snapshot and should not be treated as current Netflix availability. +- Missing human-facing metadata is stored as `"Unknown"`. +- The Chroma index must be rebuilt with `python -m rag.index` after changing the database, metadata structure, or embedding model. +- The `/ask` endpoint depends on valid Groq and Hugging Face credentials. +- The AI planner can only filter fields that exist in Chroma metadata. -If anything is unclear, ask. We'd rather answer a question than receive a -submission built on a wrong assumption. Asking good clarifying questions is -itself a positive signal. -Good luck. diff --git a/api/__init__.py b/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/ask.py b/api/ask.py new file mode 100644 index 0000000..b33e45f --- /dev/null +++ b/api/ask.py @@ -0,0 +1,34 @@ +"""FastAPI route for RAG-based Netflix catalogue Q&A.""" + +from __future__ import annotations + +from fastapi import APIRouter, HTTPException + +from api.schemas import AskRequest, AskResponse, AskSource +from rag.service import ask_question + +router = APIRouter() + + +@router.post("/ask", response_model=AskResponse) +def ask_catalogue(request: AskRequest) -> AskResponse: + """Answer a natural-language question using retrieved catalogue titles.""" + question = request.question.strip() + + if not question: + raise HTTPException(status_code=400, detail="question must not be empty") + + try: + result = ask_question(question) + except FileNotFoundError as error: + raise HTTPException(status_code=500, detail=str(error)) from error + except RuntimeError as error: + raise HTTPException(status_code=500, detail=str(error)) from error + + return AskResponse( + answer=result.answer, + sources=[ + AskSource(show_id=source.show_id, title=source.title) + for source in result.sources + ], + ) diff --git a/api/database.py b/api/database.py new file mode 100644 index 0000000..167745e --- /dev/null +++ b/api/database.py @@ -0,0 +1,22 @@ +"""Database connection helpers for the Netflix API.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +from ingestion.models import DB_PATH + + +def connect_database(db_path: Path = DB_PATH) -> sqlite3.Connection: + """Open the generated SQLite database with row access by column name.""" + if not db_path.exists(): + raise FileNotFoundError( + f"Database not found at {db_path}. Run `python ingest.py` first." + ) + + connection = sqlite3.connect(db_path) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA foreign_keys = ON") + + return connection diff --git a/api/main.py b/api/main.py new file mode 100644 index 0000000..44c7fee --- /dev/null +++ b/api/main.py @@ -0,0 +1,23 @@ +"""FastAPI application entrypoint for the Netflix catalogue API.""" + +from __future__ import annotations + +from fastapi import FastAPI + +from api.ask import router as ask_router +from api.titles import router as titles_router + +app = FastAPI( + title="Netflix Catalog Q&A System", + description="Structured search and RAG Q&A API for the cleaned Netflix catalogue.", + version="0.1.0", +) + +app.include_router(titles_router) +app.include_router(ask_router) + + +@app.get("/health") +def health_check() -> dict[str, str]: + """Return a basic health check response.""" + return {"status": "ok"} diff --git a/api/schemas.py b/api/schemas.py new file mode 100644 index 0000000..710071f --- /dev/null +++ b/api/schemas.py @@ -0,0 +1,67 @@ +"""Response schemas for the Netflix API.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class TitleResponse(BaseModel): + """A Netflix title returned by the API.""" + + show_id: str + type: str + title: str + release_year: int + rating: str + duration_value: int | None + duration_unit: str | None + date_added: str | None + description: str + countries: list[str] + genres: list[str] + cast: list[str] + directors: list[str] + + +class PaginatedTitlesResponse(BaseModel): + """Paginated list response for /titles.""" + + items: list[TitleResponse] + page: int + page_size: int + total: int + + +class CountryCount(BaseModel): + """Country count item for stats.""" + + country: str + count: int + + +class StatsResponse(BaseModel): + """Summary statistics for the catalogue.""" + + total_titles: int + count_by_type: dict[str, int] + top_countries: list[CountryCount] + + +class AskRequest(BaseModel): + """Request body for the RAG Q&A endpoint.""" + + question: str + + +class AskSource(BaseModel): + """A source title used by the RAG answer.""" + + show_id: str + title: str + + +class AskResponse(BaseModel): + """Response from the RAG Q&A endpoint.""" + + answer: str + sources: list[AskSource] diff --git a/api/titles.py b/api/titles.py new file mode 100644 index 0000000..22b87e9 --- /dev/null +++ b/api/titles.py @@ -0,0 +1,279 @@ +"""FastAPI routes for Netflix catalogue titles and stats.""" + +from __future__ import annotations + +import sqlite3 +from typing import Any + +from fastapi import APIRouter, HTTPException, Query + +from api.database import connect_database +from api.schemas import ( + CountryCount, + PaginatedTitlesResponse, + StatsResponse, + TitleResponse, +) +from ingestion.models import VALID_TYPES + +router = APIRouter() + +CHILD_TABLES = { + "countries": ("title_countries", "country"), + "genres": ("title_genres", "genre"), + "cast": ("title_cast", "actor_name"), + "directors": ("title_directors", "director_name"), +} + + +def open_connection() -> sqlite3.Connection: + """Open the SQLite database or return a clear API error.""" + try: + return connect_database() + except FileNotFoundError as error: + raise HTTPException(status_code=500, detail=str(error)) from error + + +def clean_filter_value(value: str | None) -> str | None: + """Trim optional query filter values.""" + if value is None: + return None + + cleaned = value.strip() + + return cleaned or None + + +def validate_type_filter(type_filter: str | None) -> str | None: + """Validate the optional type query parameter.""" + cleaned = clean_filter_value(type_filter) + + if cleaned is None: + return None + + normalized = cleaned.title().replace("Tv Show", "TV Show") + + if normalized not in VALID_TYPES: + raise HTTPException( + status_code=400, + detail="type must be either 'Movie' or 'TV Show'", + ) + + return normalized + + +def build_title_filters( + *, + country: str | None, + release_year: int | None, + title_type: str | None, + rating: str | None, +) -> tuple[str, list[Any]]: + """Build SQL WHERE conditions and parameters for title filters.""" + conditions = ["1 = 1"] + params: list[Any] = [] + + country = clean_filter_value(country) + rating = clean_filter_value(rating) + + if country is not None: + conditions.append( + """ + EXISTS ( + SELECT 1 + FROM title_countries country_filter + WHERE country_filter.show_id = titles.show_id + AND LOWER(country_filter.country) = LOWER(?) + ) + """ + ) + params.append(country) + + if release_year is not None: + conditions.append("titles.release_year = ?") + params.append(release_year) + + if title_type is not None: + conditions.append("titles.type = ?") + params.append(title_type) + + if rating is not None: + conditions.append("LOWER(titles.rating) = LOWER(?)") + params.append(rating) + + return " AND ".join(conditions), params + + +def fetch_child_values( + connection: sqlite3.Connection, + *, + show_id: str, + table_name: str, + value_column: str, +) -> list[str]: + """Fetch ordered child-table values for a title.""" + rows = connection.execute( + f""" + SELECT {value_column} + FROM {table_name} + WHERE show_id = ? + ORDER BY position + """, + (show_id,), + ).fetchall() + + return [str(row[value_column]) for row in rows] + + +def build_title_response( + connection: sqlite3.Connection, + row: sqlite3.Row, +) -> TitleResponse: + """Build a TitleResponse from a title row and its child-table values.""" + show_id = str(row["show_id"]) + + countries = fetch_child_values( + connection, + show_id=show_id, + table_name=CHILD_TABLES["countries"][0], + value_column=CHILD_TABLES["countries"][1], + ) + genres = fetch_child_values( + connection, + show_id=show_id, + table_name=CHILD_TABLES["genres"][0], + value_column=CHILD_TABLES["genres"][1], + ) + cast = fetch_child_values( + connection, + show_id=show_id, + table_name=CHILD_TABLES["cast"][0], + value_column=CHILD_TABLES["cast"][1], + ) + directors = fetch_child_values( + connection, + show_id=show_id, + table_name=CHILD_TABLES["directors"][0], + value_column=CHILD_TABLES["directors"][1], + ) + + return TitleResponse( + show_id=show_id, + type=str(row["type"]), + title=str(row["title"]), + release_year=int(row["release_year"]), + rating=str(row["rating"]), + duration_value=row["duration_value"], + duration_unit=row["duration_unit"], + date_added=row["date_added"], + description=str(row["description"]), + countries=countries, + genres=genres, + cast=cast, + directors=directors, + ) + + +@router.get("/titles", response_model=PaginatedTitlesResponse) +def list_titles( + country: str | None = None, + release_year: int | None = None, + type_filter: str | None = Query(default=None, alias="type"), + rating: str | None = None, + page: int = Query(default=1, ge=1), + page_size: int = Query(default=20, ge=1, le=100), +) -> PaginatedTitlesResponse: + """List titles with optional filters and pagination.""" + title_type = validate_type_filter(type_filter) + where_sql, params = build_title_filters( + country=country, + release_year=release_year, + title_type=title_type, + rating=rating, + ) + offset = (page - 1) * page_size + + with open_connection() as connection: + total = connection.execute( + f"SELECT COUNT(*) AS count FROM titles WHERE {where_sql}", + params, + ).fetchone()["count"] + + rows = connection.execute( + f""" + SELECT show_id, type, title, release_year, rating, duration_value, + duration_unit, date_added, description + FROM titles + WHERE {where_sql} + ORDER BY title, show_id + LIMIT ? OFFSET ? + """, + [*params, page_size, offset], + ).fetchall() + + items = [build_title_response(connection, row) for row in rows] + + return PaginatedTitlesResponse( + items=items, + page=page, + page_size=page_size, + total=int(total), + ) + + +@router.get("/titles/{show_id}", response_model=TitleResponse) +def get_title(show_id: str) -> TitleResponse: + """Return one title by show_id.""" + with open_connection() as connection: + row = connection.execute( + """ + SELECT show_id, type, title, release_year, rating, duration_value, + duration_unit, date_added, description + FROM titles + WHERE show_id = ? + """, + (show_id,), + ).fetchone() + + if row is None: + raise HTTPException(status_code=404, detail="Title not found") + + return build_title_response(connection, row) + + +@router.get("/stats", response_model=StatsResponse) +def get_stats() -> StatsResponse: + """Return catalogue-level summary statistics.""" + with open_connection() as connection: + total_titles = connection.execute( + "SELECT COUNT(*) AS count FROM titles" + ).fetchone()["count"] + + type_rows = connection.execute( + """ + SELECT type, COUNT(*) AS count + FROM titles + GROUP BY type + ORDER BY type + """ + ).fetchall() + + country_rows = connection.execute( + """ + SELECT country, COUNT(*) AS count + FROM title_countries + WHERE country != 'Unknown' + GROUP BY country + ORDER BY count DESC, country + LIMIT 10 + """ + ).fetchall() + + return StatsResponse( + total_titles=int(total_titles), + count_by_type={str(row["type"]): int(row["count"]) for row in type_rows}, + top_countries=[ + CountryCount(country=str(row["country"]), count=int(row["count"])) + for row in country_rows + ], + ) diff --git a/architecture_diagram.png b/architecture_diagram.png new file mode 100644 index 0000000..d9074c7 Binary files /dev/null and b/architecture_diagram.png differ diff --git a/demo/requirements.txt b/demo/requirements.txt new file mode 100644 index 0000000..d1b43aa --- /dev/null +++ b/demo/requirements.txt @@ -0,0 +1 @@ +streamlit==1.41.1 diff --git a/demo/streamlit_app.py b/demo/streamlit_app.py new file mode 100644 index 0000000..d4f3892 --- /dev/null +++ b/demo/streamlit_app.py @@ -0,0 +1,282 @@ +"""Optional Streamlit demo UI for the Netflix Catalog API.""" + +from __future__ import annotations + +import json +import os +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.parse import urlencode +from urllib.request import Request, urlopen + +import streamlit as st + +DEFAULT_API_BASE_URL = os.environ.get("API_BASE_URL", "http://127.0.0.1:8000") +DEFAULT_PAGE_SIZE = 5 + + +def call_api( + api_base_url: str, + path: str, + *, + method: str = "GET", + payload: dict[str, Any] | None = None, + query_params: dict[str, Any] | None = None, +) -> tuple[dict[str, Any] | None, str | None]: + """Call the FastAPI backend and return JSON data or an error message.""" + url = f"{api_base_url.rstrip('/')}{path}" + + if query_params: + clean_params = { + key: value + for key, value in query_params.items() + if value is not None and value != "" + } + + if clean_params: + url = f"{url}?{urlencode(clean_params)}" + + body = None + headers = {"Accept": "application/json"} + + if payload is not None: + body = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + + request = Request(url, data=body, headers=headers, method=method) + + try: + with urlopen(request, timeout=90) as response: + response_body = response.read().decode("utf-8") + return json.loads(response_body), None + except HTTPError as error: + detail = error.read().decode("utf-8") + return None, f"API error {error.code}: {detail}" + except URLError as error: + return None, f"Could not connect to API server: {error.reason}" + except TimeoutError: + return None, "API request timed out." + + +def show_health(api_base_url: str) -> None: + """Show the /health endpoint result.""" + st.subheader("Health Check") + + if st.button("Check API"): + data, error = call_api(api_base_url, "/health") + + if error: + st.error(error) + return + + st.success("API is running.") + st.json(data) + + +def show_titles(api_base_url: str) -> None: + """Show the /titles endpoint with required filters and pagination.""" + st.subheader("List Titles") + + st.write( + "This calls `GET /titles` with the required filters: " + "`country`, `release_year`, `type`, `rating`, and pagination." + ) + + col1, col2 = st.columns(2) + + with col1: + country = st.text_input("Country", value="India") + release_year = st.text_input("Release year", value="") + title_type = st.selectbox("Type", ["", "Movie", "TV Show"]) + + with col2: + rating = st.text_input("Rating", value="") + page = st.number_input("Result page", min_value=1, value=1) + page_size = st.number_input( + "Results per page", + min_value=1, + max_value=100, + value=DEFAULT_PAGE_SIZE, + ) + + if st.button("Search titles"): + data, error = call_api( + api_base_url, + "/titles", + query_params={ + "country": country, + "release_year": release_year, + "type": title_type, + "rating": rating, + "page": page, + "page_size": page_size, + }, + ) + + if error: + st.error(error) + return + + if data is None: + st.warning("No response returned.") + return + + st.write(f"Total matches: **{data['total']}**") + st.write(f"Showing result page **{data['page']}** with **{data['page_size']}**") + + if not data["items"]: + st.info("No titles found.") + return + + for item in data["items"]: + with st.expander(f"{item['title']} ({item['release_year']})"): + st.write(f"**show_id:** {item['show_id']}") + st.write(f"**Type:** {item['type']}") + st.write(f"**Rating:** {item['rating']}") + st.write( + f"**Duration:** {item['duration_value']} {item['duration_unit']}" + ) + st.write(f"**Date added:** {item['date_added']}") + st.write(f"**Countries:** {', '.join(item['countries'])}") + st.write(f"**Genres:** {', '.join(item['genres'])}") + st.write(f"**Directors:** {', '.join(item['directors'])}") + st.write(f"**Cast:** {', '.join(item['cast'][:10])}") + st.write(item["description"]) + + +def show_title_by_id(api_base_url: str) -> None: + """Show the /titles/{show_id} endpoint.""" + st.subheader("Fetch Title by ID") + + st.write("This calls `GET /titles/{show_id}`.") + + show_id = st.text_input("show_id", value="81075235") + + if st.button("Fetch title"): + data, error = call_api(api_base_url, f"/titles/{show_id}") + + if error: + st.error(error) + return + + if data is None: + st.warning("No response returned.") + return + + st.json(data) + + +def show_stats(api_base_url: str) -> None: + """Show the /stats endpoint.""" + st.subheader("Catalogue Stats") + + st.write("This calls `GET /stats`.") + + if st.button("Load stats"): + data, error = call_api(api_base_url, "/stats") + + if error: + st.error(error) + return + + if data is None: + st.warning("No response returned.") + return + + st.metric("Total titles", data["total_titles"]) + + st.markdown("### Count by type") + st.json(data["count_by_type"]) + + st.markdown("### Top countries") + st.dataframe(data["top_countries"], use_container_width=True) + + +def show_ask(api_base_url: str) -> None: + """Show the /ask RAG endpoint.""" + st.subheader("Ask the Catalogue") + + st.write("This calls `POST /ask`.") + + question = st.text_area( + "Question", + value="Suggest me an Indian comedy movie", + height=120, + ) + + if st.button("Ask"): + data, error = call_api( + api_base_url, + "/ask", + method="POST", + payload={"question": question}, + ) + + if error: + st.error(error) + return + + if data is None: + st.warning("No response returned.") + return + + st.markdown("### Answer") + st.write(data["answer"]) + + st.markdown("### Sources") + if data["sources"]: + st.dataframe(data["sources"], use_container_width=True) + else: + st.info("No sources returned.") + + +def main() -> None: + """Run the optional Streamlit demo app.""" + st.set_page_config(page_title="Netflix Catalog Demo", layout="wide") + + st.title("Netflix Catalog Q&A Demo") + st.write( + "This optional demo calls the FastAPI backend endpoints. " + "Start the API server before using this UI." + ) + + api_base_url = st.sidebar.text_input( + "API base URL", + value=DEFAULT_API_BASE_URL, + ) + + st.sidebar.markdown("### Start backend first") + st.sidebar.code( + "python ingest.py\n" + "python -m rag.index\n" + "uvicorn api.main:app --reload" + ) + + tabs = st.tabs( + [ + "Health", + "List Titles", + "Title by ID", + "Stats", + "Ask", + ] + ) + + with tabs[0]: + show_health(api_base_url) + + with tabs[1]: + show_titles(api_base_url) + + with tabs[2]: + show_title_by_id(api_base_url) + + with tabs[3]: + show_stats(api_base_url) + + with tabs[4]: + show_ask(api_base_url) + + +if __name__ == "__main__": + main() diff --git a/ingest.py b/ingest.py new file mode 100644 index 0000000..19551a6 --- /dev/null +++ b/ingest.py @@ -0,0 +1,6 @@ +"""Command-line entrypoint for Netflix catalogue ingestion.""" + +from ingestion.pipeline import run_ingestion + +if __name__ == "__main__": + run_ingestion() diff --git a/ingestion/__init__.py b/ingestion/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ingestion/cleaning.py b/ingestion/cleaning.py new file mode 100644 index 0000000..5cb7e39 --- /dev/null +++ b/ingestion/cleaning.py @@ -0,0 +1,213 @@ +"""Cleaning helpers for raw Netflix CSV rows.""" + +from __future__ import annotations + +import re +from datetime import datetime +from typing import Any + +from ingestion.models import UNKNOWN_VALUE, VALID_TYPES, CleanTitle, IngestStats + +DATE_FORMAT = "%B %d, %Y" +DURATION_RE = re.compile(r"^(?P\d+)\s+(?Pmin|Season|Seasons)$") +RATING_REPLACEMENTS = { + "NR": "Not Rated", + "UR": "Unrated", +} + + +def clean_text(value: Any, *, field_name: str, stats: IngestStats) -> str | None: + """Return a cleaned string, or None for missing/blank values.""" + if value is None: + stats.missing_values[field_name] += 1 + return None + + raw = str(value) + cleaned = raw.replace("\xa0", " ") + cleaned = cleaned.replace("\n", " ").replace("\r", " ").replace("\t", " ") + cleaned = " ".join(cleaned.split()) + + if raw != cleaned: + stats.fixes[f"cleaned_{field_name}"] += 1 + + if cleaned == "": + stats.missing_values[field_name] += 1 + return None + + return cleaned + + +def clean_metadata(value: Any, *, field_name: str, stats: IngestStats) -> str: + """Clean optional display metadata and use Unknown when it is missing.""" + cleaned = clean_text(value, field_name=field_name, stats=stats) + + if cleaned is None: + return UNKNOWN_VALUE + + return cleaned + + +def dedupe_preserve_order(values: list[str]) -> list[str]: + """Remove duplicate list values while keeping the CSV order.""" + seen: set[str] = set() + deduped: list[str] = [] + + for value in values: + key = value.casefold() + + if key in seen: + continue + + seen.add(key) + deduped.append(value) + + return deduped + + +def split_list(value: str, *, field_name: str, stats: IngestStats) -> list[str]: + """Split a comma-separated metadata field into ordered unique values.""" + if value == UNKNOWN_VALUE: + return [UNKNOWN_VALUE] + + parts = [part.strip() for part in value.split(",")] + values = [part for part in parts if part] + + empty_count = len(parts) - len(values) + if empty_count: + stats.fixes[f"removed_empty_{field_name}_items"] += empty_count + + deduped_values = dedupe_preserve_order(values) + + removed_duplicate_count = len(values) - len(deduped_values) + if removed_duplicate_count: + stats.fixes[f"removed_duplicate_{field_name}_items"] += removed_duplicate_count + + return deduped_values or [UNKNOWN_VALUE] + + +def normalize_type(value: Any, stats: IngestStats) -> str | None: + """Normalize the title type to Movie or TV Show.""" + cleaned = clean_text(value, field_name="type", stats=stats) + + if cleaned is None: + return None + + normalized = cleaned.title().replace("Tv Show", "TV Show") + + if normalized not in VALID_TYPES: + stats.anomalies["invalid_type"] += 1 + return None + + return normalized + + +def normalize_rating(value: Any, stats: IngestStats) -> str: + """Clean rating and expand unclear rating abbreviations.""" + rating = clean_metadata(value, field_name="rating", stats=stats) + normalized_rating = RATING_REPLACEMENTS.get(rating, rating) + + if normalized_rating != rating: + stats.fixes["normalized_rating"] += 1 + + return normalized_rating + + +def parse_release_year(value: Any, stats: IngestStats) -> int | None: + """Parse release_year as an integer.""" + cleaned = clean_text(value, field_name="release_year", stats=stats) + + if cleaned is None: + return None + + try: + return int(cleaned) + except ValueError: + stats.anomalies["invalid_release_year"] += 1 + return None + + +def parse_date_added(value: Any, stats: IngestStats) -> str | None: + """Parse date_added into YYYY-MM-DD, leaving missing dates as None.""" + cleaned = clean_text(value, field_name="date_added", stats=stats) + + if cleaned is None: + return None + + try: + return datetime.strptime(cleaned, DATE_FORMAT).date().isoformat() + except ValueError: + stats.anomalies["invalid_date_added"] += 1 + return None + + +def parse_duration(value: Any, stats: IngestStats) -> tuple[int | None, str | None]: + """Parse duration into a numeric value and normalized unit.""" + cleaned = clean_text(value, field_name="duration", stats=stats) + + if cleaned is None: + return None, None + + match = DURATION_RE.match(cleaned) + + if match is None: + stats.anomalies["invalid_duration"] += 1 + return None, None + + raw_unit = match.group("unit") + unit = "minutes" if raw_unit == "min" else "seasons" + + return int(match.group("value")), unit + + +def clean_row(row: dict[str, str], stats: IngestStats) -> CleanTitle | None: + """Clean one raw CSV row and return None when the row is unusable.""" + show_id = clean_text(row.get("show_id"), field_name="show_id", stats=stats) + title_type = normalize_type(row.get("type"), stats) + title = clean_text(row.get("title"), field_name="title", stats=stats) + release_year = parse_release_year(row.get("release_year"), stats) + description = clean_text( + row.get("description"), field_name="description", stats=stats + ) + + missing_required = { + "show_id": show_id, + "type": title_type, + "title": title, + "release_year": release_year, + "description": description, + } + + for field_name, cleaned_value in missing_required.items(): + if cleaned_value is None: + stats.rows_dropped += 1 + stats.dropped_reasons[f"missing_or_invalid_{field_name}"] += 1 + return None + + director_text = clean_metadata( + row.get("director"), field_name="director", stats=stats + ) + cast_text = clean_metadata(row.get("cast"), field_name="cast", stats=stats) + country_text = clean_metadata(row.get("country"), field_name="country", stats=stats) + rating = normalize_rating(row.get("rating"), stats) + genre_text = clean_metadata( + row.get("listed_in"), field_name="listed_in", stats=stats + ) + + date_added = parse_date_added(row.get("date_added"), stats) + duration_value, duration_unit = parse_duration(row.get("duration"), stats) + + return CleanTitle( + show_id=show_id, + title_type=title_type, + title=title, + release_year=release_year, + rating=rating, + duration_value=duration_value, + duration_unit=duration_unit, + date_added=date_added, + description=description, + countries=split_list(country_text, field_name="country", stats=stats), + genres=split_list(genre_text, field_name="listed_in", stats=stats), + cast_members=split_list(cast_text, field_name="cast", stats=stats), + directors=split_list(director_text, field_name="director", stats=stats), + ) diff --git a/ingestion/database.py b/ingestion/database.py new file mode 100644 index 0000000..e693cfb --- /dev/null +++ b/ingestion/database.py @@ -0,0 +1,160 @@ +"""SQLite schema and insert logic for the cleaned Netflix catalogue.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +from ingestion.models import CleanTitle + + +def create_schema(connection: sqlite3.Connection) -> None: + """Create all ingestion tables from scratch.""" + connection.executescript( + """ + DROP TABLE IF EXISTS title_directors; + DROP TABLE IF EXISTS title_cast; + DROP TABLE IF EXISTS title_genres; + DROP TABLE IF EXISTS title_countries; + DROP TABLE IF EXISTS titles; + + CREATE TABLE titles ( + show_id TEXT NOT NULL PRIMARY KEY, + type TEXT NOT NULL CHECK (type IN ('Movie', 'TV Show')), + title TEXT NOT NULL, + release_year INTEGER NOT NULL, + rating TEXT NOT NULL, + duration_value INTEGER, + duration_unit TEXT CHECK ( + duration_unit IS NULL OR duration_unit IN ('minutes', 'seasons') + ), + date_added TEXT, + description TEXT NOT NULL + ); + + CREATE TABLE title_countries ( + show_id TEXT NOT NULL REFERENCES titles(show_id) ON DELETE CASCADE, + country TEXT NOT NULL, + position INTEGER NOT NULL, + PRIMARY KEY (show_id, position) + ); + + CREATE TABLE title_genres ( + show_id TEXT NOT NULL REFERENCES titles(show_id) ON DELETE CASCADE, + genre TEXT NOT NULL, + position INTEGER NOT NULL, + PRIMARY KEY (show_id, position) + ); + + CREATE TABLE title_cast ( + show_id TEXT NOT NULL REFERENCES titles(show_id) ON DELETE CASCADE, + actor_name TEXT NOT NULL, + position INTEGER NOT NULL, + PRIMARY KEY (show_id, position) + ); + + CREATE TABLE title_directors ( + show_id TEXT NOT NULL REFERENCES titles(show_id) ON DELETE CASCADE, + director_name TEXT NOT NULL, + position INTEGER NOT NULL, + PRIMARY KEY (show_id, position) + ); + + CREATE INDEX idx_titles_type ON titles(type); + CREATE INDEX idx_titles_release_year ON titles(release_year); + CREATE INDEX idx_titles_rating ON titles(rating); + CREATE INDEX idx_title_countries_country ON title_countries(country); + CREATE INDEX idx_title_genres_genre ON title_genres(genre); + """ + ) + + +def insert_child_values( + connection: sqlite3.Connection, + *, + table_name: str, + value_column: str, + show_id: str, + values: list[str], +) -> None: + """Insert ordered values for one child table.""" + rows = [(show_id, value, index) for index, value in enumerate(values, start=1)] + + connection.executemany( + f"INSERT INTO {table_name} (show_id, {value_column}, position) VALUES (?, ?, ?)", + rows, + ) + + +def insert_title(connection: sqlite3.Connection, title: CleanTitle) -> None: + """Insert one cleaned title and its child-table values.""" + connection.execute( + """ + INSERT INTO titles ( + show_id, + type, + title, + release_year, + rating, + duration_value, + duration_unit, + date_added, + description + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + title.show_id, + title.title_type, + title.title, + title.release_year, + title.rating, + title.duration_value, + title.duration_unit, + title.date_added, + title.description, + ), + ) + + insert_child_values( + connection, + table_name="title_countries", + value_column="country", + show_id=title.show_id, + values=title.countries, + ) + insert_child_values( + connection, + table_name="title_genres", + value_column="genre", + show_id=title.show_id, + values=title.genres, + ) + insert_child_values( + connection, + table_name="title_cast", + value_column="actor_name", + show_id=title.show_id, + values=title.cast_members, + ) + insert_child_values( + connection, + table_name="title_directors", + value_column="director_name", + show_id=title.show_id, + values=title.directors, + ) + + +def write_database(titles: list[CleanTitle], db_path: Path) -> None: + """Write cleaned titles into a fresh SQLite database.""" + db_path.parent.mkdir(parents=True, exist_ok=True) + + with sqlite3.connect(db_path) as connection: + connection.execute("PRAGMA foreign_keys = ON") + create_schema(connection) + + for title in titles: + insert_title(connection, title) + + connection.commit() diff --git a/ingestion/duplicates.py b/ingestion/duplicates.py new file mode 100644 index 0000000..87738f7 --- /dev/null +++ b/ingestion/duplicates.py @@ -0,0 +1,47 @@ +"""Duplicate-content detection for cleaned Netflix titles.""" + +from __future__ import annotations + +from typing import Any + +from ingestion.models import CleanTitle, IngestStats + + +def content_key(title: CleanTitle) -> tuple[Any, ...]: + """Return all title content except show_id for duplicate comparison.""" + return ( + title.title_type, + title.title, + title.release_year, + title.rating, + title.duration_value, + title.duration_unit, + title.date_added, + title.description, + tuple(title.countries), + tuple(title.genres), + tuple(title.cast_members), + tuple(title.directors), + ) + + +def drop_duplicate_content( + titles: list[CleanTitle], + stats: IngestStats, +) -> list[CleanTitle]: + """Drop rows whose cleaned content matches a previous row except show_id.""" + seen: set[tuple[Any, ...]] = set() + unique_titles: list[CleanTitle] = [] + + for title in titles: + key = content_key(title) + + if key in seen: + stats.rows_dropped += 1 + stats.dropped_reasons["duplicate_content_except_show_id"] += 1 + continue + + seen.add(key) + unique_titles.append(title) + + return unique_titles diff --git a/ingestion/models.py b/ingestion/models.py new file mode 100644 index 0000000..0ce09f6 --- /dev/null +++ b/ingestion/models.py @@ -0,0 +1,60 @@ +"""Shared models and constants for Netflix catalogue ingestion.""" + +from __future__ import annotations + +import os +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path + +CSV_PATH = Path(os.environ.get("NETFLIX_CSV_PATH", "data/netflix_titles.csv")) +DB_PATH = Path(os.environ.get("NETFLIX_DB_PATH", "data/netflix.db")) +UNKNOWN_VALUE = "Unknown" +VALID_TYPES = {"Movie", "TV Show"} + +EXPECTED_COLUMNS = [ + "show_id", + "type", + "title", + "director", + "cast", + "country", + "date_added", + "release_year", + "rating", + "duration", + "listed_in", + "description", +] + + +@dataclass +class IngestStats: + """Counts collected during ingestion for the final summary report.""" + + rows_read: int = 0 + rows_loaded: int = 0 + rows_dropped: int = 0 + dropped_reasons: Counter[str] = field(default_factory=Counter) + missing_values: Counter[str] = field(default_factory=Counter) + fixes: Counter[str] = field(default_factory=Counter) + anomalies: Counter[str] = field(default_factory=Counter) + + +@dataclass +class CleanTitle: + """Cleaned title data ready to be inserted into SQLite.""" + + show_id: str + title_type: str + title: str + release_year: int + rating: str + duration_value: int | None + duration_unit: str | None + date_added: str | None + description: str + countries: list[str] + genres: list[str] + cast_members: list[str] + directors: list[str] diff --git a/ingestion/pipeline.py b/ingestion/pipeline.py new file mode 100644 index 0000000..45a6805 --- /dev/null +++ b/ingestion/pipeline.py @@ -0,0 +1,99 @@ +"""End-to-end ingestion pipeline for the Netflix catalogue.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +from ingestion.cleaning import clean_row +from ingestion.database import write_database +from ingestion.duplicates import drop_duplicate_content +from ingestion.models import ( + CSV_PATH, + DB_PATH, + EXPECTED_COLUMNS, + CleanTitle, + IngestStats, +) + + +def validate_columns(columns: list[str] | None) -> None: + """Fail early if the CSV does not have the expected header.""" + if columns is None: + raise ValueError("CSV file has no header row.") + + missing_columns = set(EXPECTED_COLUMNS) - set(columns) + + if missing_columns: + raise ValueError(f"Missing expected columns: {sorted(missing_columns)}") + + +def load_clean_titles(csv_path: Path) -> tuple[list[CleanTitle], IngestStats]: + """Load, clean, validate, and deduplicate title rows from the CSV.""" + stats = IngestStats() + titles: list[CleanTitle] = [] + seen_show_ids: set[str] = set() + + with csv_path.open(newline="", encoding="utf-8") as csv_file: + reader = csv.DictReader(csv_file) + validate_columns(reader.fieldnames) + + for row in reader: + stats.rows_read += 1 + title = clean_row(row, stats) + + if title is None: + continue + + if title.show_id in seen_show_ids: + stats.rows_dropped += 1 + stats.dropped_reasons["duplicate_show_id"] += 1 + continue + + seen_show_ids.add(title.show_id) + titles.append(title) + + titles = drop_duplicate_content(titles, stats) + stats.rows_loaded = len(titles) + + return titles, stats + + +def print_summary(stats: IngestStats, db_path: Path) -> None: + """Print a human-readable ingestion summary.""" + print("Ingestion complete") + print("------------------") + print(f"Database written: {db_path}") + print(f"Rows read: {stats.rows_read}") + print(f"Rows loaded: {stats.rows_loaded}") + print(f"Rows dropped: {stats.rows_dropped}") + + if stats.dropped_reasons: + print("\nDropped rows by reason:") + for reason, count in stats.dropped_reasons.most_common(): + print(f" - {reason}: {count}") + + if stats.missing_values: + print("\nMissing values handled:") + for field_name, count in stats.missing_values.most_common(): + print(f" - {field_name}: {count}") + + if stats.fixes: + print("\nCleaning fixes applied:") + for fix_name, count in stats.fixes.most_common(): + print(f" - {fix_name}: {count}") + + if stats.anomalies: + print("\nAnomalies observed:") + for anomaly, count in stats.anomalies.most_common(): + print(f" - {anomaly}: {count}") + + +def run_ingestion(csv_path: Path = CSV_PATH, db_path: Path = DB_PATH) -> None: + """Run the full CSV-to-SQLite ingestion flow.""" + if not csv_path.exists(): + raise FileNotFoundError(f"CSV file not found: {csv_path}") + + titles, stats = load_clean_titles(csv_path) + write_database(titles, db_path) + print_summary(stats, db_path) diff --git a/rag/__init__.py b/rag/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rag/config.py b/rag/config.py new file mode 100644 index 0000000..b1efba8 --- /dev/null +++ b/rag/config.py @@ -0,0 +1,21 @@ +"""Configuration values for the RAG pipeline.""" + +from __future__ import annotations + +import os +from pathlib import Path + +from dotenv import load_dotenv + +load_dotenv() + +CHROMA_PATH = Path(os.environ.get("CHROMA_PATH", "data/chroma_db")) +CHROMA_COLLECTION_NAME = os.environ.get("CHROMA_COLLECTION_NAME", "netflix_titles") +EMBEDDING_MODEL_NAME = os.environ.get( + "EMBEDDING_MODEL_NAME", + "sentence-transformers/all-MiniLM-L6-v2", +) +HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "") +GROQ_MODEL = os.environ.get("GROQ_MODEL", "llama-3.3-70b-versatile") +RAG_BATCH_SIZE = int(os.environ.get("RAG_BATCH_SIZE", "128")) +RAG_TOP_K = int(os.environ.get("RAG_TOP_K", "5")) \ No newline at end of file diff --git a/rag/documents.py b/rag/documents.py new file mode 100644 index 0000000..87bb932 --- /dev/null +++ b/rag/documents.py @@ -0,0 +1,143 @@ +"""Build searchable RAG documents from the cleaned Netflix database.""" + +from __future__ import annotations + +import sqlite3 +from dataclasses import dataclass + +from api.database import connect_database + + +@dataclass +class TitleDocument: + """A catalogue title represented as text plus metadata for retrieval.""" + + show_id: str + title: str + text: str + metadata: dict[str, str | int | float | bool | list[str] | None] + + +def fetch_child_values( + connection: sqlite3.Connection, + *, + table_name: str, + value_column: str, + show_id: str, +) -> list[str]: + """Fetch ordered child-table values for one title.""" + rows = connection.execute( + f""" + SELECT {value_column} + FROM {table_name} + WHERE show_id = ? + ORDER BY position + """, + (show_id,), + ).fetchall() + + return [str(row[value_column]) for row in rows] + + +def build_document_text( + *, + row: sqlite3.Row, + countries: list[str], + genres: list[str], + cast: list[str], + directors: list[str], +) -> str: + """Build the text that will be embedded for one Netflix title.""" + duration = "Unknown" + if row["duration_value"] is not None and row["duration_unit"] is not None: + duration = f"{row['duration_value']} {row['duration_unit']}" + + parts = [ + f"Title: {row['title']}", + f"Type: {row['type']}", + f"Release year: {row['release_year']}", + f"Rating: {row['rating']}", + f"Duration: {duration}", + f"Countries: {', '.join(countries)}", + f"Genres: {', '.join(genres)}", + f"Directors: {', '.join(directors)}", + f"Cast: {', '.join(cast)}", + f"Description: {row['description']}", + ] + + return "\n".join(parts) + +def ensure_non_empty_metadata_list(values: list[str]) -> list[str]: + """Return a non-empty metadata list for Chroma filters.""" + clean_values = [value.strip() for value in values if value.strip()] + return clean_values or ["Unknown"] + +def load_title_documents() -> list[TitleDocument]: + """Load all cleaned titles from SQLite and convert them to RAG documents.""" + documents: list[TitleDocument] = [] + + with connect_database() as connection: + rows = connection.execute( + """ + SELECT show_id, type, title, release_year, rating, duration_value, + duration_unit, date_added, description + FROM titles + ORDER BY show_id + """ + ).fetchall() + + for row in rows: + show_id = str(row["show_id"]) + countries = fetch_child_values( + connection, + table_name="title_countries", + value_column="country", + show_id=show_id, + ) + genres = fetch_child_values( + connection, + table_name="title_genres", + value_column="genre", + show_id=show_id, + ) + cast = fetch_child_values( + connection, + table_name="title_cast", + value_column="actor_name", + show_id=show_id, + ) + directors = fetch_child_values( + connection, + table_name="title_directors", + value_column="director_name", + show_id=show_id, + ) + + text = build_document_text( + row=row, + countries=countries, + genres=genres, + cast=cast, + directors=directors, + ) + metadata_countries = ensure_non_empty_metadata_list(countries) + metadata_genres = ensure_non_empty_metadata_list(genres) + + documents.append( + TitleDocument( + show_id=show_id, + title=str(row["title"]), + text=text, + metadata={ + "show_id": show_id, + "title": str(row["title"]), + "type": str(row["type"]), + "release_year": int(row["release_year"]), + "rating": str(row["rating"]), + "countries": metadata_countries, + "genres": metadata_genres, + }, + ) + ) + + return documents diff --git a/rag/embeddings.py b/rag/embeddings.py new file mode 100644 index 0000000..16b91e1 --- /dev/null +++ b/rag/embeddings.py @@ -0,0 +1,212 @@ +"""Shared embedding utilities using the Hugging Face Inference API.""" + +from __future__ import annotations + +import logging +import math +import time +from typing import Any + +from rag.config import EMBEDDING_MODEL_NAME, HF_API_TOKEN + +logger = logging.getLogger(__name__) + +MAX_RETRIES = 3 +RETRY_DELAY_SECONDS = 2 + + +def get_hf_client() -> Any: + """Create a Hugging Face InferenceClient with token validation.""" + if not HF_API_TOKEN: + raise RuntimeError("HF_API_TOKEN is not set. Add it to your .env file.") + + from huggingface_hub import InferenceClient + + return InferenceClient(token=HF_API_TOKEN) + + +def normalize_embedding(vector: list[float]) -> list[float]: + """Normalize a vector to unit length for cosine similarity. + + This matches the behavior of SentenceTransformer's + normalize_embeddings=True parameter so cosine distance + calculations in Chroma remain consistent. + """ + magnitude = math.sqrt(sum(x * x for x in vector)) + + if magnitude == 0: + return vector + + return [x / magnitude for x in vector] + + +def _to_flat_list(value: Any) -> list[float]: + """Convert a numpy array or nested list to a flat Python list of floats.""" + if hasattr(value, "tolist"): + value = value.tolist() + + return [float(x) for x in value] + + +def _extract_single_embedding(raw: Any) -> list[float]: + """Extract a single embedding vector from HF API output. + + Hugging Face feature_extraction can return different shapes: + - [384] → single vector (sentence-level) + - [[384]] → single vector wrapped in a list + - [[tok1], [tok2], ...] → per-token embeddings (need mean pooling) + + This function handles all cases and always returns a flat list. + """ + if hasattr(raw, "tolist"): + raw = raw.tolist() + + if not raw: + raise RuntimeError("Hugging Face returned an empty embedding.") + + first_element = raw[0] + + if isinstance(first_element, (int, float)): + return _to_flat_list(raw) + + if isinstance(first_element, list): + if len(raw) == 1: + return _to_flat_list(raw[0]) + + dimension = len(raw[0]) + pooled = [0.0] * dimension + + for token_vector in raw: + for i, value in enumerate(token_vector): + pooled[i] += float(value) + + token_count = len(raw) + return [value / token_count for value in pooled] + + raise RuntimeError(f"Unexpected embedding format: {type(first_element)}") + +def _extract_batch_embeddings(raw: Any, *, expected_count: int) -> list[list[float]]: + """Extract multiple embedding vectors from HF batch output. + + Possible batch shapes: + - [[384], [384], ...] + - [[[token_dim], [token_dim]], [[token_dim], [token_dim]], ...] + + Each item in the outer list should belong to one input text. + """ + if hasattr(raw, "tolist"): + raw = raw.tolist() + + if not raw: + raise RuntimeError("Hugging Face returned an empty batch embedding.") + + if expected_count == 1: + return [_extract_single_embedding(raw)] + + if not isinstance(raw, list): + raise RuntimeError(f"Unexpected batch embedding format: {type(raw)}") + + if len(raw) != expected_count: + raise RuntimeError( + f"Batch embedding count mismatch: expected {expected_count}, got {len(raw)}" + ) + + return [_extract_single_embedding(item) for item in raw] + + +def _validate_embedding(vector: list[float], *, context: str) -> None: + """Raise an error if the embedding is empty or invalid.""" + if not vector: + raise RuntimeError(f"Empty embedding returned for {context}.") + + if not all(math.isfinite(value) for value in vector): + raise RuntimeError(f"Invalid embedding values returned for {context}.") + +def _validate_embedding_dimensions(embeddings: list[list[float]]) -> int: + """Ensure all embeddings have the same dimension and return that dimension.""" + if not embeddings: + raise RuntimeError("No embeddings were generated.") + + expected_dimension = len(embeddings[0]) + + for index, vector in enumerate(embeddings): + _validate_embedding(vector, context=f"text at index {index}") + + if len(vector) != expected_dimension: + raise RuntimeError( + f"Dimension mismatch at index {index}: " + f"expected {expected_dimension}, got {len(vector)}." + ) + + return expected_dimension + + +def _call_hf_api(client: Any, inputs: str | list[str]) -> Any: + """Call the HF feature_extraction API with retry logic.""" + last_error: Exception | None = None + + for attempt in range(1, MAX_RETRIES + 1): + try: + return client.feature_extraction( + inputs, + model=EMBEDDING_MODEL_NAME, + ) + except Exception as error: + last_error = error + logger.warning( + "HF API attempt %d/%d failed: %s", attempt, MAX_RETRIES, error + ) + + if attempt < MAX_RETRIES: + time.sleep(RETRY_DELAY_SECONDS) + + raise RuntimeError( + f"Hugging Face API failed after {MAX_RETRIES} attempts: {last_error}" + ) + + +def embed_one_text(client: Any, text: str) -> list[float]: + """Embed a single text and return a normalized vector.""" + raw = _call_hf_api(client, text) + vector = _extract_single_embedding(raw) + _validate_embedding(vector, context="single text") + + return normalize_embedding(vector) + + +def embed_texts(client: Any, texts: list[str]) -> list[list[float]]: + """Embed multiple texts and return normalized vectors. + + Tries batch embedding first for speed. + Falls back to one-by-one embedding if batch mode fails. + """ + if not texts: + return [] + + try: + raw_batch = _call_hf_api(client, texts) + embeddings = [ + normalize_embedding(vector) + for vector in _extract_batch_embeddings( + raw_batch, + expected_count=len(texts), + ) + ] + + _validate_embedding_dimensions(embeddings) + return embeddings + + except Exception as error: + logger.warning( + "Batch embedding failed, falling back to one-by-one mode: %s", + error, + ) + + embeddings: list[list[float]] = [] + + for index, text in enumerate(texts): + vector = embed_one_text(client, text) + embeddings.append(vector) + + _validate_embedding_dimensions(embeddings) + return embeddings diff --git a/rag/index.py b/rag/index.py new file mode 100644 index 0000000..c7da767 --- /dev/null +++ b/rag/index.py @@ -0,0 +1,119 @@ +"""Build a persistent Chroma vector index for Netflix title documents.""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path +from typing import TypeVar + +import chromadb + +from rag.config import ( + CHROMA_COLLECTION_NAME, + CHROMA_PATH, + EMBEDDING_MODEL_NAME, + RAG_BATCH_SIZE, +) +from rag.documents import TitleDocument, load_title_documents +from rag.embeddings import get_hf_client, embed_texts + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +def reset_chroma_directory(path: Path) -> None: + """Remove an old local Chroma index before rebuilding.""" + if path.exists(): + shutil.rmtree(path) + + path.mkdir(parents=True, exist_ok=True) + + +def batched(items: list[T], batch_size: int) -> list[list[T]]: + """Split a list into fixed-size batches.""" + return [ + items[index : index + batch_size] for index in range(0, len(items), batch_size) + ] + + +def add_documents_to_collection( + collection: chromadb.Collection, + hf_client: object, + documents: list[TitleDocument], +) -> None: + """Embed documents in batches and add them to Chroma.""" + expected_dimension: int | None = None + + for batch_number, document_batch in enumerate( + batched(documents, RAG_BATCH_SIZE), start=1 + ): + texts = [document.text for document in document_batch] + embeddings = embed_texts(hf_client, texts) + + if not embeddings: + raise RuntimeError(f"No embeddings generated for batch {batch_number}.") + + batch_dimension = len(embeddings[0]) + + if expected_dimension is None: + expected_dimension = batch_dimension + + collection.modify( + metadata={ + "embedding_model": EMBEDDING_MODEL_NAME, + "embedding_dimension": expected_dimension, + } + ) + + logger.info( + "Embedding metadata saved: model=%s, dimension=%d", + EMBEDDING_MODEL_NAME, + expected_dimension, + ) + + elif batch_dimension != expected_dimension: + raise RuntimeError( + f"Embedding dimension changed in batch {batch_number}: " + f"expected {expected_dimension}, got {batch_dimension}." + ) + + collection.add( + ids=[document.show_id for document in document_batch], + documents=texts, + metadatas=[document.metadata for document in document_batch], + embeddings=embeddings, + ) + + logger.info( + "Batch %d indexed (%d documents)", batch_number, len(document_batch) + ) + + +def build_index() -> None: + """Embed all title documents and save them in a local Chroma collection.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + documents = load_title_documents() + + if not documents: + raise ValueError("No title documents found. Run `python ingest.py` first.") + + reset_chroma_directory(CHROMA_PATH) + + hf_client = get_hf_client() + client = chromadb.PersistentClient(path=str(CHROMA_PATH)) + collection = client.get_or_create_collection(name=CHROMA_COLLECTION_NAME) + + add_documents_to_collection(collection, hf_client, documents) + + logger.info("RAG index built") + logger.info("Documents indexed: %s", len(documents)) + logger.info("Chroma path: %s", CHROMA_PATH) + logger.info("Collection: %s", CHROMA_COLLECTION_NAME) + logger.info("Embedding model: %s", EMBEDDING_MODEL_NAME) + + +if __name__ == "__main__": + build_index() diff --git a/rag/llm.py b/rag/llm.py new file mode 100644 index 0000000..0a252d3 --- /dev/null +++ b/rag/llm.py @@ -0,0 +1,178 @@ +"""Groq LLM client for grounded Netflix catalogue answers.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass + +from dotenv import load_dotenv + +from rag.config import GROQ_MODEL +from rag.retriever import RetrievedTitle + +SYSTEM_MESSAGE = """ +You are a Netflix catalogue Q&A assistant. + +Rules: +- Use only the provided catalogue sources. +- Do not use outside knowledge. +- Do not invent titles, actors, countries, ratings, years, or plot details. +- If the sources do not support the question, say the catalogue does not contain enough information. +- Only cite show_ids that you actually used in the answer. +- Return valid JSON only. Do not wrap it in markdown. +- For recommendation questions, suggest up to 4 relevant titles when the sources support them. +""".strip() + +FEW_SHOT_EXAMPLES = """ +Example 1 + +User question: +Suggest an Indian comedy movie. + +Catalogue sources: +Source 1 +show_id: 111 +title: Example Comedy +Title: Example Comedy +Type: Movie +Countries: India +Genres: Comedies +Description: A light comedy about family confusion. + +Source 2 +show_id: 222 +title: Example Thriller +Title: Example Thriller +Type: Movie +Countries: India +Genres: Thrillers +Description: A dark crime story. + +Correct JSON response: +{ + "answer": "A good match is Example Comedy. It is an Indian movie listed under Comedies, so it fits your request.", + "used_show_ids": ["111"] +} + +Example 2 + +User question: +Which catalogue titles are about space travel? + +Catalogue sources: +Source 1 +show_id: 333 +title: Ocean Story +Title: Ocean Story +Type: Movie +Countries: United States +Genres: Documentaries +Description: A documentary about marine life. + +Correct JSON response: +{ + "answer": "The provided catalogue sources do not contain enough information to recommend a title about space travel.", + "used_show_ids": [] +} +""".strip() + + +@dataclass +class LLMAnswer: + """Parsed LLM answer with the source ids it used.""" + + answer: str + used_show_ids: list[str] + + +def build_context(titles: list[RetrievedTitle]) -> str: + """Format retrieved titles as grounded context for the LLM.""" + blocks = [] + + for index, title in enumerate(titles, start=1): + blocks.append( + f"Source {index}\n" + f"show_id: {title.show_id}\n" + f"title: {title.title}\n" + f"{title.text}" + ) + + return "\n\n---\n\n".join(blocks) + + +def build_prompt(question: str, titles: list[RetrievedTitle]) -> str: + """Build a guarded few-shot prompt for grounded catalogue Q&A.""" + context = build_context(titles) + + return f""" +{FEW_SHOT_EXAMPLES} + +Now answer the actual user question. + +User question: +{question} + +Catalogue sources: +{context} + +Return JSON with exactly these keys: +- "answer": string +- "used_show_ids": list of strings + +Important: +- The answer must be based only on the catalogue sources above. +- "used_show_ids" must include only show_ids that directly support the answer. +- If none of the sources support the question, use an empty list for "used_show_ids". +""".strip() + + +def parse_llm_json(content: str) -> LLMAnswer: + """Parse the JSON-only response returned by the LLM.""" + try: + payload = json.loads(content) + except json.JSONDecodeError as error: + raise RuntimeError("Groq returned invalid JSON.") from error + + answer = payload.get("answer") + used_show_ids = payload.get("used_show_ids") + + if not isinstance(answer, str) or not answer.strip(): + raise RuntimeError("Groq response did not contain a valid answer.") + + if not isinstance(used_show_ids, list): + raise RuntimeError("Groq response did not contain used_show_ids.") + + clean_show_ids = [str(show_id) for show_id in used_show_ids] + + return LLMAnswer(answer=answer.strip(), used_show_ids=clean_show_ids) + + +def answer_with_groq(question: str, titles: list[RetrievedTitle]) -> LLMAnswer: + """Send retrieved titles to Groq and return a parsed grounded answer.""" + load_dotenv() + + api_key = os.environ.get("GROQ_API_KEY") + if not api_key: + raise RuntimeError("GROQ_API_KEY is not set. Add it to your .env file.") + + from groq import Groq + + client = Groq(api_key=api_key) + + response = client.chat.completions.create( + model=GROQ_MODEL, + messages=[ + {"role": "system", "content": SYSTEM_MESSAGE}, + {"role": "user", "content": build_prompt(question, titles)}, + ], + temperature=0.1, + max_completion_tokens=500, + ) + + content = response.choices[0].message.content + + if content is None: + raise RuntimeError("Groq returned an empty response.") + + return parse_llm_json(content.strip()) diff --git a/rag/query_intelligence.py b/rag/query_intelligence.py new file mode 100644 index 0000000..a112b5c --- /dev/null +++ b/rag/query_intelligence.py @@ -0,0 +1,435 @@ +"""LLM query intelligence layer that creates direct Chroma search plans.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from typing import Any + +from dotenv import load_dotenv + +from rag.config import GROQ_MODEL + + +SYSTEM_MESSAGE = """ +You are an AI retrieval planner for a Netflix catalogue RAG system. + +Your job is NOT to answer the user. +Your job is to convert the user's natural-language question into a ChromaDB-searchable JSON plan. + +Return valid JSON only. +Do not use markdown. +Do not explain. +Do not include comments. + +The JSON must have exactly these top-level keys: +{ + "semantic_query": string, + "where": object or null +} + +The "semantic_query" will be embedded for vector search. +The "where" value will be passed directly to ChromaDB collection.query(where=...). + +Allowed Chroma metadata fields: +- countries +- genres +- type +- rating +- release_year + +Field rules: +1. countries + - Array field. + - Must use: {"countries": {"$contains": ""}} + +2. genres + - Array field. + - Must use: {"genres": {"$contains": ""}} + +3. type + - String field. + - Allowed values only: "Movie", "TV Show" + - Must use: {"type": "Movie"} or {"type": "TV Show"} + +4. rating + - String field. + - Must use: {"rating": ""} + +5. release_year + - Integer field. + - Allowed operators only: "$eq", "$gte", "$lte", "$gt", "$lt" + - Example: {"release_year": {"$gte": 2018}} + +Logical rules: +- Use "$and" only when all conditions must match. +- Use "$or" only when the user clearly asks for alternatives. +- Each item inside "$and" or "$or" must be a complete filter object. +- If there are multiple constraints, use explicit "$and". + +Strictly forbidden in "where": +- "country" +- "genre" +- "title" +- "description" +- "cast" +- "director" +- "date_added" +- any field not listed in the allowed metadata fields + +Use these exact field names: +- "countries", not "country" +- "genres", not "genre" + +Semantic query rewriting rules: +- The semantic_query should NOT simply copy the user query. +- The semantic_query should be a richer meaning-based search phrase for vector search. +- Remove strict filters from semantic_query when they are already represented in where. +- Do not include country, year, rating, or type words in semantic_query if they are already in where. +- Expand vague user intent into useful catalogue/search language. +- Keep semantic_query focused on story, mood, theme, tone, genre feel, or recommendation intent. + +Examples: +- "light heart movie" → "lighthearted feel-good entertaining comedy drama" +- "something emotional" → "emotional heartfelt family relationship drama" +- "dark crime show" → "dark crime investigation thriller" +- "romantic feel good movie" → "romantic feel-good heartwarming story" +- "scary movie" → "horror suspense scary thriller" +- "motivational movie" → "inspiring motivational uplifting story" +- "family friendly movie" → "family friendly wholesome lighthearted story" + +Domain understanding: +- "Bollywood", "desi", "Indian" means countries contains "India". +- "American", "US", "USA" means countries contains "United States". +- "Korean", "K-drama", "Kdrama" means countries contains "South Korea". +- "British", "UK" means countries contains "United Kingdom". +- "movie", "movies", "film", "films" means type "Movie". +- "series", "show", "shows", "TV series" means type "TV Show". +- "funny", "comedy", "comedies" means genre "Comedies" for movies. +- "comedy shows" or "TV comedy" means genre "TV Comedies". +- "drama movies" means genre "Dramas". +- "drama shows" means genre "TV Dramas". +- "romantic movies" means genre "Romantic Movies". +- "romantic shows" means genre "Romantic TV Shows". +- "scary" or "horror" means genre "Horror Movies". +- "thriller" means genre "Thrillers". +- "documentary" means genre "Documentaries". +- "docuseries" means genre "Docuseries". + +Year handling: +- "after 2018" means {"release_year": {"$gt": 2018}} +- "from 2018" means {"release_year": {"$gte": 2018}} +- "since 2018" means {"release_year": {"$gte": 2018}} +- "before 2020" means {"release_year": {"$lte": 2020}} +- "between 2017 and 2020" means {"release_year": {"$gte": 2017, "$lte": 2020}} +- "in 2019" means {"release_year": {"$eq": 2019}} + +Before returning JSON, silently self-check: +1. Is the output valid JSON? +2. Are the only top-level keys "semantic_query" and "where"? +3. Is "semantic_query" a non-empty string? +4. Is "where" either null or a valid Chroma filter? +5. Does "where" use only allowed fields? +6. Does "where" use "countries" and "genres" with "$contains"? +7. Did you avoid unsupported fields like "country", "genre", "title", "cast", or "director"? +8. If multiple filters are needed, did you wrap them inside "$and"? +9. Did you remove country, type, and year words from semantic_query if they are already represented in where? +10. Did you enrich vague mood words into better semantic search language? + +Examples: + +User: +Suggest Bollywood comedy movies after 2018 + +Correct JSON: +{ + "semantic_query": "funny lighthearted entertaining comedy movies", + "where": { + "$and": [ + {"countries": {"$contains": "India"}}, + {"genres": {"$contains": "Comedies"}}, + {"type": "Movie"}, + {"release_year": {"$gt": 2018}} + ] + } +} + +User: +Suggest Korean romantic shows + +Correct JSON: +{ + "semantic_query": "romantic heartfelt relationship drama shows", + "where": { + "$and": [ + {"countries": {"$contains": "South Korea"}}, + {"genres": {"$contains": "Romantic TV Shows"}}, + {"type": "TV Show"} + ] + } +} + +User: +Find American documentaries before 2020 + +Correct JSON: +{ + "semantic_query": "informative real-life documentary stories", + "where": { + "$and": [ + {"countries": {"$contains": "United States"}}, + {"genres": {"$contains": "Documentaries"}}, + {"release_year": {"$lte": 2020}} + ] + } +} + +User: +Suggest Indian or Korean dramas + +Correct JSON: +{ + "semantic_query": "emotional dramatic character-driven stories", + "where": { + "$and": [ + { + "$or": [ + {"countries": {"$contains": "India"}}, + {"countries": {"$contains": "South Korea"}} + ] + }, + {"genres": {"$contains": "Dramas"}} + ] + } +} + +User: +I want light heart movie of India after 2018 + +Correct JSON: +{ + "semantic_query": "lighthearted feel-good entertaining comedy drama", + "where": { + "$and": [ + {"countries": {"$contains": "India"}}, + {"type": "Movie"}, + {"release_year": {"$gt": 2018}} + ] + } +} + +User: +Suggest emotional family stories + +Correct JSON: +{ + "semantic_query": "emotional heartfelt family relationship stories", + "where": null +} + +User: +Suggest dark crime investigation shows + +Correct JSON: +{ + "semantic_query": "dark crime investigation suspense thriller series", + "where": { + "$and": [ + {"type": "TV Show"} + ] + } +} + +User: +Suggest family friendly movies + +Correct JSON: +{ + "semantic_query": "family friendly wholesome lighthearted feel-good story", + "where": { + "$and": [ + {"type": "Movie"} + ] + } +} +""".strip() + + +@dataclass +class SearchPlan: + """Direct Chroma search plan generated by the LLM.""" + + semantic_query: str + where: dict[str, Any] | None = None + + +ALLOWED_FIELDS = {"countries", "genres", "type", "rating", "release_year"} +ALLOWED_LOGICAL_OPERATORS = {"$and", "$or"} +ALLOWED_YEAR_OPERATORS = {"$eq", "$gte", "$lte", "$gt", "$lt"} + + +def _fallback_search_plan(question: str) -> SearchPlan: + """Return semantic-only search plan if LLM planning fails.""" + return SearchPlan(semantic_query=question, where=None) + + +def _strip_markdown_fences(content: str) -> str: + """Remove markdown JSON fences if the LLM accidentally returns them.""" + content = content.strip() + + if content.startswith("```json"): + content = content.removeprefix("```json").strip() + + if content.startswith("```"): + content = content.removeprefix("```").strip() + + if content.endswith("```"): + content = content.removesuffix("```").strip() + + return content + + +def _load_json_object(content: str) -> dict[str, Any]: + """Load a JSON object from LLM output.""" + content = _strip_markdown_fences(content) + + try: + payload = json.loads(content) + except json.JSONDecodeError as error: + raise RuntimeError("Query planner returned invalid JSON.") from error + + if not isinstance(payload, dict): + raise RuntimeError("Query planner JSON must be an object.") + + return payload + + +def _clean_semantic_query(value: Any, fallback_query: str) -> str: + """Return a valid semantic query.""" + if isinstance(value, str) and value.strip(): + return value.strip() + + return fallback_query + + +def _is_safe_where_filter(value: Any) -> bool: + """Return True only if the LLM produced a safe Chroma where filter.""" + if value is None: + return True + + if not isinstance(value, dict): + return False + + if not value: + return False + + # Chroma filters should have one root key: + # either one logical operator like "$and", + # or one metadata field like "type". + if len(value) != 1: + return False + + key, inner_value = next(iter(value.items())) + + if key in ALLOWED_LOGICAL_OPERATORS: + if not isinstance(inner_value, list) or not inner_value: + return False + + return all(_is_safe_where_filter(item) for item in inner_value) + + if key not in ALLOWED_FIELDS: + return False + + if key in {"countries", "genres"}: + if not isinstance(inner_value, dict): + return False + + if set(inner_value.keys()) != {"$contains"}: + return False + + contains_value = inner_value["$contains"] + + if not isinstance(contains_value, str): + return False + + return bool(contains_value.strip()) + + if key == "type": + return inner_value in {"Movie", "TV Show"} + + if key == "rating": + return isinstance(inner_value, str) and bool(inner_value.strip()) + + if key == "release_year": + if isinstance(inner_value, bool): + return False + + if isinstance(inner_value, int): + return True + + if not isinstance(inner_value, dict) or not inner_value: + return False + + for operator, year in inner_value.items(): + if operator not in ALLOWED_YEAR_OPERATORS: + return False + + if isinstance(year, bool) or not isinstance(year, int): + return False + + return True + + return False + + +def parse_search_plan(content: str, fallback_query: str) -> SearchPlan: + """Parse LLM output into a safe SearchPlan.""" + payload = _load_json_object(content) + + semantic_query = _clean_semantic_query( + payload.get("semantic_query"), + fallback_query=fallback_query, + ) + + where = payload.get("where") + + if not _is_safe_where_filter(where): + where = None + + return SearchPlan(semantic_query=semantic_query, where=where) + + +def create_search_plan(question: str) -> SearchPlan: + """Ask the LLM to directly generate a Chroma-searchable plan.""" + load_dotenv() + + api_key = os.environ.get("GROQ_API_KEY") + + if not api_key: + return _fallback_search_plan(question) + + from groq import Groq + + client = Groq(api_key=api_key) + + try: + response = client.chat.completions.create( + model=GROQ_MODEL, + messages=[ + {"role": "system", "content": SYSTEM_MESSAGE}, + {"role": "user", "content": question}, + ], + temperature=0, + max_completion_tokens=700, + ) + + content = response.choices[0].message.content + + if content is None: + return _fallback_search_plan(question) + + return parse_search_plan(content.strip(), fallback_query=question) + + except Exception: + return _fallback_search_plan(question) \ No newline at end of file diff --git a/rag/retriever.py b/rag/retriever.py new file mode 100644 index 0000000..e88792b --- /dev/null +++ b/rag/retriever.py @@ -0,0 +1,115 @@ +"""Retrieve relevant Netflix titles from the Chroma vector store.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from rag.config import CHROMA_COLLECTION_NAME, CHROMA_PATH, EMBEDDING_MODEL_NAME, RAG_TOP_K +from rag.embeddings import get_hf_client, embed_one_text + + +@dataclass +class RetrievedTitle: + """A title retrieved from the vector store.""" + + show_id: str + title: str + text: str + metadata: dict[str, Any] + distance: float + + +class TitleRetriever: + """Chroma-backed retriever for Netflix title documents.""" + + def __init__(self) -> None: + """Load Chroma and the HF client only when retrieval is needed.""" + if not CHROMA_PATH.exists(): + raise FileNotFoundError( + "Chroma index not found. Run `python -m rag.index` first." + ) + + import chromadb + + self.hf_client = get_hf_client() + self.client = chromadb.PersistentClient(path=str(CHROMA_PATH)) + self.collection = self.client.get_collection(name=CHROMA_COLLECTION_NAME) + + def _validate_query_embedding_dimension(self, query_embedding: list[float]) -> None: + """Ensure query embedding matches the Chroma index metadata.""" + metadata = self.collection.metadata or {} + + indexed_model = metadata.get("embedding_model") + indexed_dimension = metadata.get("embedding_dimension") + + if indexed_model and indexed_model != EMBEDDING_MODEL_NAME: + raise RuntimeError( + "Embedding model mismatch. " + f"Chroma index was built with '{indexed_model}', " + f"but current config uses '{EMBEDDING_MODEL_NAME}'. " + "Rebuild the index with `python -m rag.index`." + ) + + if indexed_dimension is None: + return + + try: + expected_dimension = int(indexed_dimension) + except (TypeError, ValueError) as error: + raise RuntimeError( + "Invalid embedding_dimension stored in Chroma metadata. " + "Rebuild the index with `python -m rag.index`." + ) from error + + actual_dimension = len(query_embedding) + + if actual_dimension != expected_dimension: + raise RuntimeError( + "Query embedding dimension does not match Chroma index. " + f"Expected {expected_dimension}, got {actual_dimension}. " + "Rebuild the index with `python -m rag.index`." + ) + + def retrieve( + self, + question: str, + top_k: int = RAG_TOP_K, + where_filter: dict[str, Any] | None = None, + ) -> list[RetrievedTitle]: + """Return the most relevant catalogue titles for a question.""" + query_embedding = embed_one_text(self.hf_client, question) + self._validate_query_embedding_dimension(query_embedding) + + query_args: dict[str, Any] = { + "query_embeddings": [query_embedding], + "n_results": top_k, + "include": ["documents", "metadatas", "distances"], + } + + if where_filter: + query_args["where"] = where_filter + + result = self.collection.query(**query_args) + + ids = result["ids"][0] + documents = result["documents"][0] + metadatas = result["metadatas"][0] + distances = result["distances"][0] + + return [ + RetrievedTitle( + show_id=str(item_id), + title=str(metadata.get("title", "")), + text=str(document), + metadata=dict(metadata), + distance=float(distance), + ) + for item_id, document, metadata, distance in zip( + ids, + documents, + metadatas, + distances, + strict=True, + ) + ] \ No newline at end of file diff --git a/rag/service.py b/rag/service.py new file mode 100644 index 0000000..1731319 --- /dev/null +++ b/rag/service.py @@ -0,0 +1,75 @@ +"""RAG service that retrieves titles and asks Groq for an answer.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache + +from rag.config import RAG_TOP_K +from rag.llm import answer_with_groq +from rag.retriever import RetrievedTitle, TitleRetriever +from rag.query_intelligence import create_search_plan + + + +@dataclass +class AnswerSource: + """A source title used in a RAG answer.""" + + show_id: str + title: str + + +@dataclass +class AskResult: + """Final RAG answer plus source titles.""" + + answer: str + sources: list[AnswerSource] + + +@lru_cache(maxsize=1) +def get_retriever() -> TitleRetriever: + """Create the retriever once and reuse it across requests.""" + return TitleRetriever() + + +def select_used_sources( + retrieved_titles: list[RetrievedTitle], + used_show_ids: list[str], +) -> list[AnswerSource]: + """Return only retrieved sources that the LLM said it used.""" + title_by_id = {title.show_id: title for title in retrieved_titles} + sources = [] + + for show_id in used_show_ids: + title = title_by_id.get(show_id) + + if title is not None: + sources.append(AnswerSource(show_id=title.show_id, title=title.title)) + + return sources + + +def ask_question(question: str, top_k: int = RAG_TOP_K) -> AskResult: + """Plan retrieval, retrieve relevant titles, ask Groq, and return answer.""" + retriever = get_retriever() + + search_plan = create_search_plan(question) + + retrieved_titles = retriever.retrieve( + search_plan.semantic_query, + top_k=top_k, + where_filter=search_plan.where, + ) + + if not retrieved_titles: + return AskResult( + answer="I could not find matching titles in the provided catalogue.", + sources=[], + ) + + llm_answer = answer_with_groq(question, retrieved_titles) + sources = select_used_sources(retrieved_titles, llm_answer.used_show_ids) + + return AskResult(answer=llm_answer.answer, sources=sources) diff --git a/requirements-demo.txt b/requirements-demo.txt new file mode 100644 index 0000000..d1b43aa --- /dev/null +++ b/requirements-demo.txt @@ -0,0 +1 @@ +streamlit==1.41.1 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7ad0251 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,13 @@ +fastapi==0.115.6 +uvicorn==0.34.0 +pytest==8.3.4 +httpx==0.28.1 +ruff==0.8.4 +python-dotenv==1.0.1 +numpy==2.2.1 +pandas==2.2.3 +ipykernel==6.29.5 +openpyxl==3.1.5 +chromadb==1.5.9 +huggingface_hub==0.32.2 +groq==1.2.0 \ No newline at end of file diff --git a/scripts/export_excel.py b/scripts/export_excel.py new file mode 100644 index 0000000..1780c14 --- /dev/null +++ b/scripts/export_excel.py @@ -0,0 +1,54 @@ +"""Export the cleaned SQLite catalogue to an Excel workbook for inspection.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +import pandas as pd + +DB_PATH = Path("data/netflix.db") +OUTPUT_PATH = Path("data/netflix_cleaned_preview.xlsx") + + +def read_table(connection: sqlite3.Connection, table_name: str) -> pd.DataFrame: + """Read a SQLite table into a dataframe.""" + return pd.read_sql_query(f"SELECT * FROM {table_name}", connection) + + +def export_excel() -> None: + """Write cleaned catalogue tables to a multi-sheet Excel workbook.""" + if not DB_PATH.exists(): + raise FileNotFoundError("Run `python ingest.py` before exporting Excel.") + + with sqlite3.connect(DB_PATH) as connection: + titles = read_table(connection, "titles") + countries = read_table(connection, "title_countries") + genres = read_table(connection, "title_genres") + cast = read_table(connection, "title_cast") + directors = read_table(connection, "title_directors") + + top_countries = pd.read_sql_query( + """ + SELECT country, COUNT(*) AS title_count + FROM title_countries + GROUP BY country + ORDER BY title_count DESC + LIMIT 20 + """, + connection, + ) + + with pd.ExcelWriter(OUTPUT_PATH, engine="openpyxl") as writer: + titles.to_excel(writer, sheet_name="titles", index=False) + countries.to_excel(writer, sheet_name="countries", index=False) + genres.to_excel(writer, sheet_name="genres", index=False) + cast.to_excel(writer, sheet_name="cast", index=False) + directors.to_excel(writer, sheet_name="directors", index=False) + top_countries.to_excel(writer, sheet_name="top_countries", index=False) + + print(f"Excel preview written: {OUTPUT_PATH}") + + +if __name__ == "__main__": + export_excel() diff --git a/scripts/inspect_db.py b/scripts/inspect_db.py new file mode 100644 index 0000000..bda8b32 --- /dev/null +++ b/scripts/inspect_db.py @@ -0,0 +1,291 @@ +"""Export SQLite database inspection output to an Excel workbook.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path +from typing import Any + +from openpyxl import Workbook +from openpyxl.styles import Font, PatternFill +from openpyxl.worksheet.worksheet import Worksheet + +DB_PATH = Path("data/netflix.db") +OUTPUT_PATH = Path("data/db_inspection.xlsx") + +TABLES = [ + "titles", + "title_countries", + "title_genres", + "title_cast", + "title_directors", +] + + +def fetch_rows(connection: sqlite3.Connection, query: str) -> list[tuple[Any, ...]]: + """Run a query and return all rows.""" + return connection.execute(query).fetchall() + + +def write_table( + sheet: Worksheet, + headers: list[str], + rows: list[tuple[Any, ...]], +) -> None: + """Write headers and rows to an Excel sheet.""" + sheet.append(headers) + + for cell in sheet[1]: + cell.font = Font(bold=True) + cell.fill = PatternFill("solid", fgColor="D9EAF7") + + for row in rows: + sheet.append(row) + + for column_cells in sheet.columns: + max_length = max(len(str(cell.value or "")) for cell in column_cells) + column_letter = column_cells[0].column_letter + sheet.column_dimensions[column_letter].width = min(max_length + 2, 45) + + +def add_overview_sheet(workbook: Workbook, connection: sqlite3.Connection) -> None: + """Add a high-level database overview sheet.""" + sheet = workbook.active + sheet.title = "Overview" + + rows = [] + + for table in TABLES: + count = connection.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + rows.append((table, count)) + + write_table(sheet, ["table_name", "row_count"], rows) + + +def add_schema_sheet(workbook: Workbook, connection: sqlite3.Connection) -> None: + """Add schema details for each database table.""" + sheet = workbook.create_sheet("Schema") + rows = [] + + for table in TABLES: + columns = connection.execute(f"PRAGMA table_info({table})").fetchall() + + for column in columns: + column_id, name, data_type, not_null, default_value, primary_key = column + rows.append( + ( + table, + column_id, + name, + data_type, + bool(not_null), + default_value, + bool(primary_key), + ) + ) + + write_table( + sheet, + [ + "table_name", + "column_id", + "column_name", + "data_type", + "not_null", + "default_value", + "primary_key", + ], + rows, + ) + + +def add_sample_titles_sheet(workbook: Workbook, connection: sqlite3.Connection) -> None: + """Add a small sample of cleaned title records.""" + rows = fetch_rows( + connection, + """ + SELECT show_id, title, type, release_year, rating, duration_value, + duration_unit, date_added, description + FROM titles + ORDER BY show_id + LIMIT 50 + """, + ) + + sheet = workbook.create_sheet("Sample Titles") + write_table( + sheet, + [ + "show_id", + "title", + "type", + "release_year", + "rating", + "duration_value", + "duration_unit", + "date_added", + "description", + ], + rows, + ) + + +def add_top_countries_sheet(workbook: Workbook, connection: sqlite3.Connection) -> None: + """Add top countries by title count.""" + rows = fetch_rows( + connection, + """ + SELECT country, COUNT(*) AS title_count + FROM title_countries + GROUP BY country + ORDER BY title_count DESC + LIMIT 25 + """, + ) + + sheet = workbook.create_sheet("Top Countries") + write_table(sheet, ["country", "title_count"], rows) + + +def add_sarkar_example_sheet( + workbook: Workbook, connection: sqlite3.Connection +) -> None: + """Add one joined example showing how normalized tables connect.""" + sheet = workbook.create_sheet("Joined Example") + + title_row = connection.execute( + """ + SELECT show_id, title, type, release_year, rating, duration_value, + duration_unit, date_added, description + FROM titles + WHERE title = 'Sarkar' + LIMIT 1 + """ + ).fetchone() + + if title_row is None: + write_table(sheet, ["message"], [("No Sarkar row found.",)]) + return + + show_id = title_row[0] + + sections = [ + ( + "title", + [ + "show_id", + "title", + "type", + "release_year", + "rating", + "duration_value", + "duration_unit", + "date_added", + "description", + ], + [title_row], + ), + ( + "countries", + ["country"], + fetch_rows( + connection, + f""" + SELECT country + FROM title_countries + WHERE show_id = '{show_id}' + ORDER BY position + """, + ), + ), + ( + "genres", + ["genre"], + fetch_rows( + connection, + f""" + SELECT genre + FROM title_genres + WHERE show_id = '{show_id}' + ORDER BY position + """, + ), + ), + ( + "directors", + ["director_name"], + fetch_rows( + connection, + f""" + SELECT director_name + FROM title_directors + WHERE show_id = '{show_id}' + ORDER BY position + """, + ), + ), + ( + "cast", + ["actor_name"], + fetch_rows( + connection, + f""" + SELECT actor_name + FROM title_cast + WHERE show_id = '{show_id}' + ORDER BY position + LIMIT 20 + """, + ), + ), + ] + + current_row = 1 + + for section_name, headers, rows in sections: + sheet.cell(row=current_row, column=1, value=section_name) + sheet.cell(row=current_row, column=1).font = Font(bold=True, size=13) + current_row += 1 + + for column_number, header in enumerate(headers, start=1): + cell = sheet.cell(row=current_row, column=column_number, value=header) + cell.font = Font(bold=True) + cell.fill = PatternFill("solid", fgColor="D9EAF7") + + current_row += 1 + + for row in rows: + for column_number, value in enumerate(row, start=1): + sheet.cell(row=current_row, column=column_number, value=value) + current_row += 1 + + current_row += 2 + + for column_cells in sheet.columns: + max_length = max(len(str(cell.value or "")) for cell in column_cells) + column_letter = column_cells[0].column_letter + sheet.column_dimensions[column_letter].width = min(max_length + 2, 45) + + +def export_excel() -> None: + """Create an Excel workbook that shows the cleaned database structure.""" + if not DB_PATH.exists(): + raise FileNotFoundError("Database not found. Run `python ingest.py` first.") + + workbook = Workbook() + + with sqlite3.connect(DB_PATH) as connection: + add_overview_sheet(workbook, connection) + add_schema_sheet(workbook, connection) + add_sample_titles_sheet(workbook, connection) + add_top_countries_sheet(workbook, connection) + add_sarkar_example_sheet(workbook, connection) + + OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True) + workbook.save(OUTPUT_PATH) + + print(f"Excel inspection file written: {OUTPUT_PATH}") + + +if __name__ == "__main__": + export_excel() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..cabb11d --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fastapi.testclient import TestClient + +from api.main import app +from ingestion.pipeline import run_ingestion + + +@pytest.fixture(scope="module", autouse=True) +def build_test_database() -> None: + """Create the SQLite database used by API tests.""" + run_ingestion() + + +client = TestClient(app) + + +def test_titles_filter_by_country_and_type() -> None: + """Filtering by country and type should return matching titles only.""" + response = client.get("/titles?country=India&type=Movie&page=1&page_size=5") + + assert response.status_code == 200 + + body = response.json() + assert body["total"] > 0 + assert len(body["items"]) > 0 + + for item in body["items"]: + assert item["type"] == "Movie" + assert "India" in item["countries"] + + +def test_get_title_returns_404_for_missing_show_id() -> None: + """Unknown show_id should return a clean 404 response.""" + response = client.get("/titles/DOES_NOT_EXIST") + + assert response.status_code == 404 + assert response.json()["detail"] == "Title not found" + + +def test_ask_returns_answer_and_sources(monkeypatch: pytest.MonkeyPatch) -> None: + """Ask endpoint should return an answer and at least one source.""" + fake_result = SimpleNamespace( + answer="Try Sarkar from the catalogue.", + sources=[ + SimpleNamespace(show_id="81075235", title="Sarkar"), + ], + ) + + def fake_ask_question(question: str) -> SimpleNamespace: + assert question == "Suggest an Indian political movie" + return fake_result + + monkeypatch.setattr("api.ask.ask_question", fake_ask_question) + + response = client.post( + "/ask", + json={"question": "Suggest an Indian political movie"}, + ) + + assert response.status_code == 200 + + body = response.json() + assert body["answer"] + assert len(body["sources"]) >= 1 + assert body["sources"][0]["show_id"] == "81075235" + assert body["sources"][0]["title"] == "Sarkar" diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..b90e33c --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,74 @@ +"""Unit tests for the embedding utility functions.""" + +from __future__ import annotations + +import math + +import pytest + +from rag.embeddings import ( + _extract_single_embedding, + _validate_embedding, + normalize_embedding, +) + + +def test_normalize_embedding_produces_unit_length() -> None: + """A normalized vector should have magnitude 1.0.""" + vector = [3.0, 4.0] + result = normalize_embedding(vector) + + magnitude = math.sqrt(sum(x * x for x in result)) + assert abs(magnitude - 1.0) < 1e-6 + + +def test_normalize_embedding_zero_vector_unchanged() -> None: + """A zero vector should be returned unchanged.""" + vector = [0.0, 0.0, 0.0] + result = normalize_embedding(vector) + + assert result == [0.0, 0.0, 0.0] + + +def test_extract_flat_vector() -> None: + """A flat list like [0.1, 0.2, 0.3] should be returned directly.""" + raw = [0.1, 0.2, 0.3] + result = _extract_single_embedding(raw) + + assert result == [0.1, 0.2, 0.3] + + +def test_extract_wrapped_vector() -> None: + """A wrapped list like [[0.1, 0.2, 0.3]] should be unwrapped.""" + raw = [[0.1, 0.2, 0.3]] + result = _extract_single_embedding(raw) + + assert result == [0.1, 0.2, 0.3] + + +def test_extract_token_embeddings_mean_pooled() -> None: + """Per-token embeddings should be mean-pooled into one vector.""" + raw = [ + [2.0, 4.0], + [4.0, 6.0], + ] + result = _extract_single_embedding(raw) + + assert result == [3.0, 5.0] + + +def test_extract_empty_raises_error() -> None: + """An empty embedding should raise a clear error.""" + with pytest.raises(RuntimeError, match="empty embedding"): + _extract_single_embedding([]) + + +def test_validate_embedding_rejects_empty() -> None: + """An empty vector should be rejected.""" + with pytest.raises(RuntimeError, match="Empty embedding"): + _validate_embedding([], context="test") + + +def test_validate_embedding_accepts_valid() -> None: + """A non-empty vector should pass validation.""" + _validate_embedding([0.1, 0.2], context="test") diff --git a/tests/test_ingestion.py b/tests/test_ingestion.py new file mode 100644 index 0000000..2ad6a66 --- /dev/null +++ b/tests/test_ingestion.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import csv +from pathlib import Path + +from ingestion.cleaning import clean_row +from ingestion.models import EXPECTED_COLUMNS, IngestStats +from ingestion.pipeline import load_clean_titles + + +def write_csv(path: Path, rows: list[dict[str, str]]) -> None: + """Write a small CSV fixture for ingestion tests.""" + path.parent.mkdir(parents=True, exist_ok=True) + + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=EXPECTED_COLUMNS) + writer.writeheader() + writer.writerows(rows) + + +def minimal_row(**overrides: str) -> dict[str, str]: + """Return a valid minimal Netflix CSV row for tests.""" + row = { + "show_id": "s1", + "type": "Movie", + "title": "Example Title", + "director": "", + "cast": "", + "country": "", + "date_added": "September 9, 2019", + "release_year": "2019", + "rating": "NR", + "duration": "90 min", + "listed_in": "Comedies", + "description": "Some description", + } + row.update(overrides) + return row + + +def test_ingestion_drops_missing_title(tmp_path: Path) -> None: + """Rows without a title should be dropped during ingestion.""" + csv_path = tmp_path / "missing_title.csv" + + write_csv(csv_path, [minimal_row(title="")]) + + titles, stats = load_clean_titles(csv_path) + + assert titles == [] + assert stats.rows_read == 1 + assert stats.rows_dropped == 1 + assert stats.dropped_reasons["missing_or_invalid_title"] == 1 + + +def test_clean_row_parses_metadata_fields() -> None: + """A valid row should parse dates, duration, ratings, and missing metadata.""" + stats = IngestStats() + + cleaned = clean_row(minimal_row(), stats) + + assert cleaned is not None + assert cleaned.date_added == "2019-09-09" + assert cleaned.duration_value == 90 + assert cleaned.duration_unit == "minutes" + assert cleaned.rating == "Not Rated" + assert cleaned.directors == ["Unknown"] + assert cleaned.cast_members == ["Unknown"] + assert cleaned.countries == ["Unknown"] + + +def test_ingestion_drops_duplicate_content_except_show_id(tmp_path: Path) -> None: + """Rows with identical cleaned content except show_id should be deduplicated.""" + csv_path = tmp_path / "duplicate_content.csv" + + write_csv( + csv_path, + [ + minimal_row(show_id="s1", title="Same Title"), + minimal_row(show_id="s2", title="Same Title"), + ], + ) + + titles, stats = load_clean_titles(csv_path) + + assert [title.show_id for title in titles] == ["s1"] + assert stats.rows_read == 2 + assert stats.rows_dropped == 1 + assert stats.dropped_reasons["duplicate_content_except_show_id"] == 1