Spaces:
Sleeping
Sleeping
File size: 2,910 Bytes
e0e5454 |
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 |
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
from smolagents import CodeAgent, HfApiModel
class PartyPlanningRetrieverTool(Tool):
name = "party_planning_retriever"
description = "Uses semantic search to retrieve relevant party planning ideas for Alfred’s superhero-themed party at Wayne Manor."
inputs = {
"query": {
"type": "string",
"description": "The query to perform. This should be a query related to party planning or superhero themes.",
}
}
output_type = "string"
def __init__(self, docs, **kwargs):
super().__init__(**kwargs)
self.retriever = BM25Retriever.from_documents(
docs, k=5 # Retrieve the top 5 documents
)
def forward(self, query: str) -> str:
assert isinstance(query, str), "Your search query must be a string"
docs = self.retriever.invoke(
query,
)
return "\nRetrieved ideas:\n" + "".join(
[
f"\n\n===== Idea {str(i)} =====\n" + doc.page_content
for i, doc in enumerate(docs)
]
)
# Simulate a knowledge base about party planning
party_ideas = [
{"text": "A superhero-themed masquerade ball with luxury decor, including gold accents and velvet curtains.", "source": "Party Ideas 1"},
{"text": "Hire a professional DJ who can play themed music for superheroes like Batman and Wonder Woman.", "source": "Entertainment Ideas"},
{"text": "For catering, serve dishes named after superheroes, like 'The Hulk's Green Smoothie' and 'Iron Man's Power Steak.'", "source": "Catering Ideas"},
{"text": "Decorate with iconic superhero logos and projections of Gotham and other superhero cities around the venue.", "source": "Decoration Ideas"},
{"text": "Interactive experiences with VR where guests can engage in superhero simulations or compete in themed games.", "source": "Entertainment Ideas"}
]
source_docs = [
Document(page_content=doc["text"], metadata={"source": doc["source"]})
for doc in party_ideas
]
# Split the documents into smaller chunks for more efficient search
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50,
add_start_index=True,
strip_whitespace=True,
separators=["\n\n", "\n", ".", " ", ""],
)
docs_processed = text_splitter.split_documents(source_docs)
# Create the retriever tool
party_planning_retriever = PartyPlanningRetrieverTool(docs_processed)
# Initialize the agent
agent = CodeAgent(tools=[party_planning_retriever], model=HfApiModel())
# Example usage
response = agent.run(
"Find ideas for a luxury superhero-themed party, including entertainment, catering, and decoration options."
)
print(response)
|