Advent of Haystack: Day 9

!pip install haystack-ai mongodb-atlas-haystack tiktoken datasets colorama -q --upgrade openai

import os
import getpass
import re

conn_str = getpass.getpass("Enter your MongoDB connection string:")
conn_str = (
    re.sub(r"appName=[^\s]*", "appName=devrel.ai.haystack_partner", conn_str)
    if "appName=" in conn_str
    else conn_str
    + ("&" if "?" in conn_str else "?")
    + "appName=devrel.ai.haystack_partner"
)
os.environ["MONGO_CONNECTION_STRING"] = conn_str
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter your OpenAI API Key:")

dataset = {
    "train": [
        {
            "title": "LEGO Star Wars Set",
            "price": "$49.99",
            "description": "Build your own galaxy with this exciting LEGO Star Wars set",
            "category": "Toys",
            "age_range": "7-12",
        },
        {
            "title": "Remote Control Car",
            "price": "$29.99",
            "description": "Fast and fun RC car with full directional control",
            "category": "Toys",
            "age_range": "6-10",
        },
        {
            "title": "Art Set",
            "price": "$24.99",
            "description": "Complete art set with paints, brushes, and canvas",
            "category": "Arts & Crafts",
            "age_range": "5-15",
        },
        {
            "title": "Science Kit",
            "price": "$34.99",
            "description": "Educational science experiments kit",
            "category": "Educational",
            "age_range": "8-14",
        },
        {
            "title": "Dollhouse",
            "price": "$89.99",
            "description": "Beautiful wooden dollhouse with furniture",
            "category": "Toys",
            "age_range": "4-10",
        },
        {
            "title": "Building Blocks Set",
            "price": "$39.99",
            "description": "Classic wooden building blocks in various shapes and colors",
            "category": "Educational",
            "age_range": "3-8",
        },
        {
            "title": "Board Game Collection",
            "price": "$44.99",
            "description": "Set of 5 classic family board games",
            "category": "Games",
            "age_range": "6-99",
        },
        {
            "title": "Puppet Theater",
            "price": "$59.99",
            "description": "Wooden puppet theater with 6 hand puppets",
            "category": "Creative Play",
            "age_range": "4-12",
        },
        {
            "title": "Robot Building Kit",
            "price": "$69.99",
            "description": "Build and program your own robot with this STEM kit",
            "category": "Educational",
            "age_range": "10-16",
        },
        {
            "title": "Play Kitchen",
            "price": "$79.99",
            "description": "Realistic play kitchen with sounds and accessories",
            "category": "Pretend Play",
            "age_range": "3-8",
        },
        {
            "title": "Musical Instrument Set",
            "price": "$45.99",
            "description": "Collection of kid-friendly musical instruments",
            "category": "Music",
            "age_range": "3-10",
        },
        {
            "title": "Sports Equipment Pack",
            "price": "$54.99",
            "description": "Complete set of kids' sports gear including ball, bat, and net",
            "category": "Sports",
            "age_range": "6-12",
        },
        {
            "title": "Magic Tricks Kit",
            "price": "$29.99",
            "description": "Professional magic set with instruction manual",
            "category": "Entertainment",
            "age_range": "8-15",
        },
        {
            "title": "Dinosaur Collection",
            "price": "$39.99",
            "description": "Set of 12 detailed dinosaur figures with fact cards",
            "category": "Educational",
            "age_range": "4-12",
        },
        {
            "title": "Craft Supply Bundle",
            "price": "$49.99",
            "description": "Comprehensive craft supplies including beads, yarn, and tools",
            "category": "Arts & Crafts",
            "age_range": "6-16",
        },
        {
            "title": "Coding for Kids Set",
            "price": "$64.99",
            "description": "Interactive coding kit with programmable robot and game cards",
            "category": "STEM",
            "age_range": "8-14",
        },
        {
            "title": "Dress Up Trunk",
            "price": "$49.99",
            "description": "Collection of costumes and accessories for imaginative play",
            "category": "Pretend Play",
            "age_range": "3-10",
        },
        {
            "title": "Microscope Kit",
            "price": "$59.99",
            "description": "Real working microscope with prepared slides and tools",
            "category": "Science",
            "age_range": "10-15",
        },
        {
            "title": "Outdoor Explorer Kit",
            "price": "$34.99",
            "description": "Nature exploration set with binoculars, compass, and field guide",
            "category": "Outdoor",
            "age_range": "7-12",
        },
        {
            "title": "Art Pottery Studio",
            "price": "$69.99",
            "description": "Complete pottery wheel set with clay and glazing materials",
            "category": "Arts & Crafts",
            "age_range": "8-16",
        },
    ]
}

# Create collection gifts and add the vector index


from pymongo import MongoClient
from bson import json_util
from pymongo.operations import SearchIndexModel
import time

client = MongoClient(os.environ["MONGO_CONNECTION_STRING"])
db = client["santa_workshop"]
collection = db["gifts"]

db.create_collection("gifts")


## create index
search_index_model = SearchIndexModel(
    definition={
        "fields": [
            {
                "type": "vector",
                "numDimensions": 1536,
                "path": "embedding",
                "similarity": "cosine",
            },
        ]
    },
    name="vector_index",
    type="vectorSearch",
)
result = collection.create_search_index(model=search_index_model)
print("New search index named " + result + " is building.")
# Wait for initial sync to complete
print("Polling to check if the index is ready. This may take up to a minute.")
predicate = None
if predicate is None:
    predicate = lambda index: index.get("queryable") is True
while True:
    indices = list(collection.list_search_indexes(result))
    if len(indices) and predicate(indices[0]):
        break
    time.sleep(5)
print(result + " is ready for querying.")
client.close()

from haystack import Pipeline, Document
from haystack.document_stores.types import DuplicatePolicy
from haystack.components.writers import DocumentWriter
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack_integrations.document_stores.mongodb_atlas import (
    MongoDBAtlasDocumentStore,
)

# Initialize document store
document_store = MongoDBAtlasDocumentStore(
    database_name="santa_workshop",
    collection_name="gifts",
    vector_search_index="vector_index",
)

# Convert dataset to Haystack documents
insert_data = [
    Document(content=gift["title"], meta=json_util.loads(json_util.dumps(gift)))
    for gift in dataset["train"]
]

# Create and run indexing pipeline
doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
doc_embedder = OpenAIDocumentEmbedder(
    model="text-embedding-3-small",
    meta_fields_to_embed=["description"],  # Include description in embeddings
)

indexing_pipe = Pipeline()
indexing_pipe.add_component(instance=doc_embedder, name="doc_embedder")
indexing_pipe.add_component(instance=doc_writer, name="doc_writer")
indexing_pipe.connect("doc_embedder.documents", "doc_writer.documents")

indexing_pipe.run({"doc_embedder": {"documents": insert_data}})

from haystack.components.generators import OpenAIGenerator
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack_integrations.components.retrievers.mongodb_atlas import (
    MongoDBAtlasEmbeddingRetriever,
)
from haystack.components.embedders import OpenAITextEmbedder
from typing import List
from haystack import component
from IPython.display import display, Markdown


# Create a GiftChecker component for self-reflection
@component
class GiftChecker:
    @component.output_types(gifts_to_check=str, gifts=str)
    def run(self, replies: List[str]):
        # Hint: Check if the selection is optimal ('DONE') or needs improvement
        latest_message = replies[0]
        print(f"{latest_message=}")
        if "DONE" in latest_message:
            display(Markdown("**OPTIMAL**"))
            return {"gifts": latest_message.replace("DONE", "")}
        else:
            display(Markdown("**NEEDS IMPROVEMENT**"))
            return {"gifts_to_check": latest_message}


# Create your prompt template
prompt_template = """
    You are Santa's gift selection assistant. Below you have a list of available gifts with their prices.
    Based on the child's wishlist and budget, suggest appropriate gifts that maximize joy while staying within budget.

    Available Gifts:
    {% for doc in documents %}
        Gift: {{ doc.content }}
        Price: {{ doc.meta['price']}}
        Age Range: {{ doc.meta['age_range']}}
    {% endfor %}

    Query: {{query}}
    {% if gifts_to_check %}
        Previous gift selection: {{gifts_to_check[0]}}
        Can we optimize this selection for better value within budget?
        If optimal, say 'DONE' and return the selection
        If not, suggest a better combination
    {% endif %}

    Gift Selection:
"""

# Create and connect your pipeline components
gift_pipeline = Pipeline(max_runs_per_component=5)

# Add your components here!
# Hint: You'll need:
# - OpenAITextEmbedder for query embedding
# - MongoDBAtlasEmbeddingRetriever for finding relevant gifts
# - PromptBuilder for creating the prompt
# - GiftChecker for self-reflection
# - OpenAIGenerator for generating responses

text_embedder = OpenAITextEmbedder(model="text-embedding-3-small")
retriever = MongoDBAtlasEmbeddingRetriever(document_store=document_store)
prompt_builder = PromptBuilder(template=prompt_template)
gift_checker = GiftChecker()
llm = OpenAIGenerator(model="gpt-4o")

# Connect your components here!
# Hint: Think about the flow of data between components

gift_pipeline.add_component(name="text_embedder", instance=text_embedder)
gift_pipeline.add_component(name="retriever", instance=retriever)
gift_pipeline.add_component(name="prompt_builder", instance=prompt_builder)
gift_pipeline.add_component(name="gift_checker", instance=gift_checker)
gift_pipeline.add_component(name="llm", instance=llm)

gift_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
gift_pipeline.connect("retriever", "prompt_builder")
gift_pipeline.connect("prompt_builder", "llm")
gift_pipeline.connect("llm", "gift_checker")
gift_pipeline.connect("gift_checker", "prompt_builder")

gift_pipeline.show()

# Test queries to try:
queries = [
    "Find gifts for a 9-year-old who loves science and building things. Budget: $100",
    "Need gifts for a creative 6-year-old interested in art. Budget: $50",
    "Looking for educational toys for a 12-year-old. Budget: $75",
]

for query in queries:
    # Run the pipeline and print the results
    output = gift_pipeline.run(
        {"text_embedder": {"text": query}, "prompt_builder": {"query": query}}
    )
    display(Markdown(output["gift_checker"]["gifts"]))
    print("-" * 100)