Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| import base64 | |
| from llama_index.core import StorageContext, load_index_from_storage | |
| from dotenv import load_dotenv | |
| from retrieve import get_latest_dir, get_latest_html_file | |
| from graph_handler import query_graph_qa, plot_subgraph | |
| from embed_handler import query_rag_qa | |
| from evaluate import evaluate_llm, reasoning_graph, get_coupon | |
| import base64 | |
| import time | |
| load_dotenv() | |
| KG_INDEX_PATH = get_latest_dir(os.getenv("GRAPH_DIR")) | |
| KG_PLOT_PATH = get_latest_html_file(os.getenv("GRAPH_VIS")) | |
| RAG_INDEX_PATH = get_latest_dir(os.getenv("EMBEDDING_DIR")) | |
| # Load Graph-RAG index | |
| graph_rag_index = load_index_from_storage( | |
| StorageContext.from_defaults(persist_dir=KG_INDEX_PATH) | |
| ) | |
| # Load RAG index | |
| rag_index = load_index_from_storage( | |
| StorageContext.from_defaults(persist_dir=RAG_INDEX_PATH) | |
| ) | |
| def query_tqa(query, search_level): | |
| """ | |
| Query the Graph-RAG and RAG models for a given query. | |
| Args: | |
| query (str): The query to ask the RAGs. | |
| search_level (int): The max search level to use for the Graph RAG. | |
| Returns: | |
| tuple: The response, reference, and reference text for the Graph-RAG and RAG models. | |
| """ | |
| if not query.strip(): | |
| raise gr.Error("Please enter a query before asking.") | |
| grag_response, grag_reference, grag_reference_text = query_graph_qa( | |
| graph_rag_index, query, search_level | |
| ) | |
| # rag_response, rag_reference, rag_reference_text = query_rag_qa( | |
| # rag_index, query, search_level | |
| # ) | |
| return ( | |
| grag_response.response, | |
| # grag_reference, | |
| # grag_reference_text, | |
| # rag_response, | |
| # rag_reference, | |
| # rag_reference_text, | |
| ) | |
| # def eval_llm(query, rag_response, grag_response): | |
| # """ | |
| # Evaluate the Graph-RAG and RAG responses using an LLM. | |
| # Args: | |
| # query (str): The query that was asked. | |
| # rag_response (str): The response from the Vanilla-RAG model. | |
| # grag_response (str): The response from the Graph-RAG model. | |
| # Returns: | |
| # str: The evaluation text on various criteria from the LLM. | |
| # """ | |
| # if not query.strip() or not rag_response.strip() or not grag_response.strip(): | |
| # raise gr.Error("Please ask a query and get responses before evaluating.") | |
| # eval_text = evaluate_llm(query, grag_response, rag_response) | |
| # return eval_text | |
| # def reason_and_plot(query, grag_response, grag_reference): | |
| # """ | |
| # Get the reasoning graph for a query and plot the knowledge graph. | |
| # Args: | |
| # query (str): The query to ask the Graph-RAG. | |
| # grag_response (str): The response from the Graph-RAG model. | |
| # grag_reference (str): The reference text from the Graph-RAG model. | |
| # Returns: | |
| # tuple: The reasoning graph and the HTML to plot the knowledge graph. | |
| # """ | |
| # if not query.strip() or not grag_response.strip() or not grag_reference.strip(): | |
| # raise gr.Error( | |
| # "Please ask a query and get a Graph-RAG response before reasoning." | |
| # ) | |
| # graph_reasoning = reasoning_graph(query, grag_response, grag_reference) | |
| # escaped_html = plot_subgraph(grag_reference) | |
| # iframe_html = f'<iframe srcdoc="{escaped_html}" width="100%" height="400px" frameborder="0"></iframe>' | |
| # return graph_reasoning, iframe_html | |
| def show_graph(): | |
| """ | |
| Show the latest graph visualization in an iframe. | |
| Returns: | |
| str: The HTML content to display the graph visualization in an iframe. | |
| """ | |
| graph_vis_dir = os.getenv("GRAPH_VIS", "graph_vis") | |
| try: | |
| latest_graph = get_latest_html_file(graph_vis_dir) | |
| if latest_graph: | |
| with open(latest_graph, "r", encoding="utf-8") as f: | |
| html_content = f.read() | |
| encoded_html = base64.b64encode(html_content.encode()).decode() | |
| iframe_html = f'<iframe src="data:text/html;base64,{encoded_html}" width="100%" height="1000px" frameborder="0"></iframe>' | |
| return iframe_html | |
| else: | |
| return "No graph visualization found." | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| def reveal_coupon(query, grag_response): | |
| """ | |
| Get the coupon from the query and response. | |
| Args: | |
| query (str): Query asked to Graph-RAG. | |
| grag_response (str): Response from the Graph-RAG model. | |
| Returns: | |
| str: Coupon with reasoning. | |
| """ | |
| if not query.strip() or not grag_response.strip(): | |
| raise gr.Error("Please ask a query and get a response before revealing the coupon.") | |
| coupon = get_coupon(query, grag_response) | |
| return coupon | |
| # from gradio import ChatMessage | |
| # def chat_function(message, history): | |
| # history.append(ChatMessage(role="user", content=message)) | |
| # history.append(ChatMessage(role="assistant", content="Hello, how can I help you?")) | |
| # return history | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Comfy Virtual Assistant") | |
| chatbot = gr.Chatbot( | |
| label="Comfy Virtual Assistant", | |
| type="messages", | |
| scale=1, | |
| # suggestions = [ | |
| # {"text": "How much iphone cost?"}, | |
| # {"text": "What phone options do i have ?"} | |
| # ], | |
| ) | |
| msg = gr.Textbox(label="Input Your Query") | |
| clear = gr.ClearButton([msg, chatbot]) | |
| def respond(message, chat_history): | |
| bot_message = query_tqa(message, 2) | |
| # chat_history.append((message, bot_message)) | |
| chat_history.append(ChatMessage(role="user", content=message)) | |
| chat_history.append(ChatMessage(role="assistant", content=bot_message)) | |
| time.sleep(1) | |
| return "", chat_history | |
| msg.submit(respond, [msg, chatbot], [msg, chatbot]) | |
| # with gr.Row(): | |
| # with gr.Column(scale=4): | |
| # query_input = gr.Textbox(label="Input Your Query", lines=3) | |
| # # with gr.Column(scale=1): | |
| # # search_level = gr.Slider( | |
| # # minimum=1, maximum=50, value=3, step=5, label="Search Level" | |
| # # ) | |
| # ask_button = gr.Button("Ask Comfy", variant="primary") | |
| # examples = gr.Examples( | |
| # examples=[ | |
| # ["Recommend me an apple phone that has more than 10MP camera."], | |
| # ["What is the price of Samsung Galaxy S24 Ultra 12/256Gb Titanium Gray"], | |
| # ["I want a phone with 5000 mAH or more battery"], | |
| # ], | |
| # inputs=[query_input], | |
| # ) | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # gr.Markdown("### Graph-RAG") | |
| # grag_output = gr.Textbox(label="Response", lines=5) | |
| # grag_reference = gr.Textbox(label="Triplets", lines=3) | |
| # with gr.Accordion("Extracted Reference (Raw)", open=False): | |
| # grag_reference_text = gr.Textbox(label="Raw Reference", lines=5) | |
| # with gr.Column(): | |
| # gr.Markdown("### Vanilla RAG") | |
| # rag_output = gr.Textbox(label="Response", lines=5) | |
| # rag_reference = gr.Textbox(label="Extracted Reference", lines=3) | |
| # with gr.Accordion("Extracted Reference (Raw)", open=False): | |
| # rag_reference_text = gr.Textbox(label="Raw Reference", lines=5) | |
| # gr.Markdown("### Coupon") | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # coupon = gr.Text(label="Coupon", lines=1) | |
| # with gr.Column(): | |
| # reveal = gr.Button("Reveal Coupon", variant="secondary") | |
| # with gr.Row(): | |
| # gr.Markdown("### Evaluate and Compare") | |
| # with gr.Row(): | |
| # eval_button = gr.Button("Evaluate LLMs", variant="secondary") | |
| # grag_performance = gr.Textbox(label="Evaluation", lines=3) | |
| # with gr.Row(): | |
| # gr.Markdown("### Graph Reasoning") | |
| # with gr.Row(): | |
| # reason_button = gr.Button("Get Graph Reasoning", variant="secondary") | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # grag_reasoning = gr.Textbox(label="Graph-RAG Reasoning", lines=5) | |
| # with gr.Column(): | |
| # subgraph_plot = gr.HTML() | |
| # with gr.Row(): | |
| # plot_button = gr.Button("Plot Knowledge Graph", variant="secondary") | |
| # kg_output = gr.HTML() | |
| # ask_button.click( | |
| # query_tqa, | |
| # inputs=[query_input, search_level], | |
| # outputs=[ | |
| # grag_output, | |
| # # grag_reference, | |
| # # grag_reference_text, | |
| # # rag_output, | |
| # # rag_reference, | |
| # # rag_reference_text, | |
| # ], | |
| # ) | |
| # eval_button.click( | |
| # eval_llm, | |
| # inputs=[query_input, rag_output, grag_output], | |
| # outputs=[grag_performance], | |
| # ) | |
| # reason_button.click( | |
| # reason_and_plot, | |
| # inputs=[query_input, grag_output, grag_reference], | |
| # outputs=[grag_reasoning, subgraph_plot], | |
| # ) | |
| # plot_button.click( | |
| # show_graph, | |
| # outputs=[kg_output], | |
| # ) | |
| # reveal.click( | |
| # reveal_coupon, | |
| # inputs=[query_input, grag_output], | |
| # outputs=[coupon], | |
| # ) | |
| demo.launch(auth=(os.getenv("ID"), os.getenv("PASS")), share=False) | |