from utils.retriever import get_query_engine
from utils.index import create_index
from utils.constant import INDEX_PATH, DATA_PATH
import os
import shutil

def set_keys(openai_api_key: str, llama_cloud_api_key: str) -> str:
    try:
        #os.environ["CO_API_KEY"] = co_api_key
        os.environ["OPENAI_API_KEY"] = openai_api_key
        os.environ["LLAMA_CLOUD_API_KEY"] = llama_cloud_api_key
        return "Keys are set successfully"
    
    except Exception as e:
        return str(e)

def handle_file(uploaded_file):
    try:

        temp_file_path = uploaded_file.name
        save_file_path = os.path.join(DATA_PATH, "paper.pdf")
        
        shutil.move(temp_file_path, save_file_path)
        
        create_index(DATA_PATH, INDEX_PATH)
        return "File uploaded and indexed"
    
    except Exception as e:
        return str(e)

def answer_query(query: str) -> str:
    
    query_engine = get_query_engine(semantic=True)
    response = query_engine.query(query)
    
    nodes = query_engine.retriever.retrieve(query)
    
    for node in nodes:

        score = node.get_score()
        text = node.text
        
        response += f"\nNode: {node.node_id}\nScore: {score:0.3f}\nText: {text}\n"
    
    return response

if __name__ == "__main__":
    
    query = "What is Few-NERD"
    response = answer_query(query)
    print(response)