Spaces:
Runtime error
Runtime error
File size: 2,504 Bytes
4cf27a6 8096509 4cf27a6 8096509 4cf27a6 8096509 4cf27a6 8096509 4cf27a6 8096509 4cf27a6 8096509 d2f2f5c 8096509 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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) |