function_tools.py

from typing import Optional
import json
import logging
from utils import iast_process, get_list_meaning_word, get_details_mantra_json, word_sentence_similarity, extract_meaning_by_language
import ast
from langchain_core.tools import tool 
from database import execute_sql_query, get_details_mantra_json
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import StructuredTool
from typing import Optional
import streamlit as st
from langchain_core.utils.function_calling import convert_to_openai_function
from langchain_core.messages import AIMessage
from langchain_core.runnables import Runnable
import os

os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"]
from langchain_openai import ChatOpenAI

#LLM
llm = ChatOpenAI(model="gpt-3.5-turbo-0125")
llm_AI4 = ChatOpenAI(model="gpt-4-1106-preview", temperature=0)



#Classes
class MantraInput(BaseModel):
    mantraid: Optional[str] = Field(None, description="The mantra id. For example, 1.1.1.1, 2.1.1,3.1.1.2,4.2.3.1, and 5.0.1.1.2")
    scripture_name: Optional[str] = Field(None, description="Name of the scripture like RigVeda, SamaVeda, AtharvaVeda, KrishnaYajurVeda, and ShuklaYajurVeda")
    KandahNumber: Optional[int] = Field(None, description="Kandah Number of Vedamantra")
    MandalaNumber: Optional[int] = Field(None, description="Mandala Number of Vedamantra")
    ArchikahNumber: Optional[int] = Field(None, description="Archikah Number of Vedamantra")
    ShuktaNumber: Optional[int] = Field(None, description="Shukta Number of Vedamantra")
    PrapatakNumber: Optional[int] = Field(None, description="Prapatak Number of Vedamantra")
    MantraNumber: Optional[int] = Field(None, description="Mantra Number of Vedamantra")    
    AnuvakNumber: Optional[int] = Field(None, description="Anuvak Number of Vedamantra")
    AdhyayaNumber: Optional[int] = Field(None, description="Adhyaya Number of Vedamantra")

class PadaMeaningInput(BaseModel):
    pada: str = Field(description="The pada or word that is being meaning checked")

class PadaAAAInput(BaseModel):
    pada: str = Field(description="The pada or word that is being meaning checked")
    mantraid: Optional[str] = Field(None, description="The mantra id. For example, 1.1.1.1, 2.1.1,3.1.1.2,4.2.3.1, and 5.0.1.1.2")
    scripture_name: Optional[str] = Field(None, description="Name of the scripture like RigVeda, SamaVeda, AtharvaVeda, KrishnaYajurVeda, and ShuklaYajurVeda")
    KandahNumber: Optional[int] = Field(None, description="Kandah Number of Vedamantra")
    MandalaNumber: Optional[int] = Field(None, description="Mandala Number of Vedamantra")
    ArchikahNumber: Optional[int] = Field(None, description="Archikah Number of Vedamantra")
    ShuktaNumber: Optional[int] = Field(None, description="Shukta Number of Vedamantra")
    PrapatakNumber: Optional[int] = Field(None, description="Prapatak Number of Vedamantra")
    MantraNumber: Optional[int] = Field(None, description="Mantra Number of Vedamantra")    
    AnuvakNumber: Optional[int] = Field(None, description="Anuvak Number of Vedamantra")
    AdhyayaNumber: Optional[int] = Field(None, description="Adhyaya Number of Vedamantra")

class NLSQLResponse(BaseModel):
    user_query: str = Field(description="user query")

class VectorResponse(BaseModel):
    query: str = Field(description="User query")

class Response(BaseModel):
    result: str = Field(description="The result based on the context. Provide the text in a readable format if there are unicode characters. Use only available context. If there is no context, return as 'unknown'. Do not use prior knowledge.")
    explanation: str = Field(description="Explanation of the steps taken to get the result")

#function tools for mantra level
def _get_mantra_details(query):
    try:
        details = get_details_mantra_json(query)
        return details['mantraHeader']['language'][1]
    except Exception as e:
        raise ValueError(f"Failed to get mantra details: {e}")


def _get_mantra_details_by_scripture(scripture_name=None, KandahNumber=None, MandalaNumber=None, ArchikahNumber=None,
                                      ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None, AnuvakNumber=None,
                                      AdhyayaNumber=None):
    try:
        # Construct the base SQL query
        query = "SELECT * FROM veda_content_details WHERE 1 = 1"
        parameters = []

        # Add conditions based on provided parameters
        if scripture_name:
            query += " AND scripture_name = %s"
            parameters.append(scripture_name.lower())
        if KandahNumber:
            query += " AND KandahNumber = %s"
            parameters.append(KandahNumber)
        if MandalaNumber:
            query += " AND MandalaNumber = %s"
            parameters.append(MandalaNumber)
        if ArchikahNumber:
            query += " AND ArchikahNumber = %s"
            parameters.append(ArchikahNumber)
        if ShuktaNumber:
            query += " AND ShuktaNumber = %s"
            parameters.append(ShuktaNumber)
        if PrapatakNumber:
            query += " AND PrapatakNumber = %s"
            parameters.append(PrapatakNumber)
        if MantraNumber:
            query += " AND MantraNumber = %s"
            parameters.append(MantraNumber)
        if AnuvakNumber:
            query += " AND AnuvakNumber = %s"
            parameters.append(AnuvakNumber)
        if AdhyayaNumber:
            query += " AND AdhyayaNumber = %s"
            parameters.append(AdhyayaNumber)

        # Execute the SQL query
        results = execute_sql_query(query, parameters)

        if results:
            return results
        else:
            return None
    except Exception as e:
        logging.error(f"Error in _get_mantra_details_by_scripture: {e}")



def get_vedamantra_details(mantraid=None, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                        ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                        AnuvakNumber=None, AdhyayaNumber=None):
    '''
    This function is used to get the vedamantra such as vedamantra, padapatha, devata, chandah, 
    and rishi, from all Vedas (RigVeda, AtharvaVeda, SamaVeda, KrishnaYajurVeda, and ShuklaYajurVeda).
    The Vedic scriptures has the following structure: \
    RigVeda->Mandala->Shukta->Mantra\
    SamaVeda->Archikah->Shukta->Mantra\
    AtharvaVeda->Kandah->Shukta->Mantra\
    ShuklaYajurVeda->Adhyaya->Mantra\
    KrishnaYajurVeda->Kandah->Prapatak->Anuvak->Mantra\
    Sample Questions:
    1. Obtain the vedamantra of the mantra whose id is 1.1.1.1?
    2. Retrieve the devata of the vedamantra from Rigveda, first mandala, first shukta, and first mantra.
    3. Provide the meaning of the vedamantra from Rigveda, first mandala, first shukta, and first mantra written by Tulsi Ram.
    4. Explain the adhibautic meaning of the first mantra from RigVeda, first mandala, and first shukta.
    5. Identify the mantraVishaya of the vedamantra from RigVeda, first mandala, first shukta, and first mantra.
    6. What is the adibhautic meaning of the mantra 1.1.1.9?
    7. What is the adhyatmic meaning of the mantra 1.1.1.7?
    8. What is the adhidyvic meaning of the 6th mantra from RigVeda, first mandala, and first shukta?
    '''
    try:
        query = ""  # Initialize query outside of the if-else block
        if mantraid:
            query = f'''SELECT mantra_json FROM veda_content WHERE mantra_number = "{mantraid}"
            '''
        else:
            filter_df = _get_mantra_details_by_scripture(scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                        ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                        AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)
            if filter_df is not None:
                mantra_id = filter_df[0]['mantra_id']
                query = f'''SELECT mantra_json FROM veda_content WHERE mantra_number = "{mantra_id}"
                '''
        return _get_mantra_details(query)
    except Exception as e:
        return json.dumps({"error": str(e)})

def get_vedamantra_summary(mantraid=None, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                        ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                        AnuvakNumber=None, AdhyayaNumber=None):
    '''
    Use the function `get_vedamantra_summary` to access the information such as adibhautic meaning of the mantra, anvaya of the mantra,  mantraVishaya of the mantra, 
    adhibautic (or adhyatmic or adhidyvic) meaning (or bhavarth) of the mantra, purpose of the mantra, usage of the mantra, and tippani of the mantra.
    Sample Query:
    1. Obtain the anvaya of the mantra whose id (mantraid) is 1.1.1.1?
    2. Retrieve tha adibhautic meaning of the first mantra from RigVeda, first mandala, and first shukta.
    '''
    try:
        if mantraid:
            query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantraid}'"
        else:
            filtered_df = _get_mantra_details_by_scripture(scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                        ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                        AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)   
            if filtered_df is not None:
                mantra_id = filtered_df[0]['mantra_id']
                query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantra_id}'"
            else:
                return None      
        json_dict = get_details_mantra_json(query)
        mantra_summary = json_dict['mantraSummary']['language']
        summary_dict = {"Roman-IAST summary of vedamantra": json_dict['mantraSummary']['language'][1]}
        for item in mantra_summary:
            if item['languageName'] == 'English':
                mahatma = item['mahatma']['mahatmaName']
                summary_dict[f"English summary of vedamantra by {mahatma}"] = item
        return summary_dict
    except Exception as e:
        return {"error": str(e)}



def get_pada_meaning(pada):
    '''
    Purpose: For given sanskrit word, you have collection of meanings for available roots and stems of it.\
    You need to process this information as context and provide possible meanings for given word.
    Sample query:
    1. What is the meaning of the word apratidhṛṣṭa-śavasam?
    '''	
    #pada=iast_process(pada)
    try:
        query = f'''
        SELECT * FROM term_details_modified WHERE Pada = "{pada}"
        '''
            # Execute the query to get details from the database
        details = execute_sql_query(query)
        #print(details)
        pada_details = details[0]
        #print(pada_details['Morphology'])
        meanings_list = []
        for morphs in ast.literal_eval(pada_details['Morphology']):
            for field in ['stem', 'root']:
                word = morphs.get(field)
                if word:
                    meanings_list.append(get_list_meaning_word(word))
        return meanings_list
    except Exception as e:
        logging.error(f"Error in get_pada_meaning: {e}")
        return {"error": f"Required meaning associated with pada is not available. {e}"} 
    


def _get_pada_details_by_scripture(pada, scripture_name=None, KandahNumber=None, MandalaNumber=None, ArchikahNumber=None,
                                   ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None, AnuvakNumber=None,
                                   AdhyayaNumber=None):
    try:
        # Construct the base SQL query
        query = "SELECT * FROM term_details_modified WHERE Pada = %s"
        parameters = [pada]

        # Add conditions based on provided parameters
        if scripture_name:
            query += " AND scripture_name = %s"
            parameters.append(scripture_name)
        if KandahNumber:
            query += " AND KandahNumber = %s"
            parameters.append(KandahNumber)
        if MandalaNumber:
            query += " AND MandalaNumber = %s"
            parameters.append(MandalaNumber)
        if ArchikahNumber:
            query += " AND ArchikahNumber = %s"
            parameters.append(ArchikahNumber)
        if ShuktaNumber:
            query += " AND ShuktaNumber = %s"
            parameters.append(ShuktaNumber)
        if PrapatakNumber:
            query += " AND PrapatakNumber = %s"
            parameters.append(PrapatakNumber)
        if MantraNumber:
            query += " AND MantraNumber = %s"
            parameters.append(MantraNumber)
        if AnuvakNumber:
            query += " AND AnuvakNumber = %s"
            parameters.append(AnuvakNumber)
        if AdhyayaNumber:
            query += " AND AdhyayaNumber = %s"
            parameters.append(AdhyayaNumber)

        # Execute the SQL query
        results = execute_sql_query(query, parameters)

        if results:
            return results
        else:
            return None

    except Exception as e:
        logging.error(f"Error in _get_pada_details_by_scripture: {e}")
        return None
    
def _get_vedamantra_meaning(mantraID, MahatmaName=None):
    try:
        query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantraID}'"
        jsonDict = get_details_mantra_json(query)
        mantraSummary = jsonDict['mantraSummary']['language']
        if MahatmaName is not None:
            filtered_summary = [data_dict for data_dict in mantraSummary if data_dict.get('mahatma', {}).get('mahatmaName') == MahatmaName]
            if filtered_summary:
                mantraSummary = filtered_summary
        best_meaning = None
        best_count = 0
        for data_dict in mantraSummary:
            if data_dict.get('languageName') == "English":
                meanings = data_dict['mahatma']['bhavartha']
                count = sum(bool(meanings.get(cat, None)) for cat in ['adibhautic', 'adidaivic', 'adhyatmic'])
                if count >= best_count:
                    best_meaning = {cat: meanings.get(cat, None) for cat in ['adibhautic', 'adidaivic', 'adhyatmic']}
                    best_count = count
        return best_meaning if best_meaning else json.dumps({"error": "Required meaning associated with vedamantra is not available."})
    except Exception as e:
        logging.error(f"Error in _get_vedamantra_meaning: {e}")
        return json.dumps({"error": f"An error occurred: {e}"})

def _get_pada_morphology(term_details, meanings):
    try:
        morphology_list = ast.literal_eval(term_details['Morphology'])
        term_morph_list = []
        for morphs in morphology_list:
            term_info = {}
            for field in ['stem', 'root']:
                morph_word = morphs.get(field)
                if morph_word:
                    meaning = word_sentence_similarity(meanings, morph_word)
                    term_info[f'{field}_word'] = morph_word
                    term_info[f'{field}_meaning'] = meaning[0][0] if meaning else None
                    term_info[f'{field}_score'] = meaning[0][1] if meaning else None
            term_info['grammar'] = morphs['grammar']
            term_morph_list.append(term_info)
        return term_morph_list
    except Exception as e:
        logging.error(f"Error in _get_pada_morphology: {e}")
        return []
    
def get_morphological_info_of_pada(pada, mantraid=None, scripture_name=None, KandahNumber=None, MandalaNumber=None,
                                   ArchikahNumber=None, ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                   AnuvakNumber=None, AdhyayaNumber=None):
    '''
    This help to get segmentation and morphological information about the word.
    '''
    try:
        if pada:
            query = f'''SELECT * FROM term_details_modified WHERE Pada = "{pada}"
            '''
            details = execute_sql_query(query)
        else:
            # Placeholder for _get_pada_details_by_scripture function call
            # Replace with your actual implementation
            details = _get_pada_details_by_scripture(pada, scripture_name=scripture_name, KandahNumber=KandahNumber,
                                                           MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber,
                                                           ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber,
                                                           MantraNumber=MantraNumber, AnuvakNumber=AnuvakNumber,
                                                           AdhyayaNumber=AdhyayaNumber)
        
        if details:
            if mantraid is not None:
                for record in details:
                    if record["mantra_id"] == mantraid:
                        segmentation = record["Segmentation"]
                        morphological_info = record["Morphology"]
                        return {"morphology_info": {"segmentation": segmentation, "morphology": morphological_info}}
                return {"error": f"No details found for mantraid '{mantraid}'"}
            else:
                pada_details = details[0]
                segmentation = pada_details["Segmentation"]
                morphological_info = pada_details["Morphology"]
                return {"morphology_info": {"segmentation": segmentation, "morphology": morphological_info}}
        else:
            return {"error": "No details found for pada."}

    except Exception as e:
        logging.error(f"Error in get_morphological_info_of_pada: {e}")
        return {"error": f"Failed to get meaning of the word {pada}. {e}"}
        
    
def get_adibauatic_adidaivic_adhyatmic_meaning_of_pada(pada, mantraid=None, scripture_name=None, 
                                                        KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                        ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                        AnuvakNumber=None, AdhyayaNumber=None,MahatmaName=None):
    '''
    Sample query:
    1. What is the adibhautic meaning of pada 'agnim' from RigVeda, first mandala, first shukta and first mantra?
    2. What is the adhyatmic meaning of the pada agnim in the context of the mantra whose id is '1.1.1.1?'
    '''
    try:
        if mantraid:
            query = f'''
            SELECT * FROM term_details_modified WHERE mantra_id = '{mantraid}' AND Pada = "{pada}"
            '''
            # Execute the query to get details from the database
            details = execute_sql_query(query)
        else:
            # Call the function to get details by scripture
            details = _get_pada_details_by_scripture(pada, scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                        ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                        AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)
    
        if details:
            pada_details = details[0]  # Assuming details is a list of dictionaries, select the first item
            mantraID = pada_details['mantra_id']
            meanings = _get_vedamantra_meaning(mantraID,MahatmaName=MahatmaName)
            if 'error' in meanings:
                return json.dumps(meanings)
            ab_term_morph_list = _get_pada_morphology(pada_details, meanings['adibhautic'])
            ad_term_morph_list = _get_pada_morphology(pada_details, meanings['adidaivic'])
            at_term_morph_list = _get_pada_morphology(pada_details, meanings['adhyatmic'])
            return json.dumps({
                f'adibhautic_info_{pada}': ab_term_morph_list,
                'vedamantra_adibhautic_meaning': meanings['adibhautic'],
                f'adidavic_info_{pada}': ad_term_morph_list,
                'vedamantra_adidavic_meaning': meanings['adidaivic'],
                f'adhyatmic_info_{pada}': at_term_morph_list,
                'vedamantra_adhyatmic_meaning': meanings['adhyatmic']
            })
        else:
            return json.dumps({"error": f"No details found for pada '{pada}'"})
    except Exception as e:
        logging.error(f"Error in get_adibauatic_adidaivic_adhyatmic_meaning_of_pada: {e}")
        return json.dumps({"error": f"Failed to get meaning of the word {pada}. {e}"})
    
# sql agent
from langchain_community.utilities.sql_database import SQLDatabase
from database import get_db
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import StructuredTool
from typing import Optional
import json


class NLSQLResponse(BaseModel):
    user_query:str = Field(description="user query")

db = get_db()
agent_executor = create_sql_agent(llm_AI4, db=db, agent_type="openai-tools", verbose=False)

def get_response(user_query):
    response = agent_executor.invoke(user_query)
    return response

sql_tool  = StructuredTool.from_function(
    func = get_response,
    name = "nl_sql_query",
    description="""To obtains a information using natural language query to sql query and then exceting sql query to get natural response.
                Sample Query:
                1. How many mantras are there in RigVeda?
                2. What is the segmentation of the word 'prathasva' from KrishnaYajurVeda?""",
    args_schema=NLSQLResponse,
    return_direct=True,
)

pada_morphological_tool = StructuredTool.from_function(
    func=get_morphological_info_of_pada,
    name="pada_morphology",
    description="""Purpose: To obtain morphological information such as segmentation, morphology, and grammar of a word.\
                Sample query:
                    1. What is the segmentation and morphology of the word 'apratidhṛṣṭa-śavasam' from RigVeda?
                    2. What is the grammar of the word 'prathasva' from KrishnaYajurVeda?
                    """,
    args_schema=PadaAAAInput,
    return_direct=False
)

pada_meaning_tool = StructuredTool.from_function(
    func=get_pada_meaning,
    name="pada_meaning",
    description="""Purpose: For given sanskrit word, you have collection of meanings for available roots and stems of it.\
    You need to process this information as context and provide possible meanings for given word.
    Sample query:
    1. What is the meaning of the word apratidhṛṣṭa-śavasam?
    """,
    args_schema=PadaMeaningInput,
    return_direct=False
    )

pada_word_sense_tool = StructuredTool.from_function(
    func=get_adibauatic_adidaivic_adhyatmic_meaning_of_pada,
    name="pada_AAA_meaning",
    description="""To obtain a complete or meaningful adibauatic/adhidaivic/adhyatmic meaning of a word or pada based on context information.\n
                    Sample query:
                    1. What is the adibhautic meaning of pada 'agnim' from RigVeda, first mandala, first shukta and first mantra?
                    2. What is the adhyatmic meaning of the pada agnim in the context of the mantra whose id is '1.1.1.1'?
                    """,
    args_schema=PadaAAAInput,
    return_direct=False
    )

vedamantra_tool = StructuredTool.from_function(
    func=get_vedamantra_details,
    name="vedamantra_details",
    description='''This function is used to get the vedamantra such as vedamantra, padapatha, devata, chandah, 
    and rishi, from all Vedas (RigVeda, AtharvaVeda, SamaVeda, KrishnaYajurVeda, and ShuklaYajurVeda).
    Sample Questions:
    1. Obtain the vedamantra of the mantra whose id is 1.1.1.1?
    2. Retrieve the devata of the vedamantra from Rigveda, first mandala, first shukta, and first mantra.
    3. Provide the meaning of the vedamantra from Rigveda, first mandala, first shukta, and first mantra written by Tulsi Ram.
    4. Explain the adhibautic meaning of the first mantra from RigVeda, first mandala, and first shukta.
    ''',
    args_schema=MantraInput,
    return_direct=False
    )
vedamantra_summary_tool =StructuredTool.from_function(
    func=get_vedamantra_summary,
    name="vedamantra_summary",
    description="""Use the function `get_vedamantra_summary` to access the information such as adibhautic meaning of the mantra, anvaya of the mantra,  mantraVishaya of the mantra, 
    adhibautic (or adhyatmic or adhidyvic) meaning (or bhavarth) of the mantra, purpose of the mantra, usage of the mantra, and tippani of the mantra.
    Sample Query:
    1. Obtain the anvaya of the mantra whose id (mantraid) is 1.1.1.1?
    2. Retrieve tha adibhautic meaning of the first mantra from RigVeda, first mandala, and first shukta.
    3. Provide the adhyatmic meaning of the mantra 1.1.1.9?
    4. What is the tippani of the mantra 1.1.1.7?
    5. What is the adhyatmic meaning of the mantra 1.1.1.7?
    6. What is the mantravishaya of the 6th mantra from RigVeda, first mandala, and first shukta?""",
    args_schema=MantraInput,
    return_direct=False
    )

## vector tool
import os
import time
import pickle
import streamlit as st
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from utils import load_pickle, initialize_embedding_model
from langchain_community.retrievers import BM25Retriever
from langchain_pinecone import PineconeVectorStore
from langchain.retrievers import EnsembleRetriever
from langchain.tools.retriever import create_retriever_tool



# Load .env file
load_dotenv()

# Constants
INDEX_NAME = "veda-index-v2"
MODEL_NAME = "BAAI/bge-large-en-v1.5"
DOCS_DIRECTORY = r"Docs\ramana_docs_ids.pkl"
CURRENT_DIRECTORY = os.getcwd()


# Initialize Pinecone client
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY_SAM")
pc = Pinecone(api_key=PINECONE_API_KEY)

#@st.cache_resource
def create_or_load_index():
    # Check if index already exists
    if INDEX_NAME not in pc.list_indexes().names():
        # Create index if it does not exist
        pc.create_index(
            INDEX_NAME,
            dimension=1024,
            metric='dotproduct',
            spec=ServerlessSpec(
                cloud="aws",
                region="us-east-1"
            )
        )
        # Wait for index to be initialized
        while not pc.describe_index(INDEX_NAME).status['ready']:
            time.sleep(1)
    # Connect to index
    return pc.Index(INDEX_NAME)

# Load documents
docs = load_pickle(DOCS_DIRECTORY)
# Initialize embedding model
embedding = initialize_embedding_model(MODEL_NAME)
# Create or load index
index = create_or_load_index()

# Initialize BM25 retriever
bm25_retriever = BM25Retriever.from_texts(
    [text['document'].page_content for text in docs],
    metadatas=[text['document'].metadata for text in docs]
)
bm25_retriever.k = 2

# Switch back to normal index for LangChain
vector_store = PineconeVectorStore(index, embedding)
retriever = vector_store.as_retriever(search_type="mmr")

# Initialize the ensemble retriever
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, retriever], weights=[0.2, 0.8]
)

class VectorResponse(BaseModel):
    query:str = Field(description="user query")

def vector_retrieve(query):
    response = retriever.get_relevant_documents(query)
    return response

vector_tool = StructuredTool.from_function(
    func = vector_retrieve,
    name = "vector_retrieve",
    description="Search and return documents related user query from the vector index.",
    args_schema=VectorResponse,
    return_direct=False
)

tools_list = [pada_morphological_tool, sql_tool, pada_meaning_tool, pada_word_sense_tool, vedamantra_tool, vedamantra_summary_tool]
#vector_tool,

# Convert tools to OpenAI functions
tools_all = [convert_to_openai_function(tool) for tool in tools_list]
# Set up the tools to execute them from the graph
from langgraph.prebuilt import ToolExecutor

tool_executor = ToolExecutor(tools_list)
#tools_response = tools_all.append(convert_to_openai_function(Response))
llm_with_tools = llm_AI4.bind_tools(tools_all)
#tool_map = {tool.name: tool for tool in tools_list}


def call_tools(msg: AIMessage) -> Runnable:
    """Simple sequential tool calling helper."""
    tool_map = {tool.name: tool for tool in tools_list}
    tool_calls = msg.tool_calls.copy()
    for tool_call in tool_calls:
        tool_call["output"] = tool_map[tool_call["name"]].invoke(tool_call["args"])
    return tool_calls

#print("Invoking the chain")
tool_chain = llm_with_tools | call_tools