This repository shows a minimal vector search workflow for movie data using:
- OpenAI embeddings
- ChromaDB as the local vector store
- A CSV dataset of movies
The script reads movie records from movies_dataset.csv, generates embeddings for each movie description, stores them in a persistent Chroma collection, and runs a semantic search query against that collection.
vector_database_example.py- main example scriptmovies_dataset.csv- input dataset used to build the collection
- Loads the movie dataset with pandas
- Builds a searchable text field from title, type, description, and categories
- Counts tokens to estimate embedding cost
- Requests embeddings from the OpenAI API
- Stores documents, embeddings, and metadata in a local Chroma database
- Queries the collection using a natural-language search string
- Prints the top matching results
- Python 3.9+
- An OpenAI API key
Install dependencies:
pip install -r requirements.txtSet your API key:
export OPENAI_API_KEY="your_api_key_here"From the project root:
python vector_database_example.py "A mind-bending science fiction movie about space or dreams"The main settings are defined near the top of vector_database_example.py:
EMBEDDING_MODEL- embedding model used for vector creationCOST_PER_MILLION_TOKENS- estimated embedding cost referenceCHROMA_DB_PATH- local folder where Chroma stores dataCOLLECTION_NAME- collection name inside ChromaN_RESULTS- number of matches returned
When the script runs, it prints:
- dataset columns and shape
- total token count
- estimated embedding cost
- confirmation that movies were inserted into Chroma
- the top semantic search matches, including distance, metadata, and document text
- The Chroma database is persisted locally in
./chroma_movies_db. - Running the script repeatedly will reuse the same collection name and upsert the same IDs.
- The embedding cost printed by the script is only an estimate based on the configured rate.
- The script expects
movies_dataset.csvto contain fields such asid,title,type,description, andlisted_in.
Pass the query as a command-line argument:
python vector_database_example.py "A mind-bending science fiction movie about space or dreams"