Code
retriever.py
from typing import Optional
from agno.agent import Agent
from agno.knowledge.embedder.openai import OpenAIEmbedder
from agno.knowledge.knowledge import Knowledge
from agno.vectordb.qdrant import Qdrant
from qdrant_client import QdrantClient
# ---------------------------------------------------------
# This section loads the knowledge base. Skip if your knowledge base was populated elsewhere.
# Define the embedder
embedder = OpenAIEmbedder(id="text-embedding-3-small")
# Initialize vector database connection
vector_db = Qdrant(
collection="thai-recipes", url="http://localhost:6333", embedder=embedder
)
# Load the knowledge base
knowledge = Knowledge(
vector_db=vector_db,
)
knowledge.add_content(
url="https://agno-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf",
)
# ---------------------------------------------------------
# Define the custom knowledge retriever
# This is the function that the agent will use to retrieve documents
def knowledge_retriever(
query: str, agent: Optional[Agent] = None, num_documents: int = 5, **kwargs
) -> Optional[list[dict]]:
"""
Custom knowledge retriever function to search the vector database for relevant documents.
Args:
query (str): The search query string
agent (Agent): The agent instance making the query
num_documents (int): Number of documents to retrieve (default: 5)
**kwargs: Additional keyword arguments
Returns:
Optional[list[dict]]: List of retrieved documents or None if search fails
"""
try:
qdrant_client = QdrantClient(url="http://localhost:6333")
query_embedding = embedder.get_embedding(query)
results = qdrant_client.query_points(
collection_name="thai-recipes",
query=query_embedding,
limit=num_documents,
)
results_dict = results.model_dump()
if "points" in results_dict:
return results_dict["points"]
else:
return None
except Exception as e:
print(f"Error during vector database search: {str(e)}")
return None
def main():
"""Main function to demonstrate agent usage."""
# Initialize agent with custom knowledge retriever
# Remember to set search_knowledge=True to use agentic_rag or add_reference=True for traditional RAG
# search_knowledge=True is default when you add a knowledge base but is needed here
agent = Agent(
knowledge_retriever=knowledge_retriever,
search_knowledge=True,
instructions="Search the knowledge base for information",
)
# Example query
query = "List down the ingredients to make Massaman Gai"
agent.print_response(query, markdown=True)
if __name__ == "__main__":
main()
Usage
1
Install libraries
pip install -U agno openai qdrant-client
2
Run Qdrant
docker run -p 6333:6333 qdrant/qdrant
3
Set OpenAI API key
export OPENAI_API_KEY=your_openai_api_key_here
4
Run the example
python cookbook/knowledge/custom_retriever/retriever.py
python cookbook/knowledge/custom_retriever/retriever.py