import os
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.agents import initialize_agent, AgentType
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.tools import Tool
from langchain.tools import DuckDuckGoSearchRun
from langchain_core.documents import Document
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv

load_dotenv()  

apikey = os.getenv("MISTRAL_API_KEY")

embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
faiss_index = FAISS.load_local("faiss_index", embedding_model, allow_dangerous_deserialization=True)
retriever = faiss_index.as_retriever(search_kwargs={"k": 3})

llm = ChatOpenAI(
    openai_api_key=apikey,
    openai_api_base="https://api.mistral.ai/v1",
    model="mistral-medium"
)

from langchain.prompts import PromptTemplate

prompt_template = PromptTemplate(
    input_variables=["user_prompt", "retriever_query"],
    template="""You are a summarizer agent tasked with generating a concise summary of academic papers retrieved from a RAG pipeline. Using the user’s prompt and the retriever query used for searching, fetch relevant document content with the RAG tool and summarize the key points, findings, or insights relevant to the user’s request.

User Input Prompt: {user_prompt}
Retriever Query: {retriever_query}

Instructions:
1. Use the RAG tool to retrieve document content based on the retriever query.
2. Analyze the user prompt to identify the focus or specific aspects of interest.
3. Summarize the main ideas, results, or trends from the retrieved documents, excluding unnecessary details.
4. Ensure the summary is clear, coherent, and no longer than 1500 words.
5. If the RAG tool returns irrelevant or insufficient documents, note this briefly and provide a general response based on the prompt.
6. Output only the summary as a string.

Example:
- User Prompt: "Recent advancements in large language models for natural language processing"
- Retriever Query: "large language models NLP advancements"
- Summary: "Recent advancements in large language models (LLMs) focus on improved efficiency and performance in NLP tasks. Techniques like fine-tuning and transformer architectures enhance accuracy in text generation and understanding."

Generate the summary for the provided user prompt and retriever query.
"""
)
search_tool = DuckDuckGoSearchRun()

tools = [
    Tool(
    name="FAISSRetriever",
    func=lambda query: "\n\n".join([doc.page_content for doc in retriever.invoke(query)]),
    description="Fetches the top 3 relevant document chunks from a FAISS index containing academic paper content based on a query. Use for retrieving scholarly text for summarization."
    ),
]
summarizer_agent = initialize_agent(
    tools=tools,
    agent_type=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    llm=llm,
    verbose=True,
    prompt=prompt_template,
    retriever=retriever
)