Upload 9 files

Changelog.md

@@ -0,0 +1,2 @@
+- 2025-01-26: 
+- 2025-01-26: 

README.md

@@ -1,12 +1,96 @@
----
-title: Mongo Vector Search Util
-emoji: 💻
-colorFrom: pink
-colorTo: purple
-sdk: gradio
-sdk_version: 5.13.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Vector Search Demo App
+
+This is a Gradio web application that demonstrates vector search capabilities using MongoDB Atlas and OpenAI embeddings.
+
+## Prerequisites
+
+1. MongoDB Atlas account with vector search enabled
+2. OpenAI API key
+3. Python 3.8+
+4. Sample movie data loaded in MongoDB Atlas (sample_mflix database)
+
+## Setup
+
+1. Clone this repository
+
+2. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+3. Set up environment variables:
+```bash
+export OPENAI_API_KEY="your-openai-api-key"
+export ATLAS_URI="your-mongodb-atlas-connection-string"
+```
+
+4. Ensure your MongoDB Atlas setup:
+   - Database name: sample_mflix
+   - Collection: embedded_movies
+   - Vector search index: idx_plot_embedding
+   - Index configuration:
+   ```json
+   {
+     "fields": [
+       {
+         "type": "vector",
+         "path": "plot_embedding",
+         "numDimensions": 1536,
+         "similarity": "dotProduct"
+       }
+     ]
+   }
+   ```
+
+## Running the App
+
+Start the application:
+```bash
+python app.py
+```
+
+The app will be available at http://localhost:7860
+
+## Usage
+
+### Generating Embeddings
+1. Select your database and collection from the dropdowns
+2. Choose the field to generate embeddings for
+3. Specify the embedding field name (defaults to "embedding")
+4. Set a document limit (0 for all documents)
+5. Click "Generate Embeddings" to start processing
+
+The app uses memory-efficient cursor-based batch processing that can handle large collections:
+- Documents are processed in batches (default 20 documents per batch)
+- Memory usage is optimized through cursor-based iteration
+- Real-time progress tracking shows completed/total documents
+- Supports processing of large collections (100,000+ documents)
+- Automatically resumes from where it left off if embeddings already exist
+
+### Searching
+1. Enter a natural language query in the text box (e.g., "humans fighting aliens")
+2. Click "Submit" to search
+3. View the results showing matching documents with their similarity scores
+
+## Example Queries
+
+- "humans fighting aliens"
+- "relationship drama between two good friends"
+- "comedy about family vacation"
+- "detective solving mysterious murder"
+
+## Performance Notes
+
+The application is optimized for handling large datasets:
+- Uses cursor-based batch processing to avoid memory issues
+- Processes documents in configurable batch sizes (default: 20)
+- Implements parallel processing with ThreadPoolExecutor
+- Provides real-time progress tracking
+- Automatically handles memory cleanup during processing
+- Supports resuming interrupted operations
+
+## Notes
+
+- The search uses OpenAI's text-embedding-ada-002 model to create embeddings
+- Results are limited to top 5 matches
+- Similarity scores range from 0 to 1, with higher scores indicating better matches

app.py

@@ -0,0 +1,276 @@
+import os
+import gradio as gr
+from openai import OpenAI
+import json
+from dotenv import load_dotenv
+from db_utils import DatabaseUtils
+from embedding_utils import parallel_generate_embeddings, get_embedding
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Initialize OpenAI client
+openai_client = OpenAI()
+
+# Initialize database utils
+db_utils = DatabaseUtils()
+
+def get_field_names(db_name: str, collection_name: str) -> list[str]:
+    """Get list of fields in the collection"""
+    return db_utils.get_field_names(db_name, collection_name)
+
+def generate_embeddings_for_field(db_name: str, collection_name: str, field_name: str, embedding_field: str, limit: int = 10, progress=gr.Progress()) -> tuple[str, str]:
+    """Generate embeddings for documents in parallel with progress tracking"""
+    try:
+        db = db_utils.client[db_name]
+        collection = db[collection_name]
+        
+        # Count documents that need embeddings
+        total_docs = collection.count_documents({field_name: {"$exists": True}})
+        if total_docs == 0:
+            return f"No documents found with field '{field_name}'", ""
+            
+        # Get total count of documents that need processing
+        query = {
+            field_name: {"$exists": True},
+            embedding_field: {"$exists": False}  # Only get docs without embeddings
+        }
+        total_to_process = collection.count_documents(query)
+        if total_to_process == 0:
+            return "No documents found that need embeddings", ""
+            
+        # Apply limit if specified
+        if limit > 0:
+            total_to_process = min(total_to_process, limit)
+            
+        print(f"\nFound {total_to_process} documents that need embeddings...")
+        
+        # Progress tracking
+        progress_text = ""
+        def update_progress(prog: float, processed: int, total: int):
+            nonlocal progress_text
+            progress_text = f"Progress: {prog:.1f}% ({processed}/{total} documents)\n"
+            print(progress_text)  # Terminal logging
+            progress(prog/100, f"Processed {processed}/{total} documents")
+            
+        # Show initial progress
+        update_progress(0, 0, total_to_process)
+        
+        # Create cursor for batch processing
+        cursor = collection.find(query)
+        if limit > 0:
+            cursor = cursor.limit(limit)
+            
+        # Generate embeddings in parallel with cursor-based batching
+        processed = parallel_generate_embeddings(
+            collection=collection,
+            cursor=cursor,
+            field_name=field_name,
+            embedding_field=embedding_field,
+            openai_client=openai_client,
+            total_docs=total_to_process,
+            callback=update_progress
+        )
+                
+        # Return completion message and final progress
+        instructions = f"""
+Successfully generated embeddings for {processed} documents using parallel processing!
+
+To create the vector search index in MongoDB Atlas:
+1. Go to your Atlas cluster
+2. Click on 'Search' tab
+3. Create an index named 'vector_index' with this configuration:
+{{
+  "fields": [
+    {{
+      "type": "vector",
+      "path": "{embedding_field}",
+      "numDimensions": 1536,
+      "similarity": "dotProduct"
+    }}
+  ]
+}}
+
+You can now use the search tab with:
+- Field to search: {field_name}
+- Embedding field: {embedding_field}
+"""
+        return instructions, progress_text
+        
+    except Exception as e:
+        return f"Error: {str(e)}", ""
+
+def vector_search(query_text: str, db_name: str, collection_name: str, embedding_field: str, index_name: str) -> str:
+    """Perform vector search using embeddings"""
+    try:
+        print(f"\nProcessing query: {query_text}")
+        
+        db = db_utils.client[db_name]
+        collection = db[collection_name]
+        
+        # Get embeddings for query
+        embedding = get_embedding(query_text, openai_client)
+        print("Generated embeddings successfully")
+        
+        results = collection.aggregate([
+            {
+                '$vectorSearch': {
+                    "index": index_name,
+                    "path": embedding_field,
+                    "queryVector": embedding,
+                    "numCandidates": 50,
+                    "limit": 5
+                }
+            },
+            {
+                "$project": {
+                    "search_score": { "$meta": "vectorSearchScore" },
+                    "document": "$$ROOT"
+                }
+            }
+        ])
+        
+        # Format results
+        results_list = list(results)
+        formatted_results = []
+        
+        for idx, result in enumerate(results_list, 1):
+            doc = result['document']
+            formatted_result = f"{idx}. Score: {result['search_score']:.4f}\n"
+            # Add all fields except _id and embeddings
+            for key, value in doc.items():
+                if key not in ['_id', embedding_field]:
+                    formatted_result += f"{key}: {value}\n"
+            formatted_results.append(formatted_result)
+            
+        return "\n".join(formatted_results) if formatted_results else "No results found"
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create Gradio interface with tabs
+with gr.Blocks(title="MongoDB Vector Search Tool") as iface:
+    gr.Markdown("# MongoDB Vector Search Tool")
+    
+    # Get available databases
+    databases = db_utils.get_databases()
+    
+    with gr.Tab("Generate Embeddings"):
+        with gr.Row():
+            db_input = gr.Dropdown(
+                choices=databases,
+                label="Select Database",
+                info="Available databases in Atlas cluster"
+            )
+            collection_input = gr.Dropdown(
+                choices=[],
+                label="Select Collection",
+                info="Collections in selected database"
+            )
+        with gr.Row():
+            field_input = gr.Dropdown(
+                choices=[],
+                label="Select Field for Embeddings",
+                info="Fields available in collection"
+            )
+            embedding_field_input = gr.Textbox(
+                label="Embedding Field Name",
+                value="embedding",
+                info="Field name where embeddings will be stored"
+            )
+            limit_input = gr.Number(
+                label="Document Limit",
+                value=10,
+                minimum=0,
+                info="Number of documents to process (0 for all documents)"
+            )
+            
+        def update_collections(db_name):
+            collections = db_utils.get_collections(db_name)
+            # If there's only one collection, select it by default
+            value = collections[0] if len(collections) == 1 else None
+            return gr.Dropdown(choices=collections, value=value)
+            
+        def update_fields(db_name, collection_name):
+            if db_name and collection_name:
+                fields = get_field_names(db_name, collection_name)
+                return gr.Dropdown(choices=fields)
+            return gr.Dropdown(choices=[])
+        
+        # Update collections when database changes
+        db_input.change(
+            fn=update_collections,
+            inputs=[db_input],
+            outputs=[collection_input]
+        )
+        
+        # Update fields when collection changes
+        collection_input.change(
+            fn=update_fields,
+            inputs=[db_input, collection_input],
+            outputs=[field_input]
+        )
+        
+        generate_btn = gr.Button("Generate Embeddings")
+        generate_output = gr.Textbox(label="Results", lines=10)
+        progress_output = gr.Textbox(label="Progress", lines=3)
+        
+        generate_btn.click(
+            generate_embeddings_for_field,
+            inputs=[db_input, collection_input, field_input, embedding_field_input, limit_input],
+            outputs=[generate_output, progress_output]
+        )
+    
+    with gr.Tab("Search"):
+        with gr.Row():
+            search_db_input = gr.Dropdown(
+                choices=databases,
+                label="Select Database",
+                info="Database containing the vectors"
+            )
+            search_collection_input = gr.Dropdown(
+                choices=[],
+                label="Select Collection",
+                info="Collection containing the vectors"
+            )
+        with gr.Row():
+            search_embedding_field_input = gr.Textbox(
+                label="Embedding Field Name",
+                value="embedding",
+                info="Field containing the vectors"
+            )
+            search_index_input = gr.Textbox(
+                label="Vector Search Index Name",
+                value="vector_index",
+                info="Index created in Atlas UI"
+            )
+            
+        # Update collections when database changes
+        search_db_input.change(
+            fn=update_collections,
+            inputs=[search_db_input],
+            outputs=[search_collection_input]
+        )
+        
+        query_input = gr.Textbox(
+            label="Search Query",
+            lines=2,
+            placeholder="What would you like to search for?"
+        )
+        search_btn = gr.Button("Search")
+        search_output = gr.Textbox(label="Results", lines=10)
+        
+        search_btn.click(
+            vector_search,
+            inputs=[
+                query_input,
+                search_db_input,
+                search_collection_input,
+                search_embedding_field_input,
+                search_index_input
+            ],
+            outputs=search_output
+        )
+
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)

db_utils.py

@@ -0,0 +1,159 @@
+import os
+from typing import List, Dict, Any, Optional
+from pymongo import MongoClient
+from pymongo.errors import (
+    ConnectionFailure, 
+    OperationFailure, 
+    ServerSelectionTimeoutError,
+    InvalidName
+)
+from dotenv import load_dotenv
+
+class DatabaseError(Exception):
+    """Base class for database operation errors"""
+    pass
+
+class ConnectionError(DatabaseError):
+    """Error when connecting to MongoDB Atlas"""
+    pass
+
+class OperationError(DatabaseError):
+    """Error during database operations"""
+    pass
+
+class DatabaseUtils:
+    """Utility class for MongoDB Atlas database operations
+    
+    This class provides methods to interact with MongoDB Atlas databases and collections,
+    including listing databases, collections, and retrieving collection information.
+    
+    Attributes:
+        atlas_uri (str): MongoDB Atlas connection string
+        client (MongoClient): MongoDB client instance
+    """
+    
+    def __init__(self):
+        """Initialize DatabaseUtils with MongoDB Atlas connection
+        
+        Raises:
+            ConnectionError: If unable to connect to MongoDB Atlas
+            ValueError: If ATLAS_URI environment variable is not set
+        """
+        # Load environment variables
+        load_dotenv()
+        
+        self.atlas_uri = os.getenv("ATLAS_URI")
+        if not self.atlas_uri:
+            raise ValueError("ATLAS_URI environment variable is not set")
+            
+        try:
+            self.client = MongoClient(self.atlas_uri)
+            # Test connection
+            self.client.admin.command('ping')
+        except (ConnectionFailure, ServerSelectionTimeoutError) as e:
+            raise ConnectionError(f"Failed to connect to MongoDB Atlas: {str(e)}")
+    
+    def get_databases(self) -> List[str]:
+        """Get list of all databases in Atlas cluster
+        
+        Returns:
+            List[str]: List of database names
+            
+        Raises:
+            OperationError: If unable to list databases
+        """
+        try:
+            return self.client.list_database_names()
+        except OperationFailure as e:
+            raise OperationError(f"Failed to list databases: {str(e)}")
+    
+    def get_collections(self, db_name: str) -> List[str]:
+        """Get list of collections in a database
+        
+        Args:
+            db_name (str): Name of the database
+            
+        Returns:
+            List[str]: List of collection names
+            
+        Raises:
+            OperationError: If unable to list collections
+            ValueError: If db_name is empty or invalid
+        """
+        if not db_name or not isinstance(db_name, str):
+            raise ValueError("Database name must be a non-empty string")
+            
+        try:
+            db = self.client[db_name]
+            return db.list_collection_names()
+        except (OperationFailure, InvalidName) as e:
+            raise OperationError(f"Failed to list collections for database '{db_name}': {str(e)}")
+    
+    def get_collection_info(self, db_name: str, collection_name: str) -> Dict[str, Any]:
+        """Get information about a collection including document count and sample document
+        
+        Args:
+            db_name (str): Name of the database
+            collection_name (str): Name of the collection
+            
+        Returns:
+            Dict[str, Any]: Dictionary containing collection information:
+                - count: Number of documents in collection
+                - sample: Sample document from collection (if exists)
+                
+        Raises:
+            OperationError: If unable to get collection information
+            ValueError: If db_name or collection_name is empty or invalid
+        """
+        if not db_name or not isinstance(db_name, str):
+            raise ValueError("Database name must be a non-empty string")
+        if not collection_name or not isinstance(collection_name, str):
+            raise ValueError("Collection name must be a non-empty string")
+            
+        try:
+            db = self.client[db_name]
+            collection = db[collection_name]
+            
+            return {
+                'count': collection.count_documents({}),
+                'sample': collection.find_one()
+            }
+        except (OperationFailure, InvalidName) as e:
+            raise OperationError(
+                f"Failed to get info for collection '{collection_name}' "
+                f"in database '{db_name}': {str(e)}"
+            )
+    
+    def get_field_names(self, db_name: str, collection_name: str) -> List[str]:
+        """Get list of fields in a collection based on sample document
+        
+        Args:
+            db_name (str): Name of the database
+            collection_name (str): Name of the collection
+            
+        Returns:
+            List[str]: List of field names (excluding _id and embedding fields)
+            
+        Raises:
+            OperationError: If unable to get field names
+            ValueError: If db_name or collection_name is empty or invalid
+        """
+        try:
+            info = self.get_collection_info(db_name, collection_name)
+            sample = info.get('sample', {})
+            
+            if sample:
+                # Get all field names except _id and any existing embedding fields
+                return [field for field in sample.keys() 
+                        if field != '_id' and not field.endswith('_embedding')]
+            return []
+        except DatabaseError as e:
+            raise OperationError(
+                f"Failed to get field names for collection '{collection_name}' "
+                f"in database '{db_name}': {str(e)}"
+            )
+    
+    def close(self):
+        """Close MongoDB connection safely"""
+        if hasattr(self, 'client'):
+            self.client.close()

embedding_utils.py

@@ -0,0 +1,122 @@
+from typing import List, Tuple
+from concurrent.futures import ThreadPoolExecutor
+from pymongo import UpdateOne
+from pymongo.collection import Collection
+import math
+
+def get_embedding(text: str, openai_client, model="text-embedding-ada-002") -> list[float]:
+    """Get embeddings for given text using OpenAI API"""
+    text = text.replace("\n", " ")
+    resp = openai_client.embeddings.create(
+        input=[text],
+        model=model
+    )
+    return resp.data[0].embedding
+
+def process_batch(docs: List[dict], field_name: str, embedding_field: str, openai_client) -> List[Tuple[str, list]]:
+    """Process a batch of documents to generate embeddings"""
+    results = []
+    for doc in docs:
+        # Skip if embedding already exists
+        if embedding_field in doc:
+            continue
+            
+        text = doc[field_name]
+        if isinstance(text, str):
+            embedding = get_embedding(text, openai_client)
+            results.append((doc["_id"], embedding))
+    return results
+
+def parallel_generate_embeddings(
+    collection: Collection,
+    cursor,
+    field_name: str,
+    embedding_field: str,
+    openai_client,
+    total_docs: int,
+    batch_size: int = 20,
+    callback=None
+) -> int:
+    """Generate embeddings in parallel using ThreadPoolExecutor with cursor-based batching
+    
+    Args:
+        collection: MongoDB collection
+        cursor: MongoDB cursor for document iteration
+        field_name: Field containing text to embed
+        embedding_field: Field to store embeddings
+        openai_client: OpenAI client instance
+        total_docs: Total number of documents to process
+        batch_size: Size of batches for parallel processing
+        callback: Optional callback function for progress updates
+        
+    Returns:
+        Number of documents processed
+    """
+    if total_docs == 0:
+        return 0
+        
+    processed = 0
+    
+    # Initial progress update
+    if callback:
+        callback(0, 0, total_docs)
+        
+    # Process documents in batches using cursor
+    with ThreadPoolExecutor(max_workers=20) as executor:
+        batch = []
+        futures = []
+        
+        # Iterate through cursor and build batches
+        for doc in cursor:
+            batch.append(doc)
+            
+            if len(batch) >= batch_size:
+                # Submit batch for processing
+                future = executor.submit(process_batch, batch.copy(), field_name, embedding_field, openai_client)
+                futures.append(future)
+                batch = []  # Clear batch for next round
+                
+                # Process completed futures to free up memory
+                completed_futures = [f for f in futures if f.done()]
+                for future in completed_futures:
+                    results = future.result()
+                    if results:
+                        # Batch update MongoDB
+                        bulk_ops = [
+                            UpdateOne({"_id": doc_id}, {"$set": {embedding_field: embedding}})
+                            for doc_id, embedding in results
+                        ]
+                        if bulk_ops:
+                            collection.bulk_write(bulk_ops)
+                            processed += len(bulk_ops)
+                            
+                        # Update progress
+                        if callback:
+                            progress = (processed / total_docs) * 100
+                            callback(progress, processed, total_docs)
+                            
+                futures = [f for f in futures if not f.done()]
+        
+        # Process any remaining documents in the last batch
+        if batch:
+            future = executor.submit(process_batch, batch, field_name, embedding_field, openai_client)
+            futures.append(future)
+        
+        # Wait for remaining futures to complete
+        for future in futures:
+            results = future.result()
+            if results:
+                bulk_ops = [
+                    UpdateOne({"_id": doc_id}, {"$set": {embedding_field: embedding}})
+                    for doc_id, embedding in results
+                ]
+                if bulk_ops:
+                    collection.bulk_write(bulk_ops)
+                    processed += len(bulk_ops)
+                    
+                    # Final progress update
+                    if callback:
+                        progress = (processed / total_docs) * 100
+                        callback(progress, processed, total_docs)
+                    
+    return processed

list_collections.py

@@ -0,0 +1,38 @@
+import os
+from pymongo import MongoClient
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Initialize MongoDB client
+atlas_uri = os.getenv("ATLAS_URI")
+client = MongoClient(atlas_uri)
+
+def list_all_collections():
+    """List all databases and their collections in the Atlas cluster"""
+    try:
+        # Get all database names
+        db_names = client.list_database_names()
+        
+        print("\nDatabases and Collections in your Atlas cluster:\n")
+        
+        # For each database, get and print collections
+        for db_name in db_names:
+            print(f"Database: {db_name}")
+            db = client[db_name]
+            collections = db.list_collection_names()
+            
+            for collection in collections:
+                # Get count of documents in collection
+                count = db[collection].count_documents({})
+                print(f"  └── Collection: {collection} ({count} documents)")
+            print()
+            
+    except Exception as e:
+        print(f"Error: {str(e)}")
+    finally:
+        client.close()
+
+if __name__ == "__main__":
+    list_all_collections()

list_db.py

@@ -0,0 +1,27 @@
+from db_utils import DatabaseUtils
+
+def main():
+    db_utils = DatabaseUtils()
+    try:
+        print("\nDatabases and Collections in your Atlas cluster:\n")
+        
+        # Get all databases
+        databases = db_utils.get_databases()
+        
+        # For each database, show collections and counts
+        for db_name in databases:
+            print(f"Database: {db_name}")
+            collections = db_utils.get_collections(db_name)
+            
+            for coll_name in collections:
+                info = db_utils.get_collection_info(db_name, coll_name)
+                print(f"  └── Collection: {coll_name} ({info['count']} documents)")
+            print()
+            
+    except Exception as e:
+        print(f"Error: {str(e)}")
+    finally:
+        db_utils.close()
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,57 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.8.0
+certifi==2024.12.14
+charset-normalizer==3.4.1
+click==8.1.8
+distro==1.9.0
+dnspython==2.7.0
+fastapi==0.115.7
+ffmpy==0.5.0
+filelock==3.17.0
+fsspec==2024.12.0
+gradio==5.13.1
+gradio_client==1.6.0
+h11==0.14.0
+httpcore==1.0.7
+httpx==0.28.1
+huggingface-hub==0.27.1
+idna==3.10
+Jinja2==3.1.5
+jiter==0.8.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+numpy==2.2.2
+openai==1.60.1
+orjson==3.10.15
+packaging==24.2
+pandas==2.2.3
+pillow==11.1.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydub==0.25.1
+Pygments==2.19.1
+pymongo==4.10.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.20
+pytz==2024.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.4
+ruff==0.9.3
+safehttpx==0.1.6
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.45.3
+tomlkit==0.13.2
+tqdm==4.67.1
+typer==0.15.1
+typing_extensions==4.12.2
+tzdata==2025.1
+urllib3==2.3.0
+uvicorn==0.34.0
+websockets==14.2

setup_data.py

@@ -0,0 +1,94 @@
+import os
+from pymongo import MongoClient
+from openai import OpenAI
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Initialize clients
+openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+atlas_uri = os.getenv("ATLAS_URI")
+client = MongoClient(atlas_uri)
+db = client['sample_mflix']
+collection = db['embedded_movies']
+
+# Sample movie data
+sample_movies = [
+    {
+        "title": "The Matrix",
+        "year": 1999,
+        "plot": "A computer programmer discovers that reality as he knows it is a simulation created by machines, and joins a rebellion to overthrow them."
+    },
+    {
+        "title": "Inception",
+        "year": 2010,
+        "plot": "A thief who enters the dreams of others to steal secrets from their subconscious is offered a chance to regain his old life in exchange for a task considered impossible: inception."
+    },
+    {
+        "title": "The Shawshank Redemption",
+        "year": 1994,
+        "plot": "Two imprisoned men bond over a number of years, finding solace and eventual redemption through acts of common decency."
+    },
+    {
+        "title": "Jurassic Park",
+        "year": 1993,
+        "plot": "A pragmatic paleontologist visiting an almost complete theme park is tasked with protecting a couple of kids after a power failure causes the park's cloned dinosaurs to run loose."
+    },
+    {
+        "title": "The Lord of the Rings: The Fellowship of the Ring",
+        "year": 2001,
+        "plot": "A young hobbit, Frodo, who has found the One Ring that belongs to the Dark Lord Sauron, begins his journey with eight companions to Mount Doom, the only place where it can be destroyed."
+    }
+]
+
+def get_embedding(text: str, model="text-embedding-ada-002") -> list[float]:
+    """Get embeddings for given text using OpenAI API"""
+    text = text.replace("\n", " ")
+    resp = openai_client.embeddings.create(
+        input=[text],
+        model=model
+    )
+    return resp.data[0].embedding
+
+def setup_data():
+    try:
+        # Drop existing collection if it exists
+        collection.drop()
+        print("Dropped existing collection")
+
+        # Add embeddings to movies and insert into collection
+        for movie in sample_movies:
+            # Generate embedding for plot
+            embedding = get_embedding(movie["plot"])
+            movie["plot_embedding"] = embedding
+            
+            # Insert movie with embedding
+            collection.insert_one(movie)
+            print(f"Inserted movie: {movie['title']}")
+
+        print("\nData setup completed successfully!")
+        print("\nIMPORTANT: You need to create the vector search index manually in the Atlas UI:")
+        print("1. Go to your Atlas cluster")
+        print("2. Click on 'Search' tab")
+        print("3. Create an index named 'idx_plot_embedding' with this definition:")
+        print("""
+{
+  "fields": [
+    {
+      "type": "vector",
+      "path": "plot_embedding",
+      "numDimensions": 1536,
+      "similarity": "dotProduct"
+    }
+  ]
+}
+        """)
+        
+    except Exception as e:
+        print(f"Error during setup: {str(e)}")
+    finally:
+        client.close()
+
+if __name__ == "__main__":
+    setup_data()