neuronslabs/ComfyKnowledgeGraph

1import os2import gradio as gr3from typing import List, Tuple4from gradio import ChatMessage5import base646from llama_index.core import StorageContext, load_index_from_storage7from dotenv import load_dotenv8from retrieve import get_latest_dir, get_latest_html_file9from graph_handler import query_graph_qa, plot_subgraph10from embed_handler import query_rag_qa11from evaluate import evaluate_llm, reasoning_graph, get_coupon12import base6413import time14 15load_dotenv()16 17KG_INDEX_PATH = get_latest_dir(os.getenv("GRAPH_DIR"))18KG_PLOT_PATH = get_latest_html_file(os.getenv("GRAPH_VIS"))19RAG_INDEX_PATH = get_latest_dir(os.getenv("EMBEDDING_DIR"))20 21# Load Graph-RAG index22graph_rag_index = load_index_from_storage(23    StorageContext.from_defaults(persist_dir=KG_INDEX_PATH)24)25 26# Load RAG index27rag_index = load_index_from_storage(28    StorageContext.from_defaults(persist_dir=RAG_INDEX_PATH)29)30 31 32def query_tqa(query, search_level):33    """34    Query the Graph-RAG and RAG models for a given query.35 36    Args:37    query (str): The query to ask the RAGs.38    search_level (int): The max search level to use for the Graph RAG.39 40    Returns:41    tuple: The response, reference, and reference text for the Graph-RAG and RAG models.42    """43 44    if not query.strip():45        raise gr.Error("Please enter a query before asking.")46 47    grag_response, grag_reference, grag_reference_text = query_graph_qa(48        graph_rag_index, query, search_level49    )50    # rag_response, rag_reference, rag_reference_text = query_rag_qa(51    #     rag_index, query, search_level52    # )53    print(str(grag_response.response))54    return (55        str(grag_response.response)56    )57 58 59 60 61# with gr.Blocks() as demo:62#     gr.Markdown("# Comfy Virtual Assistant")63#     chatbot = gr.Chatbot(64#         label="Comfy Virtual Assistant",65#         type="messages",66#         scale=1,67#         # suggestions = [68#         #         {"text": "How much iphone cost?"},69#         #         {"text": "What phone options do i have ?"}70#         #         ],71           72#     )73#     msg = gr.Textbox(label="Input Your Query")74#     clear = gr.ClearButton([msg, chatbot])75 76#     def respond(message, chat_history):77#             bot_message = query_tqa(message, 2)78#             # chat_history.append((message, bot_message))79#             chat_history.append(ChatMessage(role="user", content=message))80#             chat_history.append(ChatMessage(role="assistant", content=bot_message))81#             time.sleep(1)82#             return "", chat_history83    84#     msg.submit(respond, [msg, chatbot], [msg, chatbot])85 86def chatbot_response(message: str, history: List[Tuple[str, str]]) -> str:87    # Use the query_tqa function to get the response88    search_level = 2  # You can adjust this or make it configurable89    response = query_tqa(message, search_level)90    return response91 92# Create the Gradio interface93with gr.Blocks() as demo:94    chatbot = gr.Chatbot()95    msg = gr.Textbox()96    clear = gr.Button("Clear")97 98    def user(user_message, history):99        return "", history + [[user_message, None]]100 101    def bot(history):102        bot_message = chatbot_response(history[-1][0], history)103        history[-1][1] = bot_message104        return history105 106    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(107        bot, chatbot, chatbot108    )109    clear.click(lambda: None, None, chatbot, queue=False)110 111    112demo.launch(auth=(os.getenv("ID"), os.getenv("PASS")), share=False)113