Spaces:
Sleeping
Sleeping
File size: 1,360 Bytes
08e2c16 |
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 |
# Knowledge retrieval tool using BM25
from langchain_community.retrievers import BM25Retriever
from smolagents import Tool
class GaiaRetrieverTool(Tool):
name = "gaia_retriever"
description = "Semantic search for retrieving relevant information for GaiaAgent."
inputs = {
"query": {
"type": "string",
"description": "Query for semantic search."
}
}
output_type = "string"
def __init__(self, docs, **kwargs):
super().__init__(**kwargs)
self.retriever = BM25Retriever.from_documents(docs, k=3)
self.docs = docs # Store docs for fallback
def forward(self, query: str) -> str:
assert isinstance(query, str), "Query must be a string."
try:
docs = self.retriever.invoke(query)
if not docs:
return "\nNo specific information found. Here's some general knowledge:\n" + "".join([
f"\n- {self.docs[i].page_content}" for i in range(min(3, len(self.docs)))
])
return "\nRetrieved Information:\n" + "".join([
f"\n- {doc.page_content}" for doc in docs
])
except Exception as e:
print(f"Error in retriever: {str(e)}")
return f"Unable to retrieve specific information. The agent will rely on its general knowledge."
|