Spaces:
Runtime error
Runtime error
File size: 4,859 Bytes
cd5abb0 a0ae9e2 cd5abb0 336f222 cd5abb0 |
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 |
from typing import Any, Optional, List, Dict
from smolagents.tools import Tool
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
import json
class RAGSearchTool(Tool):
name = "rag_search"
description = "A RAG (Retrieval Augmented Generation) tool that can store documents and perform semantic search using FAISS and GTE-small embeddings."
inputs = {
'action': {'type': 'string', 'description': 'The action to perform: "add" to add documents, "search" to search existing documents, or "clear" to clear the database'},
'content': {'type': 'string', 'description': 'For "add" action: the text content to add to the database. For "search" action: the query to search for. For "clear" action: can be empty.'},
'metadata': {'type': 'object', 'description': 'Optional metadata for the documents when adding them', 'nullable': True}
}
output_type = "string"
def __init__(self, persist_dir="rag_db", **kwargs):
super().__init__()
self.persist_dir = persist_dir
try:
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
except ImportError as e:
raise ImportError(
"You must install packages `faiss-cpu` and `sentence-transformers` to run this tool: "
"run `pip install faiss-cpu sentence-transformers`."
) from e
# Initialize the embedding model
self.embeddings = HuggingFaceEmbeddings(
model_name="thenlper/gte-small",
model_kwargs={'device': 'cpu'},
encode_kwargs={'normalize_embeddings': True}
)
# Initialize text splitter for chunking
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50,
length_function=len,
)
# Try to load existing database or create new one
if os.path.exists(os.path.join(persist_dir, "index.faiss")):
self.db = FAISS.load_local(persist_dir, self.embeddings, allow_dangerous_deserialization=True)
else:
self.db = None
self.is_initialized = True
def forward(self, action: str, content: str, metadata: Optional[Dict] = None) -> str:
try:
if action == "add":
if not content:
return "Error: No content provided to add to the database."
# Split text into chunks
chunks = self.text_splitter.split_text(content)
# Create metadata for each chunk if provided
metadatas = [metadata] * len(chunks) if metadata else None
# Initialize db if it doesn't exist
if self.db is None:
self.db = FAISS.from_texts(chunks, self.embeddings, metadatas=metadatas)
else:
self.db.add_texts(chunks, metadatas=metadatas)
# Save the updated database
self.db.save_local(self.persist_dir)
return f"Successfully added {len(chunks)} chunks to the database."
elif action == "search":
if not content:
return "Error: No search query provided."
if self.db is None:
return "Error: No documents have been added to the database yet."
# Perform similarity search
results = self.db.similarity_search_with_score(content, k=3)
# Format results
formatted_results = []
for doc, score in results:
result = {
'content': doc.page_content,
'metadata': doc.metadata,
'similarity_score': float(score) # Convert numpy float to Python float
}
formatted_results.append(result)
return json.dumps(formatted_results, indent=2)
elif action == "clear":
if os.path.exists(self.persist_dir):
import shutil
shutil.rmtree(self.persist_dir)
self.db = None
return "Database cleared successfully."
else:
return f"Error: Invalid action '{action}'. Valid actions are 'add', 'search', or 'clear'."
except Exception as e:
return f"Error performing {action} operation: {str(e)}"
|