from llmConnect import managerAgent, vectorDBAgent, searchAgent, sqlAgent, answerAgent
from json import loads




def generateAnswer(question:str) -> str:
    managerResponse: list= managerAgent(question)
    managerResponse.sort(key=lambda x: x['agent'])
    #print(managerResponse)
    context = ""
    agentsResponseCummalative =[{"Question to be answered":question}]
    for i in managerResponse:
        if i['agent'] == 'vectorDBAgent':
            if len(context) == 0:
                context = "No context Required"
            prompt = f"""
            context -> {context}
            question -> {i['question']}
            """
            agentResponse = vectorDBAgent(prompt)
            context += " " + str(agentResponse)
            agentsResponseCummalative.append({"agent":"vectorDBAgent","prompt":prompt,"answer":agentResponse})
        elif i['agent'] == 'searchAgent':
            prompt = f"""
            context -> {context}
            question -> {i['question']}
            """
            agentResponse = searchAgent(prompt)
            context += " " + str(agentResponse)  
            agentsResponseCummalative.append({"agent":"searchAgent","prompt":prompt,"answer":agentResponse})
        elif i['agent'] == 'sqlAgent':
            prompt = f"""{i['question']}
            """
            agentResponse: dict = sqlAgent(prompt)
            context += " " + str(agentResponse)
            agentsResponseCummalative.append({"agent":"sqlAgent","prompt":prompt,"answer":agentResponse})
            
        else:
            agentsResponseCummalative.append(i)
    
    for i in agentsResponseCummalative:
        print(i)
        print("\n")
        print("."*100)
        print("\n")
           
    return answerAgent(str(agentsResponseCummalative))


if __name__ == "__main__":
    question = "What is the highest score of MS Dhoni in the IPL"
    print(generateAnswer(question))


    
        #if "requestingAgent" in agentResponse.keys():
            #if agentResponse["requestingAgent"] == "searchAgent":
                #ans = vectorDBAgent(agentResponse["question"])
                #newPrompt = f"""
                #context -> {agentResponse["question"]} + " " + {ans}
                #question -> {question}
                #"""

                #agentResponse = loads(sqlAgent(newPrompt))
                #print(agentResponse)


#if "requestingAgent" in agentResponse.keys():
                #if agentResponse["requestingAgent"] == "searchAgent":
                    #ans = vectorDBAgent(agentResponse["question"])
                    #newPrompt = f"""
                    #context -> {agentResponse["question"]} + " " + {ans}
                    #question -> {question}
                    #"""

                    #agentResponse = loads(sqlAgent(newPrompt))
                    #print(agentResponse)