import pandas as pd
import requests
from pydantic import Field, BaseModel

from omegaconf import OmegaConf

from vectara_agentic.agent import Agent
from vectara_agentic.tools import ToolsFactory, VectaraToolFactory

def create_assistant_tools(cfg):
    class QueryDocsArgs(BaseModel):
        query: str = Field(..., description="The user query, always in the form of a question", 
                           examples=["Based on uploaded documents, what are the top four challenges of the Fintech sector in Saudi Arabia? list them in bullet points."])        
        
    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key,
                                     vectara_corpus_key=cfg.corpus_key)
    summarizer = 'mockingbird-1.0-2024-07-16'
    ask_docs = vec_factory.create_rag_tool(
        tool_name = "ask_docs",
        tool_description = """
        Responds to an user question about a particular analysis, based on the documentation provide.
        """,
        tool_args_schema = QueryDocsArgs,
        reranker = "chain", rerank_k = 100, 
        rerank_chain = [
            {
                "type": "multilingual_reranker_v1",
#                "cutoff": 0.2
            },
            {
                "type": "mmr",
                "diversity_bias": 0.2,
                "limit": 50
            }
        ],
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 15,
        vectara_summarizer = summarizer,
        include_citations = True,
        #vectara_prompt_text=prompt,
        save_history = True,
        verbose=False
    )
    tools_factory = ToolsFactory()
    return (
            tools_factory.standard_tools() +
            [ask_docs]
    )

def initialize_agent(_cfg, agent_progress_callback=None):
    stc_bank_bot_instructions = """
    - Call the the ask_docs tool to retrieve the information to answer the user query.
    - If the question has an 'Excel' or 'excel' word only fetch for the documents with 'type_file' equals to 'excel'.
    - Always print the title of the References
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="STC Bank questions",
        custom_instructions=stc_bank_bot_instructions,
        agent_progress_callback=agent_progress_callback,
    )
    agent.report()
    return agent