[Feature] apache#185: Improvement of Prompt Configuration for rag_dem…

…o Graph Index generation
Kryst4lDem0ni4s · Mar 2, 2025 · 50a8966 · 50a8966
1 parent 3e0bf46
commit 50a8966
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 53 deletions.
diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py
@@ -23,12 +23,23 @@
 import gradio as gr
 import pandas as pd
 from gradio.utils import NamedString
+import uuid
 
 from hugegraph_llm.config import resource_path, prompt, huge_settings, llm_settings
 from hugegraph_llm.operators.graph_rag_task import RAGPipeline
 from hugegraph_llm.utils.log import log
 
+# A decorator to wrap functions with task id generation and logging
+def with_task_id(func):
+    def wrapper(*args, **kwargs):
+        task_id = str(uuid.uuid4())
+        log.info(f"New task created with id: {task_id}")
+        # Optionally, you could also pass the task_id to the function if needed:
+        # kwargs['task_id'] = task_id
+        return func(*args, **kwargs)
+    return wrapper
 
+@with_task_id
 def rag_answer(
     text: str,
     raw_answer: bool,
@@ -113,67 +124,79 @@ def rag_answer(
 
 def create_rag_block():
     # pylint: disable=R0915 (too-many-statements),C0301
+    log.info("Creating RAG block...")
     gr.Markdown("""## 1. HugeGraph RAG Query""")
     with gr.Row():
         with gr.Column(scale=2):
-            inp = gr.Textbox(value=prompt.default_question, label="Question", show_copy_button=True, lines=3)
+            with gr.Blocks().queue(max_size=20, default_concurrency_limit=5):
+                log.info(f"Queue created...")
 
-            # TODO: Only support inline formula now. Should support block formula
-            gr.Markdown("Basic LLM Answer", elem_classes="output-box-label")
-            raw_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
-            gr.Markdown("Vector-only Answer", elem_classes="output-box-label")
-            vector_only_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
-            gr.Markdown("Graph-only Answer", elem_classes="output-box-label")
-            graph_only_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
-            gr.Markdown("Graph-Vector Answer", elem_classes="output-box-label")
-            graph_vector_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
+                inp = gr.Textbox(value=prompt.default_question, label="Question", show_copy_button=True, lines=3)
 
-            answer_prompt_input = gr.Textbox(
-                value=prompt.answer_prompt, label="Query Prompt", show_copy_button=True, lines=7
-            )
-            keywords_extract_prompt_input = gr.Textbox(
-                value=prompt.keywords_extract_prompt,
-                label="Keywords Extraction Prompt",
-                show_copy_button=True,
-                lines=7,
-            )
-        with gr.Column(scale=1):
-            with gr.Row():
-                raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer")
-                vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer")
-            with gr.Row():
-                graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer")
-                graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer")
+                # TODO: Only support inline formula now. Should support block formula
+                gr.Markdown("Basic LLM Answer", elem_classes="output-box-label")
+                raw_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
+                gr.Markdown("Vector-only Answer", elem_classes="output-box-label")
+                vector_only_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
+                gr.Markdown("Graph-only Answer", elem_classes="output-box-label")
+                graph_only_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
+                gr.Markdown("Graph-Vector Answer", elem_classes="output-box-label")
+                graph_vector_out = gr.Markdown(elem_classes="output-box", show_copy_button=True, latex_delimiters=[{"left":"$", "right":"$", "display":False}])
+
+                answer_prompt_input = gr.Textbox(
+                    value=prompt.answer_prompt, label="Query Prompt", show_copy_button=True, lines=7
+                )
+                keywords_extract_prompt_input = gr.Textbox(
+                    value=prompt.keywords_extract_prompt,
+                    label="Keywords Extraction Prompt",
+                    show_copy_button=True,
+                    lines=7,
+                )
 
-            def toggle_slider(enable):
-                return gr.update(interactive=enable)
+                # try:
+                #     log.debug(f"Current queue size: {queueinstance.size()}")
+                #     log.debug(f"Number of concurrent processes: {queueinstance.concurrent_count()}")
+                # except Exception as e:
+                #     log.error(f"Error: {e}")
 
-            with gr.Column():
+            with gr.Column(scale=1):
+                with gr.Row():
+                    raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer")
+                    vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer")
                 with gr.Row():
-                    online_rerank = llm_settings.reranker_type
-                    rerank_method = gr.Dropdown(
-                        choices=["bleu", ("rerank (online)", "reranker")] if online_rerank else ["bleu"],
-                        value="reranker" if online_rerank else "bleu",
-                        label="Rerank method",
+                    graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer")
+                    graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer")
+
+                def toggle_slider(enable):
+                    return gr.update(interactive=enable)
+
+                with gr.Column():
+                    with gr.Row():
+                        online_rerank = llm_settings.reranker_type
+                        rerank_method = gr.Dropdown(
+                            choices=["bleu", ("rerank (online)", "reranker")] if online_rerank else ["bleu"],
+                            value="reranker" if online_rerank else "bleu",
+                            label="Rerank method",
+                        )
+                        example_num = gr.Number(value=2, label="Template Num (0 to disable it) ", precision=0)
+                        graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False)
+
+                    graph_vector_radio.change(
+                        toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio
+                    )  # pylint: disable=no-member
+                    near_neighbor_first = gr.Checkbox(
+                        value=False,
+                        label="Near neighbor first(Optional)",
+                        info="One-depth neighbors > two-depth neighbors",
                     )
-                    example_num = gr.Number(value=2, label="Template Num (0 to disable it) ", precision=0)
-                    graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False)
+                    custom_related_information = gr.Text(
+                        prompt.custom_rerank_info,
+                        label="Query related information(Optional)",
+                    )
+                    btn = gr.Button("Answer Question", variant="primary")
 
-                graph_vector_radio.change(
-                    toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio
-                )  # pylint: disable=no-member
-                near_neighbor_first = gr.Checkbox(
-                    value=False,
-                    label="Near neighbor first(Optional)",
-                    info="One-depth neighbors > two-depth neighbors",
-                )
-                custom_related_information = gr.Text(
-                    prompt.custom_rerank_info,
-                    label="Query related information(Optional)",
-                )
-                btn = gr.Button("Answer Question", variant="primary")
 
-    btn.click(  # pylint: disable=no-member
+    btn.click( # pylint: disable=no-member
         fn=rag_answer,
         inputs=[
             inp,
@@ -190,6 +213,9 @@ def toggle_slider(enable):
             example_num,
         ],
         outputs=[raw_out, vector_only_out, graph_only_out, graph_vector_out],
+        # queue=True,                       # Enable queueing for this event
+        # concurrency_limit=5,               # Maximum of 5 concurrent executions
+        concurrency_id="rag_answer_task"
     )
 
     gr.Markdown(

diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py
@@ -16,6 +16,7 @@
 # under the License.
 
 import re
+import sys
 import traceback
 from dataclasses import dataclass, field
 from typing import List, Optional
@@ -63,8 +64,12 @@ def __post_init__(self):
                     )
 
             except Exception as e:  # pylint: disable=broad-exception-caught
-                traceback.print_exception(e)
-                self.gs_supported = False
-                log.warning(
+                try:
+                    traceback.print_exception(e)
+                    self.gs_supported = False
+                except Exception as e:
+                    exc_type, exc_value, tb = sys.exc_info()
+                    traceback.print_exception(exc_type, exc_value, tb)
+                    log.warning(
                     "Failed to retrieve API version information from the server, reverting to default v1."
                 )