Spaces:
Sleeping
Sleeping
File size: 1,504 Bytes
2a51e7d |
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 |
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 apps.agent.constant import INDEX_NAME_WEWEB, INDEX_NAME_XANO
# os.environ["PINECONE_API_KEY"] = "a526d62f-ccca-40d6-859b-3d878c8d288b"
embeddings = FastEmbedEmbeddings(model_name="BAAI/bge-small-en-v1.5")
compressor = FlashrankRerank()
def create_compressed_retriever(index_name: str, embeddings, compressor) -> ContextualCompressionRetriever:
vectorstore = Pinecone.from_existing_index(embedding=embeddings, index_name=index_name)
retriever = vectorstore.as_retriever()
return ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
reranker_xano = create_compressed_retriever(INDEX_NAME_XANO, embeddings, compressor)
reranker_weweb = create_compressed_retriever(INDEX_NAME_WEWEB, embeddings, compressor)
@tool
def tool_xano(query: str):
"""
Searches and returns excerpts from the Xano documentation
"""
docs = reranker_xano.invoke(query)
return "\n\n".join([doc["page_content"] for doc in docs])
@tool
def tool_weweb(query: str):
"""
Searches and returns excerpts from the Weweb documentation
"""
docs = reranker_weweb.invoke(query)
return "\n\n".join([doc["page_content"] for doc in docs])
|