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])