apps/agent/tools.py

import os
from langchain_community.vectorstores.pinecone import Pinecone
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import FlashrankRerank
from langchain_core.tools import tool
from langchain.retrievers.multi_query import MultiQueryRetriever

from apps.agent.multi_query_chain import llm_chain
from apps.agent.constant import INDEX_NAME_WEWEB, INDEX_NAME_XANO


embeddings = FastEmbedEmbeddings(model_name="jinaai/jina-embeddings-v2-small-en")
compressor = FlashrankRerank(model="ms-marco-MiniLM-L-12-v2")

def multiquery_retriever(index_name: str, embeddings, compressor) -> ContextualCompressionRetriever:
    vectorstore = Pinecone.from_existing_index(embedding=embeddings, index_name=index_name)
    retriever = vectorstore.as_retriever()
    reranker_retriever = ContextualCompressionRetriever(
        base_compressor=compressor, base_retriever=retriever
    )
    print("Initialize MultiQuery...")
    return MultiQueryRetriever(
        retriever=reranker_retriever, llm_chain=llm_chain, parser_key="lines"
    )

retriever_xano = multiquery_retriever(INDEX_NAME_XANO, embeddings, compressor)
retriever_weweb = multiquery_retriever(INDEX_NAME_WEWEB, embeddings, compressor)

@tool
def tool_xano(query: str):
    """
    Searches and returns excerpts from the Xano documentation
    """
    return retriever_xano.get_relevant_documents(query)
    

@tool
def tool_weweb(query: str):
    """
    Searches and returns excerpts from the Weweb documentation
    """
    return retriever_weweb.get_relevant_documents(query)