Add guest stories and auto-build Chroma database

data/guest_stories/ada_lovelace.md

@@ -0,0 +1,7 @@
+# Ada Lovelace
+
+Ada Lovelace is considered the world's first computer programmer.
+
+She collaborated with Charles Babbage on the Analytical Engine.
+
+At the AI Gala she will speak about creativity, mathematics, and artificial intelligence.

data/guest_stories/alan_turning.md

@@ -0,0 +1,7 @@
+# Alan Turing
+
+Alan Turing is known as the father of theoretical computer science.
+
+He helped break the Enigma code during World War II.
+
+At the AI Gala he will discuss machine intelligence and the Turing Test.

data/guest_stories/grace_hopper.md

@@ -0,0 +1,7 @@
+# Grace Hopper
+
+Grace Hopper developed one of the first compilers.
+
+She popularized machine-independent programming languages.
+
+At the AI Gala she will speak about programming languages and software engineering.

requirements.txt

@@ -6,4 +6,8 @@ llama-index
 llama-index-vector-stores-chroma
 llama-index-embeddings-huggingface
 sentence-transformers
-openai
+huggingface-hub
+openai
+gradio
+pandas
+requests

tools/rag_tool.py

@@ -1,37 +1,79 @@
+import os
 import chromadb
 
 from smolagents import tool
 
-from llama_index.core import StorageContext, VectorStoreIndex
+from llama_index.core import (
+    StorageContext,
+    VectorStoreIndex,
+    SimpleDirectoryReader,
+)
+from llama_index.core.node_parser import SentenceSplitter
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.vector_stores.chroma import ChromaVectorStore
 
 
-# Load the embedding model once
+# -----------------------------
+# Embedding model
+# -----------------------------
 embed_model = HuggingFaceEmbedding(
     model_name="BAAI/bge-small-en-v1.5"
 )
 
-# Connect to the existing ChromaDB
+# -----------------------------
+# Chroma DB
+# -----------------------------
 db = chromadb.PersistentClient(path="./chroma_db")
 
-collection = db.get_collection("guest_stories")
+# -----------------------------
+# Create collection if needed
+# -----------------------------
+try:
+    collection = db.get_collection("guest_stories")
+except Exception:
+    print("Chroma DB not found. Creating it...")
+
+    documents = SimpleDirectoryReader(
+        "data/guest_stories"
+    ).load_data()
+
+    splitter = SentenceSplitter(
+        chunk_size=512,
+        chunk_overlap=50,
+    )
+
+    nodes = splitter.get_nodes_from_documents(documents)
+
+    collection = db.get_or_create_collection("guest_stories")
+
+    vector_store = ChromaVectorStore(
+        chroma_collection=collection
+    )
+
+    storage_context = StorageContext.from_defaults(
+        vector_store=vector_store
+    )
 
+    VectorStoreIndex(
+        nodes,
+        storage_context=storage_context,
+        embed_model=embed_model,
+    )
+
+    print("Chroma DB created successfully.")
+
+# -----------------------------
+# Load vector store
+# -----------------------------
 vector_store = ChromaVectorStore(
     chroma_collection=collection
 )
 
-storage_context = StorageContext.from_defaults(
-    vector_store=vector_store
-)
-
-# Load the index from the vector store
 index = VectorStoreIndex.from_vector_store(
     vector_store=vector_store,
     embed_model=embed_model,
 )
 
-# Create the retriever
 retriever = index.as_retriever(similarity_top_k=1)
 
 
@@ -53,12 +95,6 @@ def search_guest_story(query: str) -> str:
     - event planning
     - laws or regulations
     - general internet searches
-
-    Args:
-        query: A question about an AI Gala guest.
-
-    Returns:
-        Relevant guest story text.
     """
 
     results = retriever.retrieve(query)