import os
from omegaconf import OmegaConf
from vectara_agentic.agent import Agent

from dotenv import load_dotenv
load_dotenv(override=True)

initial_prompt = "How can I help you today?"


def initialize_agent(_cfg, agent_progress_callback=None):
    agent = Agent.from_corpus(
        vectara_corpus_key=_cfg.corpus_key,
        vectara_api_key=_cfg.api_key,
        tool_name="ask_ucsf_ortho",
        data_description="UCSF Orthopedic Website",
        assistant_specialty="UCSF Orthopedic department",
        vectara_summarizer="vectara-summary-ext-24-05-med-omni",
        # vectara_reranker = "chain", vectara_rerank_k = 100, 
        # vectara_rerank_chain = [
        #     {
        #         "type": "multilingual_reranker_v1",
        #         "cutoff": 0.5,
        #         "limit": 100
        #     },
        #     {
        #         "type": "mmr",
        #         "diversity_bias": 0.1
        #     }
        # ],
        vectara_reranker="multilingual_reranker_v1",
        vectara_rerank_k=100,
        vectara_lambda_val=0.005,
        vectara_summary_num_results=25,
        verbose=False,
        agent_progress_callback=agent_progress_callback,
    )
    agent.report()
    return agent


def get_agent_config() -> OmegaConf:
    cfg = OmegaConf.create({
        'corpus_key': str(os.environ['VECTARA_CORPUS_KEY']),
        'api_key': str(os.environ['VECTARA_API_KEY']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "UCSF Ortho Demo",
        'demo_welcome': "",
        'demo_description': "This assistant can help you with any questions about UCSF Orthopedic department."
    })
    return cfg