tools/rag_tool.py

import os
import chromadb

from smolagents import tool

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


# -----------------------------
# Embedding model
# -----------------------------
embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-en-v1.5"
)

# -----------------------------
# Chroma DB
# -----------------------------
db = chromadb.PersistentClient(path="./chroma_db")

# -----------------------------
# 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
)

index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store,
    embed_model=embed_model,
)

retriever = index.as_retriever(similarity_top_k=1)


@tool
def search_guest_story(query: str) -> str:
    """
    Search AI Gala guest stories and return relevant information.

    Use this tool ONLY for:
    - AI Gala guests
    - speakers
    - guest biographies
    - guest backgrounds
    - guest stories

    Do NOT use this tool for:
    - weather information
    - event locations
    - 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)

    if not results:
        return "No relevant guest story found."

    return "\n\n".join(node.text for node in results)